code: def make_time_major(policy, seq_lens, tensor):
      """Swaps batch and trajectory axis.
  
      Args:
          policy: Policy reference
          seq_lens: Sequence lengths if recurrent or None
          tensor: A tensor or list of tensors to reshape.
  
      Returns:
          res: A tensor with swapped axes or a list of tensors with
          swapped axes.
      """
      if isinstance(tensor, (list, tuple)):
          return [make_time_major(policy, seq_lens, t) for t in tensor]
  
      if policy.is_recurrent():
          B = seq_lens.shape[0]
          T = tensor.shape[0] // B
      else:
          # Important: chop the tensor into batches at known episode cut
          # boundaries.
          # TODO: (sven) this is kind of a hack and won't work for
          #  batch_mode=complete_episodes.
          T = policy.config["rollout_fragment_length"]
          B = tensor.shape[0] // T
      rs = torch.reshape(tensor, [B, T] + list(tensor.shape[1:]))
  
      # Swap B and T axes.
      res = torch.transpose(rs, 1, 0)
  
      return res
  
imports: import gymnasium
  import logging
  import numpy
  from typing import Dict, List, Optional, Type, Union
  import ray
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.evaluation.postprocessing import compute_bootstrap_value
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.action_dist import ActionDistribution
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import EntropyCoeffSchedule, LearningRateSchedule, ValueNetworkMixin
  from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.torch_utils import apply_grad_clipping, explained_variance, global_norm, sequence_mask
  from ray.rllib.utils.typing import TensorType
label_desc: Use image node_id 1 for calling a global function with example usage: make_time_major(policy, seq_lens, tensor) and returns: res, unknown