code: def _initialize_torch_state(self) -> None:
      self._model_parallel_mc_buffer_name_to_sharded_tensor = (
          OrderedDict()
      )
      for (
          table_name,
          mc_module,
      ) in self._managed_collision_modules.items():
          assert (
              self._table_name_to_parameter_sharding[
                  table_name
              ].sharding_type
              == ShardingType.ROW_WISE.value
          )
          mc_module_state_dict = mc_module.state_dict(
              prefix=table_name + "."
          )
          shardable_buffers = set.intersection(
              {
                  name
                  for name, _ in mc_module.named_buffers(
                      prefix=table_name
                  )
              },
              set(mc_module_state_dict.keys()),
          )
          (
              shard_offset,
              shard_size,
              global_size,
          ) = self._mc_module_name_shard_metadata[table_name]
          for name, tensor in mc_module_state_dict.items():
              if name not in shardable_buffers:
                  continue
  
              self._model_parallel_mc_buffer_name_to_sharded_tensor[
                  name
              ] = ShardedTensor._init_from_local_shards(
                  [
                      Shard(
                          tensor=tensor,
                          metadata=ShardMetadata(
                              # pyre-ignore [6]
                              shard_offsets=[shard_offset],
                              # pyre-ignore [6]
                              shard_sizes=[shard_size],
                              placement=(
                                  f"rank:{self._env.rank}/cuda:"
                                  f"{get_local_rank(self._env.world_size, self._env.rank)}"
                              ),
                          ),
                      )
                  ],
                  # pyre-ignore [6]
                  torch.Size([global_size]),
                  process_group=self._env.process_group,
              )
  
      def _post_state_dict_hook(
          module: ShardedManagedCollisionCollection,
          destination: Dict[str, torch.Tensor],
          prefix: str,
          _local_metadata: Dict[str, Any],
      ) -> None:
          for (
              mc_buffer_name,
              sharded_tensor,
          ) in (
              module._model_parallel_mc_buffer_name_to_sharded_tensor.items()
          ):
              destination_key = (
                  f"{prefix}_managed_collision_modules.{mc_buffer_name}"
              )
              destination[destination_key] = sharded_tensor
  
      def _load_state_dict_pre_hook(
          module: "ShardedManagedCollisionCollection",
          state_dict: Dict[str, Any],
          prefix: str,
          *args: Any,
      ) -> None:
          for (
              mc_buffer_name,
              _sharded_tensor,
          ) in (
              module._model_parallel_mc_buffer_name_to_sharded_tensor.items()
          ):
              key = (
                  f"{prefix}_managed_collision_modules.{mc_buffer_name}"
              )
              if key in state_dict:
                  if isinstance(state_dict[key], ShardedTensor):
                      local_shards = state_dict[key].local_shards()
                      state_dict[key] = local_shards[0].tensor
                  else:
                      raise RuntimeError(
                          f"Unexpected state_dict key type {type(state_dict[key])} found for {key}"
                      )
  
      self._register_state_dict_hook(_post_state_dict_hook)
      self._register_load_state_dict_pre_hook(
          _load_state_dict_pre_hook, with_module=True
      )
  
imports: import copy
  from collections import defaultdict, OrderedDict
  from dataclasses import dataclass
  from typing import Any, DefaultDict, Dict, Iterator, List, Optional, Type
  import torch
  import torch.distributed
  from torch import nn
  from torch.distributed._shard.sharded_tensor import Shard
  from torchrec.distributed.comm import get_local_rank
  from torchrec.distributed.embedding import _construct_jagged_tensors, EmbeddingCollectionContext
  from torchrec.distributed.embedding_sharding import EmbeddingSharding, EmbeddingShardingContext, EmbeddingShardingInfo, KJTListSplitsAwaitable
  from torchrec.distributed.embedding_types import BaseEmbeddingSharder, GroupedEmbeddingConfig, KJTList
  from torchrec.distributed.sharding.rw_sequence_sharding import RwSequenceEmbeddingDist, RwSequenceEmbeddingSharding
  from torchrec.distributed.sharding.rw_sharding import BaseRwEmbeddingSharding, RwSparseFeaturesDist
  from torchrec.distributed.sharding.sequence_sharding import SequenceShardingContext
  from torchrec.distributed.types import Awaitable, LazyAwaitable, ParameterSharding, QuantizedCommCodecs, ShardedModule, ShardedTensor, ShardingEnv, ShardingType, ShardMetadata
  from torchrec.distributed.utils import append_prefix
  from torchrec.modules.mc_modules import apply_mc_method_to_jt_dict, ManagedCollisionCollection
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor
label_desc: Use image node_id 2 for calling the ShardedManagedCollisionCollection obj's underlying member method code with example usage: obj._initialize_torch_state() without return types