code: def swap_time_and_batch_axes(inputs):
    
    transposed_axes = tf.concat([[1, 0], tf.range(2, tf.rank(inputs))], axis=0)
    return tf.transpose(inputs, transposed_axes)
docs: Swaps time and batch axis (the first two axis).