code: def get_tensor_shapes_for_tp(transformer_config, tp_size):
      hs = transformer_config.hidden_size
      return {
          "0.mlp.linear_fc1.layer_norm_weight": (hs,),
          "0.mlp.linear_fc1.layer_norm_bias": (hs,),
          "0.mlp.linear_fc1.weight": (hs * 4 // tp_size, hs),
          "0.mlp.linear_fc1.bias": (hs * 4 // tp_size,),
          "0.mlp.linear_fc2.weight": (hs, hs * 4 // tp_size),
          "0.mlp.linear_fc2.bias": (hs,),
          "0.self_attention.linear_proj.weight": (hs, hs // tp_size),
          "0.self_attention.linear_proj.bias": (hs,),
          "0.self_attention.linear_qkv.layer_norm_weight": (hs,),
          "0.self_attention.linear_qkv.layer_norm_bias": (hs,),
          "0.self_attention.linear_qkv.weight": (hs * 3 // tp_size, hs),
          "0.self_attention.linear_qkv.bias": (hs * 3 // tp_size,),
      }
  
file_path: megatron-lm/tests/unit_tests/transformer/test_transformer_layer.py
imports: import pytest
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.mapping import ShardedObject, ShardedTensor
  from megatron.core.transformer.transformer_layer import TransformerLayer
  from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
  from tests.unit_tests.test_utilities import Utils
label_desc: Use image node_id 1 for calling a global function with example usage: get_tensor_shapes_for_tp(transformer_config, tp_size) and returns: dict