code: def test_column_wise(self) -> None:
      embedding_bag_config = [
          EmbeddingBagConfig(
              name=f"table_{idx}",
              feature_names=[f"feature_{idx}"],
              embedding_dim=64,
              num_embeddings=4096,
          )
          for idx in range(2)
      ]
      module_sharding_plan = construct_module_sharding_plan(
          EmbeddingBagCollection(tables=embedding_bag_config),
          per_param_sharding={
              "table_0": column_wise(ranks=[0, 1]),
              "table_1": column_wise(ranks=[0, 1]),
          },
          local_size=2,
          world_size=2,
          device_type="cuda",
      )
      expected = {
          "table_0": ParameterSharding(
              sharding_type="column_wise",
              compute_kernel="dense",
              ranks=[0, 1],
              sharding_spec=EnumerableShardingSpec(
                  shards=[
                      ShardMetadata(
                          shard_offsets=[0, 0],
                          shard_sizes=[4096, 32],
                          placement="rank:0/cuda:0",
                      ),
                      ShardMetadata(
                          shard_offsets=[0, 32],
                          shard_sizes=[4096, 32],
                          placement="rank:1/cuda:1",
                      ),
                  ]
              ),
          ),
          "table_1": ParameterSharding(
              sharding_type="column_wise",
              compute_kernel="dense",
              ranks=[0, 1],
              sharding_spec=EnumerableShardingSpec(
                  shards=[
                      ShardMetadata(
                          shard_offsets=[0, 0],
                          shard_sizes=[4096, 32],
                          placement="rank:0/cuda:0",
                      ),
                      ShardMetadata(
                          shard_offsets=[0, 32],
                          shard_sizes=[4096, 32],
                          placement="rank:1/cuda:1",
                      ),
                  ]
              ),
          ),
      }
      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 2 for calling the ConstructParameterShardingTest obj's underlying member method code with example usage: obj.test_column_wise() without return types