code: def get_tensor_shapes(
      rank: int,
      model_type: ModelType,
      *,
      tensor_shape: Union[List[int], torch.Size],
      decoder_sequence_length: Optional[int] = None,
      sequence_parallel_enabled: bool = False,
  ) -> Sequence[Sequence[int]]:
      """Get tensors shapes
  
      Args:
          rank: pipeline parallel rank
          model_type:
  
      Keyword Args:
          tensor_shape:
          decoder_sequence_length:
          sequence_parallel_enabled:
      """
      # Determine right tensor sizes (based on position of rank with respect to split
      # rank) and model size.
      # Send two tensors if model is T5 and rank is in decoder stage:
      #     first tensor is decoder (pre-transpose),
      #     second tensor is encoder (post-transpose).
      # If model is T5 and rank is at the boundary:
      #     send one tensor (post-transpose from encoder).
      # Otherwise, send one tensor (pre-transpose).
      assert (
          len(tensor_shape) == 3
      ), f"`tensor_shape` should be [sequence_length, micro_batch_size, hidden_size] but {tensor_shape}"
  
      sequence_length, micro_batch_size, hidden_size = tensor_shape
  
      tensor_shapes = []
  
      if sequence_parallel_enabled:
          seq_length = (
              sequence_length
              // parallel_state.get_tensor_model_parallel_world_size()
          )
      else:
          seq_length = sequence_length
  
      if model_type == ModelType.encoder_and_decoder:
          if sequence_parallel_enabled:
              dec_seq_length = (
                  decoder_sequence_length
                  // parallel_state.get_tensor_model_parallel_world_size()
              )
          else:
              dec_seq_length = decoder_sequence_length
  
          if parallel_state.is_pipeline_stage_before_split(rank):
              tensor_shapes.append(
                  (seq_length, micro_batch_size, hidden_size)
              )
          else:
              tensor_shapes.append(
                  (dec_seq_length, micro_batch_size, hidden_size)
              )
              tensor_shapes.append(
                  (seq_length, micro_batch_size, hidden_size)
              )
      else:
          tensor_shapes.append(
              (seq_length, micro_batch_size, hidden_size)
          )
  
      return tensor_shapes
  
imports: import contextlib
  from typing import Any, List, Optional, Sequence, Union
  import warnings
  import torch
  from apex.transformer import parallel_state
  from apex.transformer.enums import ModelType
  from apex.transformer.pipeline_parallel import p2p_communication
  from apex.transformer.pipeline_parallel.p2p_communication import FutureTensor
  from apex.transformer.pipeline_parallel.utils import get_kth_microbatch
  from apex.transformer.pipeline_parallel.utils import listify_model
  from apex.transformer.pipeline_parallel.utils import get_num_microbatches
  from apex.transformer.pipeline_parallel.utils import get_model_type
  from apex.transformer.pipeline_parallel.schedules.common import Batch
  from apex.transformer.pipeline_parallel.schedules.common import FwdStepFunc
  from apex.transformer.pipeline_parallel.schedules.common import backward_step
  from apex.transformer.pipeline_parallel.schedules.common import forward_step
  from apex.transformer.pipeline_parallel.schedules.common import free_output_tensor
  from apex.transformer.log_util import get_transformer_logger
label_desc: Use image node_id 1 for calling a global function with example usage: get_tensor_shapes(rank, model_type) and returns: tensor_shapes