code: def test_load_state_dict_cw_multiple_shards(
      self, sharder_type: str, sharding_type: str, kernel_type: str
  ) -> None:
      sharders = [
          cast(
              ModuleSharder[nn.Module],
              create_test_sharder(
                  sharder_type,
                  sharding_type,
                  kernel_type,
                  fused_params={
                      "learning_rate": 0.2,
                      "optimizer": EmbOptimType.EXACT_ROWWISE_ADAGRAD,
                  },
              ),
          ),
      ]
  
      constraints = defaultdict(
          lambda: trec_dist.planner.ParameterConstraints()
      )
      num_cw_shards_per_table = {}
      for table in self.tables + self.weighted_tables:
          constraints[table.name].min_partition = 4
          num_cw_shards_per_table[table.name] = table.embedding_dim // 4
  
      (m1, m2), batch = self._generate_dmps_and_batch(
          sharders, constraints=constraints
      )
  
      # load the second's (m2's) with the first (m1's) state_dict
      m2.load_state_dict(
          cast("OrderedDict[str, torch.Tensor]", m1.state_dict())
      )
  
      # load optimizer state dict
  
      # Check to see that we can load optimizer state
      src_optimizer = m1.fused_optimizer
      dst_optimizer = m2.fused_optimizer
  
      src_optimizer_state_dict = src_optimizer.state_dict()
      dst_optimizer_state_dict = dst_optimizer.state_dict()
      m2.fused_optimizer.load_state_dict(src_optimizer_state_dict)
  
      # validate the models are equivalent
      loss1, pred1 = m1(batch)
      loss2, pred2 = m2(batch)
      self.assertTrue(torch.equal(loss1, loss2))
      self.assertTrue(torch.equal(pred1, pred2))
  
      sd1 = m1.state_dict()
      for key, value in m2.state_dict().items():
          if "." in key:
              table_name = key.split(".")[-2]
              v2 = sd1[key]
              if isinstance(value, ShardedTensor):
                  self.assertEqual(
                      len(value.local_shards()),
                      num_cw_shards_per_table[table_name],
                  )
                  dst = value.local_shards()[0].tensor
              else:
                  dst = value
  
              if isinstance(v2, ShardedTensor):
                  self.assertEqual(
                      len(value.local_shards()),
                      num_cw_shards_per_table[table_name],
                  )
  
                  for src_local_shard, dst_local_shard in zip(
                      value.local_shards(), v2.local_shards()
                  ):
                      self.assertTrue(
                          torch.equal(
                              src_local_shard.tensor,
                              dst_local_shard.tensor,
                          )
                      )
              else:
                  src = v2
                  self.assertTrue(torch.equal(src, dst))
  
      for (
          param_name,
          dst_param_group,
      ) in dst_optimizer_state_dict.items():
          src_param_group = src_optimizer_state_dict[param_name]
  
          for state_key, dst_opt_state in dst_param_group.items():
              table_name = state_key.split(".")[-2]
              src_opt_state = src_param_group[state_key]
              if isinstance(dst_opt_state, ShardedTensor):
                  self.assertIsInstance(
                      src_param_group[state_key], ShardedTensor
                  )
  
                  self.assertEqual(
                      len(dst_opt_state.local_shards()),
                      num_cw_shards_per_table[table_name],
                  )
  
                  self.assertEqual(
                      len(src_opt_state.local_shards()),
                      num_cw_shards_per_table[table_name],
                  )
  
                  for src_local_shard, dst_local_shard in zip(
                      src_opt_state.local_shards(),
                      dst_opt_state.local_shards(),
                  ):
                      self.assertTrue(
                          torch.equal(
                              src_local_shard.tensor,
                              dst_local_shard.tensor,
                          )
                      )
              elif isinstance(dst_opt_state, torch.Tensor):
                  self.assertIsInstance(src_opt_state, torch.Tensor)
  
imports: import os
  import unittest
  from collections import defaultdict
  from typing import Any, Callable, cast, Dict, List, Optional, OrderedDict, Tuple
  import numpy
  import torch
  import torch.nn
  from fbgemm_gpu.split_embedding_configs import EmbOptimType
  from hypothesis import given, settings, strategies, Verbosity
  from torch import distributed
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torchrec import distributed
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel, EmbeddingTableConfig
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder, ShardedEmbeddingBagCollection
  from torchrec.distributed.fused_embeddingbag import ShardedFusedEmbeddingBagCollection
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.planner import EmbeddingShardingPlanner, ParameterConstraints, Topology
  from torchrec.distributed.sharding_plan import get_default_sharders
  from torchrec.distributed.test_utils.test_model import _get_default_rtol_and_atol, ModelInput, TestSparseNN, TestSparseNNBase
  from torchrec.distributed.test_utils.test_sharding import copy_state_dict, create_test_sharder, gen_model_and_input, ModelInputCallable, SharderType
  from torchrec.distributed.types import ModuleSharder, ShardingEnv, ShardingPlan, ShardingType
  from torchrec.modules.embedding_configs import BaseEmbeddingConfig, EmbeddingBagConfig, PoolingType
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.modules.fused_embedding_modules import FusedEmbeddingBagCollection
  from torchrec.test_utils import get_free_port, seed_and_log
label_desc: Use image node_id 9 for calling the ModelParallelStateDictBase obj's underlying member method code with example usage: obj.test_load_state_dict_cw_multiple_shards(sharder_type, sharding_type, kernel_type) without return types