code: def _test_sharding(  # noqa C901
      tables: List[EmbeddingBagConfig],
      initial_state_dict: Dict[str, Any],
      rank: int,
      world_size: int,
      kjt_input_per_rank: List[KeyedJaggedTensor],
      sharder: ModuleSharder[nn.Module],
      backend: str,
      constraints: Optional[Dict[str, ParameterConstraints]] = None,
      local_size: Optional[int] = None,
      is_data_parallel: bool = False,
      use_apply_optimizer_in_backward: bool = False,
  ) -> 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()
          }
  
          planner = EmbeddingShardingPlanner(
              topology=Topology(
                  world_size,
                  ctx.device.type,
                  local_world_size=ctx.local_size,
              ),
              constraints=constraints,
          )
          model = EmbeddingBagCollection(
              tables=tables,
              device=ctx.device,
          )
          unsharded_model = EmbeddingBagCollection(
              tables=tables,
              device=ctx.device,
          )
  
          if use_apply_optimizer_in_backward:
              apply_optimizer_in_backward(
                  torch.optim.SGD,
                  model.embedding_bags["table_0"].parameters(),
                  {"lr": 1.0},
              )
              apply_optimizer_in_backward(
                  torch.optim.SGD,
                  model.embedding_bags["table_1"].parameters(),
                  {"lr": 4.0},
              )
              apply_optimizer_in_backward(
                  torch.optim.SGD,
                  unsharded_model.embedding_bags[
                      "table_0"
                  ].parameters(),
                  {"lr": 1.0},
              )
              apply_optimizer_in_backward(
                  torch.optim.SGD,
                  unsharded_model.embedding_bags[
                      "table_1"
                  ].parameters(),
                  {"lr": 4.0},
              )
          plan: ShardingPlan = planner.collective_plan(
              model, [sharder], ctx.pg
          )
          sharded_model = shard(
              module=model,
              env=ShardingEnv.from_process_group(ctx.pg),
              plan=plan.get_plan_for_module(""),
              sharder=sharder,
              device=ctx.device,
          )
  
          if not use_apply_optimizer_in_backward:
              unsharded_model_optimizer = torch.optim.SGD(
                  unsharded_model.parameters(), lr=0.01
              )
              sharded_model_optimizer = torch.optim.SGD(
                  sharded_model.parameters(), lr=0.01
              )
  
          assert isinstance(
              sharded_model, ShardedEmbeddingBagCollection
          )
  
          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)
  
          for _it in range(5):
              unsharded_model_params = dict(
                  unsharded_model.named_parameters()
              )
  
              if not use_apply_optimizer_in_backward:
                  unsharded_model_optimizer.zero_grad()
                  sharded_model_optimizer.zero_grad()
  
              if is_data_parallel:
                  for fqn, param in sharded_model.named_parameters():
                      assert _optional_equals(
                          param.grad, unsharded_model_params[fqn].grad
                      )
  
              unsharded_model_pred_kt = []
              for unsharded_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[unsharded_rank]
                      )
                  )
  
              all_unsharded_preds = []
              for unsharded_rank in range(ctx.world_size):
                  unsharded_model_pred_kt_mini_batch = (
                      unsharded_model_pred_kt[unsharded_rank].to_dict()
                  )
  
                  all_unsharded_preds.extend(
                      [
                          unsharded_model_pred_kt_mini_batch[feature]
                          for feature in feature_keys
                      ]
                  )
                  if unsharded_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)
              _sum = all_unsharded_preds.sum()
              if is_data_parallel:
                  _sum /= world_size
              _sum.backward()
  
              if is_data_parallel:
                  for fqn, param in sharded_model.named_parameters():
                      assert _optional_equals(
                          param.grad, unsharded_model_params[fqn].grad
                      )
  
              if not use_apply_optimizer_in_backward:
                  unsharded_model_optimizer.step()
                  sharded_model_optimizer.step()
  
          # check nn.Module APIs look the same
          assert_state_buffers_parameters_equal(
              unsharded_model, sharded_model
          )
  
          for fqn in unsharded_model.state_dict():
              unsharded_state = unsharded_model.state_dict()[fqn]
              sharded_state = sharded_model.state_dict()[fqn]
              if is_data_parallel:
                  torch.testing.assert_close(
                      unsharded_state, sharded_state
                  )
              else:
                  out = (
                      torch.zeros(
                          size=unsharded_state.shape, device=ctx.device
                      )
                      if ctx.rank == 0
                      else None
                  )
                  sharded_state.gather(out=out)
                  if ctx.rank == 0:
                      torch.testing.assert_close(
                          unsharded_state,
                          out,
                      )
  
imports: import copy
  import unittest
  from functools import partial
  from typing import Any, Dict, List, Optional
  import hypothesis.strategies
  import torch
  import torch.nn
  from hypothesis import assume, given, settings, Verbosity
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torchrec import distributed
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder, ShardedEmbeddingBagCollection
  from torchrec.distributed.planner import EmbeddingShardingPlanner, ParameterConstraints, Topology
  from torchrec.distributed.shard import shard
  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 BoundsCheckMode, CacheAlgorithm, CacheParams, DataType, ModuleSharder, QuantizedCommCodecs, ShardingEnv, ShardingPlan, ShardingType
  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 assert_state_buffers_parameters_equal, skip_if_asan_class
label_desc: Use image node_id 2 for calling a global function with example usage: _test_sharding(tables, initial_state_dict, rank, world_size, kjt_input_per_rank, sharder, backend, constraints, local_size, is_data_parallel, use_apply_optimizer_in_backward) without return types