def METHOD_NAME(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):
        """Swaps the timestep and batch dimensions of a tensor."""
        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.METHOD_NAME(inputs[i], **kwargs)
    # Implementation #1: is the time axis is static, use a Python for loop.
    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)
    # Implementation #2: use backend.vectorized_map.
    outputs = backend.vectorized_map(step_function, ops.arange(timesteps))
    return time_distributed_transpose(outputs)