code: def get_tensor_shapes(
      *,
      rank: int,
      model_type: ModelType,
      seq_length: int,
      micro_batch_size: int,
      decoder_seq_length: int,
      config,
  ):
      # 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).
      tensor_shapes = []
  
      seq_length = (
          seq_length // parallel_state.get_context_parallel_world_size()
      )
      if model_type == ModelType.encoder_and_decoder:
          decoder_seq_length = (
              decoder_seq_length
              // parallel_state.get_context_parallel_world_size()
          )
  
      if config.sequence_parallel:
          seq_length = (
              seq_length
              // parallel_state.get_tensor_model_parallel_world_size()
          )
          if model_type == ModelType.encoder_and_decoder:
              decoder_seq_length = (
                  decoder_seq_length
                  // parallel_state.get_tensor_model_parallel_world_size()
              )
  
      if model_type == ModelType.encoder_and_decoder:
          if parallel_state.is_pipeline_stage_before_split(rank):
              tensor_shapes.append(
                  (seq_length, micro_batch_size, config.hidden_size)
              )
          else:
              tensor_shapes.append(
                  (
                      decoder_seq_length,
                      micro_batch_size,
                      config.hidden_size,
                  )
              )
              tensor_shapes.append(
                  (seq_length, micro_batch_size, config.hidden_size)
              )
      else:
          tensor_shapes.append(
              (seq_length, micro_batch_size, config.hidden_size)
          )
      return tensor_shapes
  
imports: import contextlib
  from typing import Callable, Iterator, List, Optional, Union
  import torch
  from torch.autograd.variable import Variable
  from megatron.core import parallel_state
  from megatron.core.enums import ModelType
  from megatron.core.pipeline_parallel import p2p_communication
  from megatron.core.utils import get_attr_wrapped_model, get_model_config, get_model_type
label_desc: Use image node_id 8 for calling a global function with example usage: get_tensor_shapes() and returns: tensor_shapes