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,
      parameter_sharding_plan: Dict[str, ParameterSharding],
      sharder: ModuleSharder[nn.Module],
      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,
          )
  
          apply_optimizer_in_backward(
              torch.optim.SGD,
              model.parameters(),
              {"lr": 1.0},
          )
  
          unsharded_model = model
          sharded_model = sharder.shard(
              module=model,
              params=parameter_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
              ]
          )
  
          for _it in range(1):
              unsharded_model_pred_kt = []
              for rank in range(ctx.world_size):
                  # simulate the unsharded model run on the entire batch
                  unsharded_model_pred_kt.append(
                      unsharded_model(kjt_input_per_rank[rank])
                  )
  
              all_unsharded_preds = []
              for rank in range(ctx.world_size):
                  unsharded_model_pred_kt_mini_batch = (
                      unsharded_model_pred_kt[rank].to_dict()
                  )
  
                  all_unsharded_preds.extend(
                      [
                          unsharded_model_pred_kt_mini_batch[feature]
                          for feature in feature_keys
                      ]
                  )
                  if rank == ctx.rank:
                      unsharded_model_pred = torch.stack(
                          [
                              unsharded_model_pred_kt_mini_batch[
                                  feature
                              ]
                              for feature in feature_keys
                          ]
                      )
  
              # sharded model
              # each rank gets a subbatch
              sharded_model_pred_kt = sharded_model(
                  kjt_input_per_rank[ctx.rank]
              ).to_dict()
              sharded_model_pred = torch.stack(
                  [
                      sharded_model_pred_kt[feature]
                      for feature in feature_keys
                  ]
              )
  
              # cast to CPU because when casting unsharded_model.to on the same module, there could some race conditions
              # in normal author modelling code this won't be an issue because each rank would individually create
              # their model. output from sharded_pred is correctly on the correct device.
              # Compare predictions of sharded vs unsharded models.
              torch.testing.assert_close(
                  sharded_model_pred.cpu(),
                  unsharded_model_pred.cpu(),
              )
  
              sharded_model_pred.sum().backward()
  
              all_unsharded_preds = torch.stack(all_unsharded_preds)
              all_unsharded_preds.sum().backward()
  
imports: import copy
  import unittest
  from typing import Any, Dict, List, Optional
  import hypothesis.strategies
  import torch
  import torch.nn
  import torchrec.distributed
  from hypothesis import given, settings, Verbosity
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder
  from torchrec.distributed.fbgemm_qcomm_codec import CommType, get_qcomm_codecs_registry, QCommsConfig
  from torchrec.distributed.sharding_plan import column_wise, construct_module_sharding_plan, ParameterShardingGenerator, 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 ModuleSharder, ParameterSharding, ShardingEnv
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  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, parameter_sharding_plan, sharder, local_size) without return types