code: def _test_sharding(  # noqa C901
      tables: List[EmbeddingBagConfig],
      rank: int,
      world_size: int,
      kjt_input_per_rank: List[KeyedJaggedTensor],
      sharder: ModuleSharder[nn.Module],
      backend: str,
      set_gradient_division: bool,
      local_size: Optional[int] = None,
      use_dmp: bool = False,
      use_fp_collection: bool = False,
  ) -> None:
      with MultiProcessContext(
          rank, world_size, backend, local_size
      ) as ctx:
          trec_dist.comm_ops.set_gradient_division(
              set_gradient_division
          )
  
          kjt_input_per_rank = [
              kjt.to(ctx.device) for kjt in kjt_input_per_rank
          ]
  
          (
              sparse_arch,
              sharded_sparse_arch,
          ) = get_unsharded_and_sharded_module(
              tables,
              sharder,
              use_dmp,
              use_fp_collection,
              init_device=ctx.device,
              ctx=ctx,
          )
  
          copy_state_dict(
              sharded_sparse_arch.state_dict(),
              copy.deepcopy(sparse_arch.state_dict()),
          )
  
          unsharded_model_preds = []
          for unsharded_rank in range(ctx.world_size):
              # simulate the unsharded model run on the entire batch
              unsharded_model_preds.append(
                  sparse_arch(kjt_input_per_rank[unsharded_rank])[0]
              )
  
          unsharded_model_pred_this_rank = unsharded_model_preds[
              ctx.rank
          ]
  
          # sharded model
          # each rank gets a subbatch
          sharded_model_pred = sharded_sparse_arch(
              kjt_input_per_rank[ctx.rank]
          )[0]
  
          torch.testing.assert_close(
              sharded_model_pred.cpu(),
              unsharded_model_pred_this_rank.cpu(),
          )
  
          torch.stack(unsharded_model_preds).mean().backward()
          sharded_model_pred.mean().backward()
  
          unsharded_named_parameters = dict(
              sparse_arch.named_parameters()
          )
          sharded_named_parameters = dict(
              sharded_sparse_arch.named_parameters()
          )
  
          for fqn, param in unsharded_named_parameters.items():
              if "_feature_processors" not in fqn:
                  continue
  
              replicated_param = sharded_named_parameters[fqn]
  
              torch.testing.assert_close(
                  # pyre-ignore
                  param.grad.cpu(),
                  replicated_param.grad.cpu(),
                  msg=f"Did not match for {fqn} {param.grad=} {replicated_param.grad=}",
              )
  
          assert (
              sparse_arch.state_dict().keys()
              == sharded_sparse_arch.state_dict().keys()
          ), "State dict keys are not the same"
  
imports: import copy
  import unittest
  from operator import xor
  from typing import List, Optional, Tuple
  import torch
  import torch.nn
  from hypothesis import given, settings, strategies, Verbosity
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torchrec import distributed
  from torchrec.distributed.fp_embeddingbag import FeatureProcessedEmbeddingBagCollectionSharder
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.shard import _shard_modules
  from torchrec.distributed.sharding_plan import column_wise, construct_module_sharding_plan, data_parallel, 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.tests.test_fp_embeddingbag_utils import create_module_and_freeze
  from torchrec.distributed.types import ModuleSharder, ShardingEnv, ShardingPlan
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  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 a global function with example usage: _test_sharding(tables, rank, world_size, kjt_input_per_rank, sharder, backend, set_gradient_division, local_size, use_dmp, use_fp_collection) without return types