code: def make_time_major(
      tensor: Union["tf.Tensor", List["tf.Tensor"]],
      *,
      trajectory_len: int = None,
      recurrent_seq_len: int = None,
  ):
      """Swaps batch and trajectory axis.
  
      Args:
          tensor: A tensor or list of tensors to swap the axis of.
              NOTE: Each tensor must have the shape [B * T] where B is the batch size and
              T is the trajectory length.
          trajectory_len: The length of each trajectory being transformed.
              If None then `recurrent_seq_len` must be set.
          recurrent_seq_len: Sequence lengths if recurrent.
              If None then `trajectory_len` must be set.
  
      Note: Either `trajectory_len` or `recurrent_seq_len` must be set. `trajectory_len`
          should be used in cases where tensor is not produced from a
          RNN/recurrent module. `recurrent_seq_len` should be used in those cases instead.
  
      Returns:
          A tensor with swapped axes or a list of tensors with swapped axes.
      """
      if isinstance(tensor, list):
          return [
              make_time_major(
                  _tensor, trajectory_len, recurrent_seq_len
              )
              for _tensor in tensor
          ]
  
      assert (trajectory_len != recurrent_seq_len) and (
          trajectory_len is None or recurrent_seq_len is None
      ), "Either trajectory_len or recurrent_seq_len must be set."
  
      if recurrent_seq_len:
          B = tf.shape(recurrent_seq_len)[0]
          T = tf.shape(tensor)[0] // B
      else:
          T = trajectory_len
          B = tf.shape(tensor)[0] // T
      rs = tf.reshape(
          tensor, tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0)
      )
  
      # swap B and T axes
      res = tf.transpose(
          rs,
          [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0]))),
      )
  
      return res
  
label_desc: Use image node_id 1 for calling a global function with example usage: make_time_major(tensor) and returns: res, unknown