code: def test_construct_module_sharding_plan(self) -> None:
      embedding_bag_config = [
          EmbeddingBagConfig(
              name=f"table_{idx}",
              feature_names=[f"feature_{idx}"],
              embedding_dim=256,
              num_embeddings=32 * 32,
          )
          for idx in range(5)
      ]
  
      expected = {
          "table_0": ParameterSharding(
              sharding_type="data_parallel",
              compute_kernel="dense",
              ranks=[
                  0,
                  1,
                  2,
                  3,
                  4,
                  5,
                  6,
                  7,
                  8,
                  9,
                  10,
                  11,
                  12,
                  13,
                  14,
                  15,
                  16,
                  17,
                  18,
                  19,
                  20,
                  21,
                  22,
                  23,
                  24,
                  25,
                  26,
                  27,
                  28,
                  29,
                  30,
                  31,
              ],
              sharding_spec=None,
          ),
          "table_1": ParameterSharding(
              sharding_type="table_wise",
              compute_kernel="dense",
              ranks=[1],
              sharding_spec=EnumerableShardingSpec(
                  shards=[
                      ShardMetadata(
                          shard_offsets=[0, 0],
                          shard_sizes=[1024, 256],
                          placement="rank:1/cuda:1",
                      )
                  ]
              ),
          ),
          "table_2": ParameterSharding(
              sharding_type="row_wise",
              compute_kernel="dense",
              ranks=[
                  0,
                  1,
                  2,
                  3,
                  4,
                  5,
                  6,
                  7,
                  8,
                  9,
                  10,
                  11,
                  12,
                  13,
                  14,
                  15,
                  16,
                  17,
                  18,
                  19,
                  20,
                  21,
                  22,
                  23,
                  24,
                  25,
                  26,
                  27,
                  28,
                  29,
                  30,
                  31,
              ],
              sharding_spec=EnumerableShardingSpec(
                  shards=[
                      ShardMetadata(
                          shard_offsets=[0, 0],
                          shard_sizes=[32, 256],
                          placement="rank:0/cuda:0",
                      ),
                      ShardMetadata(
                          shard_offsets=[32, 0],
                          shard_sizes=[32, 256],
                          placement="rank:1/cuda:1",
                      ),
                      ShardMetadata(
                          shard_offsets=[64, 0],
                          shard_sizes=[32, 256],
                          placement="rank:2/cuda:2",
                      ),
                      ShardMetadata(
                          shard_offsets=[96, 0],
                          shard_sizes=[32, 256],
                          placement="rank:3/cuda:3",
                      ),
                      ShardMetadata(
                          shard_offsets=[128, 0],
                          shard_sizes=[32, 256],
                          placement="rank:4/cuda:4",
                      ),
                      ShardMetadata(
                          shard_offsets=[160, 0],
                          shard_sizes=[32, 256],
                          placement="rank:5/cuda:5",
                      ),
                      ShardMetadata(
                          shard_offsets=[192, 0],
                          shard_sizes=[32, 256],
                          placement="rank:6/cuda:6",
                      ),
                      ShardMetadata(
                          shard_offsets=[224, 0],
                          shard_sizes=[32, 256],
                          placement="rank:7/cuda:7",
                      ),
                      ShardMetadata(
                          shard_offsets=[256, 0],
                          shard_sizes=[32, 256],
                          placement="rank:8/cuda:0",
                      ),
                      ShardMetadata(
                          shard_offsets=[288, 0],
                          shard_sizes=[32, 256],
                          placement="rank:9/cuda:1",
                      ),
                      ShardMetadata(
                          shard_offsets=[320, 0],
                          shard_sizes=[32, 256],
                          placement="rank:10/cuda:2",
                      ),
                      ShardMetadata(
                          shard_offsets=[352, 0],
                          shard_sizes=[32, 256],
                          placement="rank:11/cuda:3",
                      ),
                      ShardMetadata(
                          shard_offsets=[384, 0],
                          shard_sizes=[32, 256],
                          placement="rank:12/cuda:4",
                      ),
                      ShardMetadata(
                          shard_offsets=[416, 0],
                          shard_sizes=[32, 256],
                          placement="rank:13/cuda:5",
                      ),
                      ShardMetadata(
                          shard_offsets=[448, 0],
                          shard_sizes=[32, 256],
                          placement="rank:14/cuda:6",
                      ),
                      ShardMetadata(
                          shard_offsets=[480, 0],
                          shard_sizes=[32, 256],
                          placement="rank:15/cuda:7",
                      ),
                      ShardMetadata(
                          shard_offsets=[512, 0],
                          shard_sizes=[32, 256],
                          placement="rank:16/cuda:0",
                      ),
                      ShardMetadata(
                          shard_offsets=[544, 0],
                          shard_sizes=[32, 256],
                          placement="rank:17/cuda:1",
                      ),
                      ShardMetadata(
                          shard_offsets=[576, 0],
                          shard_sizes=[32, 256],
                          placement="rank:18/cuda:2",
                      ),
                      ShardMetadata(
                          shard_offsets=[608, 0],
                          shard_sizes=[32, 256],
                          placement="rank:19/cuda:3",
                      ),
                      ShardMetadata(
                          shard_offsets=[640, 0],
                          shard_sizes=[32, 256],
                          placement="rank:20/cuda:4",
                      ),
                      ShardMetadata(
                          shard_offsets=[672, 0],
                          shard_sizes=[32, 256],
                          placement="rank:21/cuda:5",
                      ),
                      ShardMetadata(
                          shard_offsets=[704, 0],
                          shard_sizes=[32, 256],
                          placement="rank:22/cuda:6",
                      ),
                      ShardMetadata(
                          shard_offsets=[736, 0],
                          shard_sizes=[32, 256],
                          placement="rank:23/cuda:7",
                      ),
                      ShardMetadata(
                          shard_offsets=[768, 0],
                          shard_sizes=[32, 256],
                          placement="rank:24/cuda:0",
                      ),
                      ShardMetadata(
                          shard_offsets=[800, 0],
                          shard_sizes=[32, 256],
                          placement="rank:25/cuda:1",
                      ),
                      ShardMetadata(
                          shard_offsets=[832, 0],
                          shard_sizes=[32, 256],
                          placement="rank:26/cuda:2",
                      ),
                      ShardMetadata(
                          shard_offsets=[864, 0],
                          shard_sizes=[32, 256],
                          placement="rank:27/cuda:3",
                      ),
                      ShardMetadata(
                          shard_offsets=[896, 0],
                          shard_sizes=[32, 256],
                          placement="rank:28/cuda:4",
                      ),
                      ShardMetadata(
                          shard_offsets=[928, 0],
                          shard_sizes=[32, 256],
                          placement="rank:29/cuda:5",
                      ),
                      ShardMetadata(
                          shard_offsets=[960, 0],
                          shard_sizes=[32, 256],
                          placement="rank:30/cuda:6",
                      ),
                      ShardMetadata(
                          shard_offsets=[992, 0],
                          shard_sizes=[32, 256],
                          placement="rank:31/cuda:7",
                      ),
                  ]
              ),
          ),
          "table_3": ParameterSharding(
              sharding_type="column_wise",
              compute_kernel="dense",
              ranks=[8, 9],
              sharding_spec=EnumerableShardingSpec(
                  shards=[
                      ShardMetadata(
                          shard_offsets=[0, 0],
                          shard_sizes=[1024, 128],
                          placement="rank:8/cuda:0",
                      ),
                      ShardMetadata(
                          shard_offsets=[0, 128],
                          shard_sizes=[1024, 128],
                          placement="rank:9/cuda:1",
                      ),
                  ]
              ),
          ),
          "table_4": ParameterSharding(
              sharding_type="table_row_wise",
              compute_kernel="dense",
              ranks=[24, 25, 26, 27, 28, 29, 30, 31],
              sharding_spec=EnumerableShardingSpec(
                  shards=[
                      ShardMetadata(
                          shard_offsets=[0, 0],
                          shard_sizes=[128, 256],
                          placement="rank:24/cuda:0",
                      ),
                      ShardMetadata(
                          shard_offsets=[128, 0],
                          shard_sizes=[128, 256],
                          placement="rank:25/cuda:1",
                      ),
                      ShardMetadata(
                          shard_offsets=[256, 0],
                          shard_sizes=[128, 256],
                          placement="rank:26/cuda:2",
                      ),
                      ShardMetadata(
                          shard_offsets=[384, 0],
                          shard_sizes=[128, 256],
                          placement="rank:27/cuda:3",
                      ),
                      ShardMetadata(
                          shard_offsets=[512, 0],
                          shard_sizes=[128, 256],
                          placement="rank:28/cuda:4",
                      ),
                      ShardMetadata(
                          shard_offsets=[640, 0],
                          shard_sizes=[128, 256],
                          placement="rank:29/cuda:5",
                      ),
                      ShardMetadata(
                          shard_offsets=[768, 0],
                          shard_sizes=[128, 256],
                          placement="rank:30/cuda:6",
                      ),
                      ShardMetadata(
                          shard_offsets=[896, 0],
                          shard_sizes=[128, 256],
                          placement="rank:31/cuda:7",
                      ),
                  ]
              ),
          ),
      }
  
      module_sharding_plan = construct_module_sharding_plan(
          EmbeddingBagCollection(tables=embedding_bag_config),
          per_param_sharding={
              "table_0": data_parallel(),
              "table_1": table_wise(rank=1),
              "table_2": row_wise(),
              "table_3": column_wise(ranks=[8, 9]),
              "table_4": table_row_wise(host_index=3),
          },
          local_size=8,
          world_size=32,
          device_type="cuda",
      )
      self.assertDictEqual(expected, module_sharding_plan)
  
imports: import copy
  import unittest
  from typing import Any, Dict, List, Optional
  import hypothesis.strategies
  import torch
  from hypothesis import given, settings, Verbosity
  from torchrec import distributed
  from torchrec.distributed.sharding_plan import column_wise, construct_module_sharding_plan, data_parallel, EmbeddingBagCollectionSharder, EmbeddingCollectionSharder, FeatureProcessedEmbeddingBagCollectionSharder, FusedEmbeddingBagCollectionSharder, get_module_to_default_sharders, ManagedCollisionEmbeddingBagCollectionSharder, ManagedCollisionEmbeddingCollectionSharder, ParameterShardingGenerator, QuantEmbeddingBagCollectionSharder, QuantEmbeddingCollectionSharder, row_wise, table_row_wise, table_wise
  from torchrec.distributed.test_utils.multi_process import MultiProcessContext, MultiProcessTestBase
  from torchrec.distributed.test_utils.test_sharding import copy_state_dict
  from torchrec.distributed.types import EmbeddingModuleShardingPlan, EnumerableShardingSpec, ParameterSharding, ShardingEnv, ShardingPlan, ShardingType, ShardMetadata
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection, EmbeddingCollection
  from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
  from torchrec.modules.fused_embedding_modules import FusedEmbeddingBagCollection
  from torchrec.modules.mc_embedding_modules import ManagedCollisionEmbeddingBagCollection, ManagedCollisionEmbeddingCollection
  from torchrec.quant.embedding_modules import EmbeddingBagCollection, EmbeddingCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
  from torchrec.test_utils import skip_if_asan_class
label_desc: Use image node_id 1 for calling the ConstructParameterShardingTest obj's underlying member method code with example usage: obj.test_construct_module_sharding_plan() without return types