code: def _test_sharding(
      tables: List[EmbeddingBagConfig],
      initial_state_dict: Dict[str, Any],
      rank: int,
      world_size: int,
      kjt_input_per_rank: List[KeyedJaggedTensor],
      backend: str,
      module_sharding_plan: EmbeddingModuleShardingPlan,
      local_size: Optional[int] = None,
  ) -> None:
      trec_dist.comm_ops.set_gradient_division(False)
      with MultiProcessContext(
          rank, world_size, backend, local_size
      ) as ctx:
          kjt_input_per_rank = [
              kjt.to(ctx.device) for kjt in kjt_input_per_rank
          ]
          initial_state_dict = {
              fqn: tensor.to(ctx.device)
              for fqn, tensor in initial_state_dict.items()
          }
  
          model = EmbeddingBagCollection(
              tables=tables,
              device=ctx.device,
          )
  
          sharder = get_module_to_default_sharders()[type(model)]
  
          unsharded_model = model
          sharded_model = sharder.shard(
              module=model,
              params=module_sharding_plan,
              env=ShardingEnv.from_process_group(ctx.pg),
              device=ctx.device,
          )
  
          unsharded_model.load_state_dict(
              copy.deepcopy(initial_state_dict)
          )
          copy_state_dict(
              sharded_model.state_dict(),
              copy.deepcopy(initial_state_dict),
          )
  
          feature_keys = []
          for table in tables:
              feature_keys.extend(table.feature_names)
  
          # each rank gets a subbatch
          sharded_model_pred_kt = sharded_model(
              kjt_input_per_rank[ctx.rank]
          ).to_dict()
          _sharded_model_pred = torch.stack(  # noqa
              [
                  sharded_model_pred_kt[feature]
                  for feature in feature_keys
              ]
          )
  
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 a global function with example usage: _test_sharding(tables, initial_state_dict, rank, world_size, kjt_input_per_rank, backend, module_sharding_plan, local_size) without return types