code: def get_unsharded_and_sharded_module(
      tables: List[EmbeddingBagConfig],
      sharder: ModuleSharder[nn.Module],
      use_dmp: bool,
      use_fp_collection: bool,
      init_device: torch.device,
      ctx: MultiProcessContext,
  ) -> Tuple[nn.Module, nn.Module]:
      sparse_arch = create_module_and_freeze(
          tables,
          use_fp_collection=use_fp_collection,
          device=init_device,
      )
  
      apply_optimizer_in_backward(
          torch.optim.SGD,
          sparse_arch._fp_ebc._embedding_bag_collection.embedding_bags.parameters(),
          {"lr": 1.0},
      )
  
      module_sharding_plan = construct_module_sharding_plan(
          sparse_arch._fp_ebc,
          per_param_sharding={
              "table_0": column_wise(ranks=[0, 1]),
              "table_1": table_wise(rank=0),
              "table_2": data_parallel(),
              "table_3": column_wise(ranks=[0, 0, 1]),
          },
          local_size=ctx.local_size,
          world_size=ctx.world_size,
          device_type="cuda" if torch.cuda.is_available() else "cpu",
          sharder=sharder,
      )
  
      if use_dmp:
          sharded_sparse_arch = DistributedModelParallel(
              module=copy.deepcopy(sparse_arch),
              plan=ShardingPlan({"_fp_ebc": module_sharding_plan}),
              # pyre-ignore
              env=ShardingEnv.from_process_group(ctx.pg),
              sharders=[sharder],
              device=ctx.device,
          )
      else:
          sharded_sparse_arch = _shard_modules(
              module=copy.deepcopy(sparse_arch),
              plan=ShardingPlan({"._fp_ebc": module_sharding_plan}),
              # pyre-ignore
              env=ShardingEnv.from_process_group(ctx.pg),
              sharders=[sharder],
              device=ctx.device,
          )
          from torch.distributed._composable.replicate import replicate
  
          replicate(
              sharded_sparse_arch,
              ignored_modules=[
                  sharded_sparse_arch._fp_ebc._embedding_bag_collection
              ],
              process_group=ctx.pg,
              gradient_as_bucket_view=True,
              device_ids=None
              if ctx.device.type == "cpu"
              else [ctx.device],
              broadcast_buffers=False,
          )
      return sparse_arch, sharded_sparse_arch
  
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 1 for calling a global function with example usage: get_unsharded_and_sharded_module(tables, sharder, use_dmp, use_fp_collection, init_device, ctx) and returns: sparse_arch,  sharded_sparse_arch