class_docstr: Self-attention layer class
  
  Self-attention layer takes input with size [s, b, h]
  and returns output of the same size.
class_docstr_tok: ["Self-attention","layer","class","Self-attention","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","output","of","the","same","size","."]
code: def sharded_state_dict(
      self, prefix="", sharded_key_prefix=None, sharded_offsets=()
  ):
      sharded_key_prefix = (
          prefix if sharded_key_prefix is None else sharded_key_prefix
      )
      sharded_state_dict = {}
      for name, module in (
          ("linear_qkv", self.linear_qkv),
          ("linear_proj", self.linear_proj),
      ):
          sub_sd = module.sharded_state_dict(
              prefix=f"{prefix}{name}.",
              sharded_key_prefix=f"{sharded_key_prefix}{name}.",
              sharded_offsets=sharded_offsets,
          )
          sharded_state_dict.update(sub_sd)
      return sharded_state_dict
  
imports: from abc import ABC, abstractmethod
  from dataclasses import dataclass
  from typing import Union
  import torch
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.models.common.embeddings.rotary_pos_embedding import apply_rotary_pos_emb
  from megatron.core.transformer.custom_layers.transformer_engine import SplitAlongDim
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
  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.utils import divide
  from .enums import AttnMaskType
  from .transformer_config import TransformerConfig
  from .utils import make_sharded_tensors_for_checkpoint
label_desc: Use image node_id 3 for calling the SelfAttention obj's underlying member method code with example usage: obj.sharded_state_dict(prefix, sharded_key_prefix, sharded_offsets) and returns: sharded_state_dict