code: def _test_sharded_fused_optimizer_state_dict(
      tables: List[EmbeddingBagConfig],
      rank: int,
      local_size: int,
      world_size: int,
      backend: str,
  ) -> None:
      with MultiProcessContext(
          rank, world_size, backend, local_size
      ) as ctx:
          ebc = EmbeddingBagCollection(
              tables=tables, device=torch.device("meta")
          )
          apply_optimizer_in_backward(
              RowWiseAdagrad,
              [
                  ebc.embedding_bags["table_0"].weight,
                  ebc.embedding_bags["table_1"].weight,
              ],
              {"lr": 0.01},
          )
          apply_optimizer_in_backward(
              PartialRowWiseAdam,
              [
                  ebc.embedding_bags["table_2"].weight,
                  ebc.embedding_bags["table_3"].weight,
              ],
              {"lr": 0.02},
          )
  
          parameter_sharding_plan = construct_module_sharding_plan(
              ebc,
              per_param_sharding={
                  "table_0": column_wise(ranks=[0, 0, 1, 1]),
                  "table_1": row_wise(),
                  "table_2": column_wise(ranks=[0, 1, 0, 1]),
                  "table_3": row_wise(),
              },
              world_size=ctx.world_size,
              local_size=ctx.local_size,
              device_type=ctx.device.type,
          )
  
          ebc = shard(
              module=ebc,
              plan=parameter_sharding_plan,
              device=ctx.device,
          )
  
          ebc.embedding_bags["table_0"].weight._in_backward_optimizers[
              0
          ].state_dict()["state"][""]["table_0.momentum1"].gather(
              dst=0,
              out=None if ctx.rank != 0
              # sharded column, each shard will have rowwise state
              else torch.empty(
                  (4 * tables[0].num_embeddings,), device=ctx.device
              ),
          )
  
          ebc.embedding_bags["table_1"].weight._in_backward_optimizers[
              0
          ].state_dict()["state"][""]["table_1.momentum1"].gather(
              dst=0,
              out=None if ctx.rank != 0
              # sharded rowwise
              else torch.empty(
                  (tables[1].num_embeddings,), device=ctx.device
              ),
          )
  
          ebc.embedding_bags["table_2"].weight._in_backward_optimizers[
              0
          ].state_dict()["state"][""]["table_2.momentum1"].gather(
              dst=0,
              out=None if ctx.rank != 0
              # Column wise - with partial rowwise adam, first state is point wise
              else torch.empty(
                  (tables[2].num_embeddings, tables[2].embedding_dim),
                  device=ctx.device,
              ),
          )
  
          ebc.embedding_bags["table_2"].weight._in_backward_optimizers[
              0
          ].state_dict()["state"][""]["table_2.exp_avg_sq"].gather(
              dst=0,
              out=None if ctx.rank != 0
              # Column wise - with partial rowwise adam, first state is point wise
              else torch.empty(
                  (4 * tables[2].num_embeddings,), device=ctx.device
              ),
          )
  
          ebc.embedding_bags["table_3"].weight._in_backward_optimizers[
              0
          ].state_dict()["state"][""]["table_3.momentum1"].gather(
              dst=0,
              out=None if ctx.rank != 0
              # Row wise - with partial rowwise adam, first state is point wise
              else torch.empty(
                  (tables[3].num_embeddings, tables[3].embedding_dim),
                  device=ctx.device,
              ),
          )
  
          ebc.embedding_bags["table_3"].weight._in_backward_optimizers[
              0
          ].state_dict()["state"][""]["table_3.exp_avg_sq"].gather(
              dst=0,
              out=None if ctx.rank != 0
              # Column wise - with partial rowwise adam, first state is point wise
              else torch.empty(
                  (tables[2].num_embeddings,), device=ctx.device
              ),
          )
  
file_path: torchrec/torchrec/distributed/composable/tests/test_fused_optim_nccl.py
imports: import unittest
  from typing import List
  import torch
  from torchrec.distributed.shard import shard
  from torchrec.distributed.sharding_plan import column_wise, construct_module_sharding_plan, row_wise
  from torchrec.distributed.test_utils.multi_process import MultiProcessContext, MultiProcessTestBase
  from torchrec.modules.embedding_configs import DataType, EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.optim.apply_optimizer_in_backward import apply_optimizer_in_backward
  from torchrec.optim.optimizers import PartialRowWiseAdam
  from torchrec.optim.rowwise_adagrad import RowWiseAdagrad
label_desc: Use image node_id 1 for calling the ShardedFusedOptimizerStateDictTest obj's underlying member method code with example usage: obj._test_sharded_fused_optimizer_state_dict(tables, rank, local_size, world_size, backend) without return types