code: def _test_sharding(  # noqa C901
      tables: List[EmbeddingConfig],
      rank: int,
      world_size: int,
      kjt_input_per_rank: List[KeyedJaggedTensor],
      backend: str,
      local_size: Optional[int] = None,
      use_apply_optimizer_in_backward: bool = False,
      use_index_dedup: bool = False,
  ) -> None:
      trec_dist.comm_ops.set_gradient_division(False)
      with MultiProcessContext(
          rank, world_size, backend, local_size
      ) as ctx:
          sharder = EmbeddingCollectionSharder(
              use_index_dedup=use_index_dedup
          )
          kjt_input_per_rank = [
              kjt.to(ctx.device) for kjt in kjt_input_per_rank
          ]
  
          unsharded_model = EmbeddingCollection(
              tables=tables,
              device=ctx.device,
              need_indices=True,
          )
  
          # syncs model across ranks
          torch.manual_seed(0)
          for param in unsharded_model.parameters():
              nn.init.uniform_(param, -1, 1)
          torch.manual_seed(0)
  
          if use_apply_optimizer_in_backward:
              apply_optimizer_in_backward(
                  torch.optim.SGD,
                  unsharded_model.embeddings.parameters(),
                  {"lr": 1.0},
              )
          else:
              unsharded_model_optimizer = torch.optim.SGD(
                  unsharded_model.parameters(), lr=1.0
              )
  
          module_sharding_plan = construct_module_sharding_plan(
              unsharded_model,
              per_param_sharding={
                  "table_0": table_wise(rank=0),
                  "table_1": row_wise(),
                  "table_2": column_wise(ranks=[0, 1]),
              },
              local_size=local_size,
              world_size=world_size,
              device_type=ctx.device.type,
              # pyre-ignore
              sharder=sharder,
          )
  
          sharded_model = _shard_modules(
              module=unsharded_model,
              plan=ShardingPlan({"": module_sharding_plan}),
              env=ShardingEnv.from_process_group(ctx.pg),
              # pyre-ignore
              sharders=[sharder],
              device=ctx.device,
          )
  
          if not use_apply_optimizer_in_backward:
              sharded_model_optimizer = torch.optim.SGD(
                  sharded_model.parameters(), lr=1.0
              )
  
          assert isinstance(sharded_model, ShardedEmbeddingCollection)
  
          if not use_apply_optimizer_in_backward:
              unsharded_model_optimizer.zero_grad()
              sharded_model_optimizer.zero_grad()
  
          unsharded_model_pred_jt_dict = []
          for unsharded_rank in range(ctx.world_size):
              # simulate the unsharded model run on the entire batch
              unsharded_model_pred_jt_dict.append(
                  unsharded_model(kjt_input_per_rank[unsharded_rank])
              )
  
          # sharded model
          # each rank gets a subbatch
          sharded_model_pred_jts_dict: Dict[
              str, JaggedTensor
          ] = sharded_model(kjt_input_per_rank[ctx.rank])
  
          unsharded_model_pred_jt_dict_this_rank: Dict[
              str, JaggedTensor
          ] = unsharded_model_pred_jt_dict[ctx.rank]
  
          embedding_names = (
              unsharded_model_pred_jt_dict_this_rank.keys()
          )
          assert set(
              unsharded_model_pred_jt_dict_this_rank.keys()
          ) == set(sharded_model_pred_jts_dict.keys())
  
          unsharded_loss = []
          sharded_loss = []
          for embedding_name in embedding_names:
              unsharded_jt = unsharded_model_pred_jt_dict_this_rank[
                  embedding_name
              ]
              sharded_jt = sharded_model_pred_jts_dict[embedding_name]
  
              torch.testing.assert_close(
                  unsharded_jt.values(), sharded_jt.values()
              )
              torch.testing.assert_close(
                  unsharded_jt.lengths(), sharded_jt.lengths()
              )
              torch.testing.assert_close(
                  unsharded_jt.offsets(), sharded_jt.offsets()
              )
              torch.testing.assert_close(
                  unsharded_jt.weights_or_none(),
                  sharded_jt.weights_or_none(),
              )
  
              sharded_loss.append(sharded_jt.values().view(-1))
  
          for rank in range(ctx.world_size):
              for embedding_name in embedding_names:
                  unsharded_loss.append(
                      unsharded_model_pred_jt_dict[rank][embedding_name]
                      .values()
                      .view(-1)
                  )
  
          torch.cat(sharded_loss).sum().backward()
          torch.cat(unsharded_loss).sum().backward()
  
          if not use_apply_optimizer_in_backward:
              unsharded_model_optimizer.step()
              sharded_model_optimizer.step()
  
          for fqn in unsharded_model.state_dict():
              unsharded_state = unsharded_model.state_dict()[fqn]
              sharded_state = sharded_model.state_dict()[fqn]
  
              sharded_param = (
                  torch.zeros(
                      size=unsharded_state.shape, device=ctx.device
                  )
                  if ctx.rank == 0
                  else None
              )
              if isinstance(sharded_state, ShardedTensor):
                  sharded_state.gather(out=sharded_param)
              else:
                  sharded_param = sharded_state
  
              if ctx.rank == 0:
                  torch.testing.assert_close(
                      unsharded_state,
                      sharded_param,
                      msg=f"Did not match for {fqn=} after backward",
                  )
  
imports: import unittest
  from typing import Dict, List, Optional
  import hypothesis.strategies
  import torch
  import torch.nn
  from hypothesis import given, settings, Verbosity
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torchrec import distributed
  from torchrec.distributed.embedding import EmbeddingCollectionSharder, ShardedEmbeddingCollection
  from torchrec.distributed.shard import _shard_modules
  from torchrec.distributed.sharding_plan import column_wise, construct_module_sharding_plan, row_wise, table_wise
  from torchrec.distributed.test_utils.multi_process import MultiProcessContext, MultiProcessTestBase
  from torchrec.distributed.types import ShardedTensor, ShardingEnv, ShardingPlan
  from torchrec.modules.embedding_configs import EmbeddingConfig
  from torchrec.modules.embedding_modules import EmbeddingCollection
  from torchrec.sparse.jagged_tensor import JaggedTensor, 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, rank, world_size, kjt_input_per_rank, backend, local_size, use_apply_optimizer_in_backward, use_index_dedup) without return types