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)
Swaps the timestep and batch dimensions of a tensor.
keras/src/layers/rnn/time_distributed.py