code: def shard(
      module: nn.Module,
      plan: Union[
          ModuleShardingPlan,
          Dict[str, ParameterShardingGenerator],
          ParameterShardingGenerator,
      ],
      env: Optional[ShardingEnv] = None,
      device: Optional[torch.device] = None,
      sharder: Optional[ModuleSharder[nn.Module]] = None,
  ) -> nn.Module:
      """
      Replaces this module with its sharded variant
  
      This will leave the other parts of the model unaffected.
  
      It returns the original module
  
      Args:
          module (nn.Module): module to wrap.
          env (Optional[ShardingEnv]): sharding environment that has the process group.
          device (Optional[torch.device]): compute device, defaults to cpu.
          plan (Union[ModuleShardingPlan, Dict[str, ParameterShardingGenerator], ParameterShardingGenerator]):
              dict of ParameterSharding (materized plan) or ParameterShardingGenerator (which will be run to produce ParameterSharding).
              If single ParameterShardingGenerator is supplied, it will be applied to all module parameters.
          sharder (Optional[List[ModuleSharder[nn.Module]]]): sharder to use, default is picked from `get_default_sharders()`.
  
      Example:
  
          ebc = EmbeddingBagCollection()
          sharded_ebc = shard(ebc, table_row_wise(host_index=0))
          assert isinstance(sharded_ebc, ShardedEmbeddingBagCollection)
      """
      torch._C._log_api_usage_once("torchrec.distributed.shard")
      return _shard(module, plan, env, device, sharder)
  
file_path: torchrec/torchrec/distributed/shard.py
imports: from typing import Callable, Dict, List, Optional, Type, Union
  import torch
  import torch.distributed
  from torch import nn
  from torch.distributed._composable.contract import contract
  from torchrec.distributed.comm import get_local_size
  from torchrec.distributed.model_parallel import get_default_sharders
  from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
  from torchrec.distributed.sharding_plan import get_module_to_default_sharders, ParameterShardingGenerator
  from torchrec.distributed.types import ModuleSharder, ModuleShardingPlan, ShardingEnv, ShardingPlan
  from torchrec.distributed.utils import init_parameters
label_desc: Use image node_id 2 for calling a global function with example usage: shard(module, plan, env, device, sharder) and returns: _shard