code: def _test_sharding_and_remapping(  # noqa C901
      tables: List[EmbeddingConfig],
      rank: int,
      world_size: int,
      kjt_input_per_rank: List[KeyedJaggedTensor],
      kjt_out_per_iter_per_rank: List[List[KeyedJaggedTensor]],
      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:
          kjt_input = kjt_input_per_rank[rank].to(ctx.device)
          kjt_out_per_iter = [
              kjt[rank].to(ctx.device)
              for kjt in kjt_out_per_iter_per_rank
          ]
          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_ec._embedding_collection.embeddings[
                      "table_0"
                  ].weight,
                  sparse_arch._mc_ec._embedding_collection.embeddings[
                      "table_1"
                  ].weight,
              ],
              {"lr": 0.01},
          )
          module_sharding_plan = construct_module_sharding_plan(
              sparse_arch._mc_ec,
              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_ec": module_sharding_plan}),
              env=ShardingEnv.from_process_group(ctx.pg),
              sharders=[sharder],
              device=ctx.device,
          )
  
          assert isinstance(
              sharded_sparse_arch._mc_ec,
              ShardedManagedCollisionEmbeddingCollection,
          )
          assert isinstance(
              sharded_sparse_arch._mc_ec._embedding_collection,
              ShardedEmbeddingCollection,
          )
          assert (
              sharded_sparse_arch._mc_ec._embedding_collection._has_uninitialized_input_dist
              is False
          )
          assert (
              not hasattr(
                  sharded_sparse_arch._mc_ec._embedding_collection,
                  "_input_dists",
              )
              or len(
                  sharded_sparse_arch._mc_ec._embedding_collection._input_dists
              )
              == 0
          )
  
          assert isinstance(
              sharded_sparse_arch._mc_ec._managed_collision_collection,
              ShardedManagedCollisionCollection,
          )
  
          test_state_dict = sharded_sparse_arch.state_dict()
          sharded_sparse_arch.load_state_dict(test_state_dict)
  
          # sharded model
          # each rank gets a subbatch
          loss1, remapped_ids1 = sharded_sparse_arch(kjt_input)
          loss1.backward()
          loss2, remapped_ids2 = sharded_sparse_arch(kjt_input)
          loss2.backward()
          remapped_ids = [remapped_ids1, remapped_ids2]
          for key in kjt_input.keys():
              for i, kjt_out in enumerate(kjt_out_per_iter):
                  assert torch.equal(
                      remapped_ids[i][key].values(),
                      kjt_out[key].values(),
                  ), f"feature {key} on {ctx.rank} iteration {i} does not match, got {remapped_ids[i][key].values()}, expect {kjt_out[key].values()}"
  
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.embedding import ShardedEmbeddingCollection
  from torchrec.distributed.mc_embedding import ManagedCollisionEmbeddingCollectionSharder, ShardedManagedCollisionEmbeddingCollection
  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 EmbeddingConfig
  from torchrec.modules.embedding_modules import EmbeddingCollection
  from torchrec.modules.mc_embedding_modules import ManagedCollisionEmbeddingCollection
  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 2 for calling a global function with example usage: _test_sharding_and_remapping(tables, rank, world_size, kjt_input_per_rank, kjt_out_per_iter_per_rank, sharder, backend, local_size, mch_size) without return types