code: def _test_sharding(  # noqa C901
      tables: List[EmbeddingBagConfig],
      rank: int,
      world_size: int,
      sharder: ModuleSharder[nn.Module],
      backend: str,
      local_size: Optional[int] = None,
      mch_size: Optional[int] = None,
  ) -> None:
      with MultiProcessContext(
          rank, world_size, backend, local_size
      ) as ctx:
          return_remapped: bool = True
          sparse_arch = SparseArch(
              tables,
              torch.device("meta"),
              return_remapped=return_remapped,
              mch_size=mch_size,
          )
  
          apply_optimizer_in_backward(
              RowWiseAdagrad,
              [
                  sparse_arch._mc_ebc._embedding_bag_collection.embedding_bags[
                      "table_0"
                  ].weight,
                  sparse_arch._mc_ebc._embedding_bag_collection.embedding_bags[
                      "table_1"
                  ].weight,
              ],
              {"lr": 0.01},
          )
          module_sharding_plan = construct_module_sharding_plan(
              sparse_arch._mc_ebc,
              per_param_sharding={
                  "table_0": row_wise(),
                  "table_1": row_wise(),
              },
              local_size=local_size,
              world_size=world_size,
              device_type="cuda"
              if torch.cuda.is_available()
              else "cpu",
              sharder=sharder,
          )
  
          sharded_sparse_arch = _shard_modules(
              module=copy.deepcopy(sparse_arch),
              plan=ShardingPlan({"_mc_ebc": module_sharding_plan}),
              env=ShardingEnv.from_process_group(ctx.pg),
              sharders=[sharder],
              device=ctx.device,
          )
  
          assert isinstance(
              sharded_sparse_arch._mc_ebc,
              ShardedManagedCollisionEmbeddingBagCollection,
          )
          assert isinstance(
              sharded_sparse_arch._mc_ebc._managed_collision_collection,
              ShardedManagedCollisionCollection,
          )
  
imports: import copy
  import unittest
  from typing import Dict, List, Optional, Tuple
  import torch
  import torch.nn
  from hypothesis import given, settings, strategies
  from torchrec.distributed.embeddingbag import ShardedEmbeddingBagCollection
  from torchrec.distributed.mc_embeddingbag import ManagedCollisionEmbeddingBagCollectionSharder, ShardedManagedCollisionEmbeddingBagCollection
  from torchrec.distributed.mc_modules import ShardedManagedCollisionCollection
  from torchrec.distributed.shard import _shard_modules
  from torchrec.distributed.sharding_plan import construct_module_sharding_plan, row_wise
  from torchrec.distributed.test_utils.multi_process import MultiProcessContext, MultiProcessTestBase
  from torchrec.distributed.types import ModuleSharder, ShardingEnv, ShardingPlan
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.modules.mc_embedding_modules import ManagedCollisionEmbeddingBagCollection
  from torchrec.modules.mc_modules import DistanceLFU_EvictionPolicy, ManagedCollisionCollection, MCHManagedCollisionModule
  from torchrec.optim.apply_optimizer_in_backward import apply_optimizer_in_backward
  from torchrec.optim.rowwise_adagrad import RowWiseAdagrad
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor
  from torchrec.test_utils import skip_if_asan_class
label_desc: Use image node_id 1 for calling a global function with example usage: _test_sharding(tables, rank, world_size, sharder, backend, local_size, mch_size) without return types