def transpose_batch_time(tensor):
  """Transposes the batch and time dimensions of the `tensor`.

  If the `tensor` has less than 2 dimensions, then no operation is performed.
  Otherwise, swaps the first 2 dimensions, which are assumed to be the batch and
  time dimensions.

  Args:
    tensor: Tensor to transpose.

  Returns:
    Possibly-transposed version of `tensor`.

  Raises:
    ValueError: If the `tensor` has unknown rank.
  """
  rank = tensor.shape.ndims
  if rank is None:
    raise ValueError("Tensor with unknown rank")
  if rank < 2:
    return tensor
  return tf.transpose(tensor, perm=[1, 0] + list(range(2, rank)))
Transposes the batch and time dimensions of the `tensor`.

  If the `tensor` has less than 2 dimensions, then no operation is performed.
  Otherwise, swaps the first 2 dimensions, which are assumed to be the batch and
  time dimensions.

  Args:
    tensor: Tensor to transpose.

  Returns:
    Possibly-transposed version of `tensor`.

  Raises:
    ValueError: If the `tensor` has unknown rank.
  
language/common/utils/tensor_utils.py