class_docstr: Transformer class.
class_docstr_tok: ["Transformer","class","."]
code: def sharded_state_dict(self, prefix: str = ""):
      sharded_state_dict = {}
  
      layer_prefix = f"{prefix}layers."
      for layer in self.layers:
          sharded_state_dict.update(
              layer.sharded_state_dict(prefix=layer_prefix)
          )
  
      if self.post_process and self.post_layer_norm:
          state_dict = self.state_dict(keep_vars=True)
  
          tensor = state_dict["final_layernorm.weight"]
          layer_name = f"{prefix}final_layernorm.weight"
          sharded_state_dict[
              layer_name
          ] = make_sharded_tensor_for_checkpoint(tensor, layer_name)
  
          # RMSNorm doesn't have bias.
          if "final_layernorm.bias" in state_dict.keys():
              tensor = state_dict["final_layernorm.bias"]
              layer_name = f"{prefix}final_layernorm.bias"
              sharded_state_dict[
                  layer_name
              ] = make_sharded_tensor_for_checkpoint(tensor, layer_name)
  
      return sharded_state_dict
  
imports: import re
  from contextlib import nullcontext
  from dataclasses import dataclass
  from typing import List, Union
  import torch
  from torch import Tensor
  from megatron.core import InferenceParams, parallel_state, tensor_parallel
  from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
  from megatron.core.transformer.custom_layers.transformer_engine import TENorm
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec, build_module
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.transformer.transformer_layer import TransformerLayer
  from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_viewless_tensor
label_desc: Use image node_id 7 for calling the TransformerBlock obj's underlying member method code with example usage: obj.sharded_state_dict(prefix) and returns: sharded_state_dict