wrapper layer to apply every temporal slice of an input from keras import backend from keras import ops from keras apiexport import kerasexport from keras layers core wrapper import wrapper from keras layers layer import layer kerasexportkeras layers timedistributed class timedistributedwrapper def initself layer kwargs if not isinstancelayer layer raise valueerror please initialize timedistributed layer with a fkeras layers layer instance received layer super initlayer kwargs self supportsmasking true def getchildinputshapeself inputshape if not isinstanceinputshape tuple list or leninputshape 3 raise valueerror timedistributed layer should be passed an inputshape fwith at least 3 dimensions received inputshape return inputshape0 inputshape2 def computeoutputshapeself inputshape childinputshape self getchildinputshapeinputshape childoutputshape self layer computeoutputshapechildinputshape return childoutputshape0 inputshape1 childoutputshape1 def buildself inputshape childinputshape self getchildinputshapeinputshape super buildchildinputshape self built true def callself inputs trainingnone masknone inputshape inputs shape maskshape none if mask is none else tuplemask shape batchsize inputshape0 timesteps inputshape1 if maskshape is not none and maskshape 2 batchsize timesteps raise valueerror timedistributed layer should be passed a mask of shape fbatchsize timesteps freceived mask shapemaskshape def timedistributedtransposedata implementation 1 is the time axis is static use a python for loop implementation 2 use backend vectorizedmap this wrapper allows to apply a layer to every temporal slice of an input every input should be at least 3d and the dimension of index one of the first input will be considered to be the temporal dimension consider a batch of 32 video samples where each sample is a 128x128 rgb image with channels_last data format across 10 timesteps the batch input shape is 32 10 128 128 3 you can then use timedistributed to apply the same conv2d layer to each of the 10 timesteps independently inputs layers input shape 10 128 128 3 batch_size 32 conv_2d_layer layers conv2d 64 3 3 outputs layers timedistributed conv_2d_layer inputs outputs shape 32 10 126 126 64 because timedistributed applies the same instance of conv2d to each of the timestamps the same set of weights are used at each timestamp args layer a keras layers layer instance call arguments inputs input tensor of shape batch time or nested tensors and each of which has shape batch time training python boolean indicating whether the layer should behave in training mode or in inference mode this argument is passed to the wrapped layer only if the layer supports this argument mask binary tensor of shape samples timesteps indicating whether a given timestep should be masked this argument is passed to the wrapped layer only if the layer supports this argument swaps the timestep and batch dimensions of a tensor implementation 1 is the time axis is static use a python for loop implementation 2 use backend vectorized_map
from keras import backend
from keras import ops
from keras.api_export import keras_export
from keras.layers.core.wrapper import Wrapper
from keras.layers.layer import Layer


@keras_export("keras.layers.TimeDistributed")
class TimeDistributed(Wrapper):
    

    def __init__(self, layer, **kwargs):
        if not isinstance(layer, Layer):
            raise ValueError(
                "Please initialize `TimeDistributed` layer with a "
                f"`keras.layers.Layer` instance. Received: {layer}"
            )
        super().__init__(layer, **kwargs)
        self.supports_masking = True

    def _get_child_input_shape(self, input_shape):
        if not isinstance(input_shape, (tuple, list)) or len(input_shape) < 3:
            raise ValueError(
                "`TimeDistributed` Layer should be passed an `input_shape` "
                f"with at least 3 dimensions, received: {input_shape}"
            )
        return (input_shape[0], *input_shape[2:])

    def compute_output_shape(self, input_shape):
        child_input_shape = self._get_child_input_shape(input_shape)
        child_output_shape = self.layer.compute_output_shape(child_input_shape)
        return (child_output_shape[0], input_shape[1], *child_output_shape[1:])

    def build(self, input_shape):
        child_input_shape = self._get_child_input_shape(input_shape)
        super().build(child_input_shape)
        self.built = True

    def call(self, inputs, training=None, mask=None):
        input_shape = inputs.shape
        mask_shape = None if mask is None else tuple(mask.shape)
        batch_size = input_shape[0]
        timesteps = input_shape[1]

        if mask_shape is not None and mask_shape[:2] != (batch_size, timesteps):
            raise ValueError(
                "`TimeDistributed` Layer should be passed a `mask` of shape "
                f"({batch_size}, {timesteps}, ...), "
                f"received: mask.shape={mask_shape}"
            )

        def time_distributed_transpose(data):
            
            axes = [1, 0, *range(2, len(data.shape))]
            return ops.transpose(data, axes=axes)

        inputs = time_distributed_transpose(inputs)
        if mask is not None:
            mask = time_distributed_transpose(mask)

        def step_function(i):
            kwargs = {}
            if self.layer._call_has_mask_arg and mask is not None:
                kwargs["mask"] = mask[i]
            if self.layer._call_has_training_arg:
                kwargs["training"] = training
            return self.layer.call(inputs[i], **kwargs)

        

        if inputs.shape[0] is not None:
            outputs = ops.stack(
                [step_function(i) for i in range(inputs.shape[0])]
            )
            return time_distributed_transpose(outputs)

        

        outputs = backend.vectorized_map(step_function, ops.arange(timesteps))
        return time_distributed_transpose(outputs)