code: def _initialize_torch_state(  # noqa: C901
      # Union[ShardedQuantEmbeddingBagCollection, ShardedQuantEmbeddingCollection]
      self,
      tbes: Dict[
          IntNBitTableBatchedEmbeddingBagsCodegen,
          GroupedEmbeddingConfig,
      ],
      table_name_to_parameter_sharding: Dict[str, ParameterSharding],
      tables_weights_prefix: str,  # "embedding_bags" or "embeddings"
  ) -> None:  # noqa
      # State is prepared only in "quant_state_dict_split_scale_bias" mode
      assert (
          tables_weights_prefix == "embedding_bags"
          or tables_weights_prefix == "embeddings"
      )
  
      # weight
      self._table_name_to_local_shards: Dict[str, List[Shard]] = {}
      self._table_name_to_sharded_tensor: Dict[
          str, Union[torch.Tensor, ShardedTensorBase]
      ] = {}
  
      # weight_qscale
      self._table_name_to_local_shards_qscale: Dict[
          str, List[Shard]
      ] = {}
      self._table_name_to_sharded_tensor_qscale: Dict[
          str, Union[torch.Tensor, ShardedTensorBase]
      ] = {}
      self._table_name_to_tensors_list_qscale: Dict[
          str, List[torch.Tensor]
      ] = {}
  
      # weight_qbias
      self._table_name_to_local_shards_qbias: Dict[
          str, List[Shard]
      ] = {}
      self._table_name_to_sharded_tensor_qbias: Dict[
          str, Union[torch.Tensor, ShardedTensorBase]
      ] = {}
      self._table_name_to_tensors_list_qbias: Dict[
          str, List[torch.Tensor]
      ] = {}
  
      # pruning_index_remappings
      self._table_name_to_local_shards_pruning_index_remappings: Dict[
          str, List[Shard]
      ] = {}
      self._table_name_to_sharded_tensor_pruning_index_remappings: Dict[
          str, Union[torch.Tensor, ShardedTensorBase]
      ] = {}
      self._table_name_to_tensors_list_pruning_index_remappings: Dict[
          str, List[torch.Tensor]
      ] = {}
  
      for tbe, config in tbes.items():
          for (
              tbe_split_w,
              tbe_split_qscale,
              tbe_split_qbias,
          ), table in zip(
              tbe.split_embedding_weights_with_scale_bias(
                  split_scale_bias_mode=2
              ),
              config.embedding_tables,
          ):
              if table.data_type == DataType.INT4:
                  tbe_split_w = UInt4Tensor(tbe_split_w)
              elif table.data_type == DataType.INT2:
                  tbe_split_w = UInt2Tensor(tbe_split_w)
  
              # weight shards section:
              assert table.local_metadata
              metadata: ShardMetadata = copy.deepcopy(
                  table.local_metadata
              )
              metadata.shard_sizes = [
                  tbe_split_w.size(0),
                  tbe_split_w.size(1),
              ]
  
              # TODO(ivankobzarev): "meta" sharding support: cleanup when copy to "meta" moves all tensors to "meta"
              # pyre-ignore
              if metadata.placement.device != tbe_split_w.device:
                  metadata.placement = _remote_device(
                      tbe_split_w.device
                  )
              _append_table_shard(
                  self._table_name_to_local_shards,
                  table.name,
                  Shard(tensor=tbe_split_w, metadata=metadata),
              )
              # end of weight shards section
  
              # weight_qscale & weight_qbias section:
              # For RW - ShardedTensorBase
              # For CW - List[Tensor] that logically corresponds to the same unsharded Tensor, but present on each sharded rank
              for (
                  tbe_split_qparam,
                  table_name_to_local_shards,
                  table_name_to_tensors_list,
              ) in [
                  (
                      tbe_split_qscale,
                      self._table_name_to_local_shards_qscale,
                      self._table_name_to_tensors_list_qscale,
                  ),
                  (
                      tbe_split_qbias,
                      self._table_name_to_local_shards_qbias,
                      self._table_name_to_tensors_list_qbias,
                  ),
              ]:
                  assert table.local_metadata
                  metadata: ShardMetadata = copy.deepcopy(
                      table.local_metadata
                  )
                  shard_sizes = metadata.shard_sizes
                  shard_offsets = metadata.shard_offsets
  
                  shard_sizes_cols = shard_sizes[1]
                  shard_offsets_cols = shard_offsets[1]
  
                  parameter_sharding: ParameterSharding = (
                      table_name_to_parameter_sharding[table.name]
                  )
                  sharding_type: str = parameter_sharding.sharding_type
  
                  if sharding_type == ShardingType.COLUMN_WISE.value:
                      if table.name not in table_name_to_tensors_list:
                          assert parameter_sharding.ranks
                          num_shards: int = len(
                              parameter_sharding.ranks
                          )
                          table_name_to_tensors_list[table.name] = [
                              torch.empty([])
                          ] * num_shards
  
                      column_idx = int(
                          shard_offsets_cols / shard_sizes_cols
                      )
                      table_name_to_tensors_list[table.name][
                          column_idx
                      ] = tbe_split_qparam
                  else:
                      qmetadata = ShardMetadata(
                          shard_offsets=metadata.shard_offsets,
                          shard_sizes=[
                              tbe_split_qparam.shape[0],
                              tbe_split_qparam.shape[1],
                          ],
                          # pyre-ignore
                          placement=table.local_metadata.placement,
                      )
                      # TODO(ivankobzarev): "meta" sharding support: cleanup when copy to "meta" moves all tensors to "meta"
                      if (
                          qmetadata.placement.device
                          != tbe_split_qparam.device
                      ):
                          qmetadata.placement = _remote_device(
                              tbe_split_qparam.device
                          )
                      _append_table_shard(
                          table_name_to_local_shards,
                          table.name,
                          Shard(
                              tensor=tbe_split_qparam,
                              metadata=qmetadata,
                          ),
                      )
                  # end of weight_qscale & weight_qbias section
          if table.pruning_indices_remapping is not None:
              for (
                  qparam,
                  table_name_to_local_shards,
                  _,
              ) in [
                  (
                      table.pruning_indices_remapping,
                      self._table_name_to_local_shards_pruning_index_remappings,
                      self._table_name_to_tensors_list_pruning_index_remappings,
                  )
              ]:
                  parameter_sharding: ParameterSharding = (
                      table_name_to_parameter_sharding[table.name]
                  )
                  sharding_type: str = parameter_sharding.sharding_type
  
                  assert sharding_type in [
                      ShardingType.TABLE_WISE.value,
                      ShardingType.COLUMN_WISE.value,
                  ]
  
                  qmetadata = ShardMetadata(
                      shard_offsets=[0],
                      shard_sizes=[
                          qparam.shape[0],
                      ],
                      placement=table.local_metadata.placement,
                  )
                  # TODO(ivankobzarev): "meta" sharding support: cleanup when copy to "meta" moves all tensors to "meta"
                  if qmetadata.placement.device != qparam.device:
                      qmetadata.placement = _remote_device(
                          qparam.device
                      )
                  _append_table_shard(
                      table_name_to_local_shards,
                      table.name,
                      Shard(tensor=qparam, metadata=qmetadata),
                  )
  
      for table_name_to_local_shards, table_name_to_sharded_tensor in [
          (
              self._table_name_to_local_shards,
              self._table_name_to_sharded_tensor,
          ),
          (
              self._table_name_to_local_shards_qscale,
              self._table_name_to_sharded_tensor_qscale,
          ),
          (
              self._table_name_to_local_shards_qbias,
              self._table_name_to_sharded_tensor_qbias,
          ),
      ]:
          for (
              table_name,
              local_shards,
          ) in table_name_to_local_shards.items():
              if len(local_shards) == 1:
                  # Single Tensor per table (TW sharding)
                  table_name_to_sharded_tensor[
                      table_name
                  ] = local_shards[0].tensor
                  continue
  
              # ShardedTensor per table
              global_rows = max(
                  [
                      ls.metadata.shard_offsets[0]
                      + ls.metadata.shard_sizes[0]
                      for ls in local_shards
                  ]
              )
              global_cols = max(
                  [
                      ls.metadata.shard_offsets[1]
                      + ls.metadata.shard_sizes[1]
                      for ls in local_shards
                  ]
              )
              global_metadata: ShardedTensorMetadata = (
                  ShardedTensorMetadata(
                      shards_metadata=[
                          ls.metadata for ls in local_shards
                      ],
                      size=torch.Size([global_rows, global_cols]),
                  )
              )
              table_name_to_sharded_tensor[
                  table_name
              ] = ShardedTensorBase._init_from_local_shards_and_global_metadata(
                  local_shards=local_shards,
                  sharded_tensor_metadata=global_metadata,
              )
  
      for table_name_to_local_shards, table_name_to_sharded_tensor in [
          (
              self._table_name_to_local_shards_pruning_index_remappings,
              self._table_name_to_sharded_tensor_pruning_index_remappings,
          ),
      ]:
          for (
              table_name,
              local_shards,
          ) in table_name_to_local_shards.items():
              # Single Tensor per table (TW sharding)
              table_name_to_sharded_tensor[table_name] = local_shards[
                  0
              ].tensor
              continue
  
      def post_state_dict_hook(
          # Union["ShardedQuantEmbeddingBagCollection", "ShardedQuantEmbeddingCollection"]
          module: ShardedQuantEmbeddingModuleState[
              CompIn, DistOut, Out, ShrdCtx
          ],
          destination: Dict[str, torch.Tensor],
          prefix: str,
          _local_metadata: Dict[str, Any],
      ) -> None:
          for (
              table_name,
              sharded_t,
          ) in module._table_name_to_sharded_tensor.items():
              destination[
                  f"{prefix}{tables_weights_prefix}.{table_name}.weight"
              ] = sharded_t
  
          for sfx, dict_sharded_t, dict_t_list in [
              (
                  "weight_qscale",
                  module._table_name_to_sharded_tensor_qscale,
                  module._table_name_to_tensors_list_qscale,
              ),
              (
                  "weight_qbias",
                  module._table_name_to_sharded_tensor_qbias,
                  module._table_name_to_tensors_list_qbias,
              ),
              (
                  "index_remappings_array",
                  module._table_name_to_sharded_tensor_pruning_index_remappings,
                  module._table_name_to_tensors_list_pruning_index_remappings,
              ),
          ]:
              for (
                  table_name,
                  sharded_t,
              ) in dict_sharded_t.items():
                  destination[
                      f"{prefix}{tables_weights_prefix}.{table_name}.{sfx}"
                  ] = sharded_t
              for (
                  table_name,
                  t_list,
              ) in dict_t_list.items():
                  destination[
                      f"{prefix}{tables_weights_prefix}.{table_name}.{sfx}"
                  ] = t_list
  
      self._register_state_dict_hook(post_state_dict_hook)
  
imports: import copy
  from dataclasses import dataclass
  from typing import Any, Dict, List, Mapping, Optional, TypeVar, Union
  import torch
  from fbgemm_gpu.split_table_batched_embeddings_ops_inference import IntNBitTableBatchedEmbeddingBagsCodegen
  from torch.distributed import _remote_device
  from torch.distributed._shard.sharded_tensor import Shard, ShardedTensorBase, ShardedTensorMetadata, ShardMetadata
  from torchrec.distributed.embedding_sharding import EmbeddingShardingInfo
  from torchrec.distributed.embedding_types import GroupedEmbeddingConfig, ShardedEmbeddingModule
  from torchrec.distributed.types import ParameterSharding, ShardingType
  from torchrec.modules.embedding_configs import DataType
  from torchrec.streamable import Multistreamable
  from torchrec.tensor_types import UInt2Tensor, UInt4Tensor
label_desc: Use image node_id 1 for calling the ShardedQuantEmbeddingModuleState obj's underlying member method code with example usage: obj._initialize_torch_state(tbes, table_name_to_parameter_sharding, tables_weights_prefix) without return types