code: def shard_modules(
      module: nn.Module,
      env: Optional[ShardingEnv] = None,
      device: Optional[torch.device] = None,
      plan: Optional[ShardingPlan] = None,
      sharders: Optional[List[ModuleSharder[torch.nn.Module]]] = None,
      init_params: bool = False,
  ) -> nn.Module:
      """
      Replaces all sub_modules that are embedding modules with their sharded variants. This embedding_module -> sharded_embedding_module mapping
      is derived from the passed in sharders.
  
      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 (Optional[ShardingPlan]): plan to use when sharding, defaults to
              `EmbeddingShardingPlanner.collective_plan()`.
          sharders (Optional[List[ModuleSharder[nn.Module]]]): `ModuleSharders` available
              to shard with, defaults to `get_default_sharders()`.
          init_params: (Optional[bool]): If ``True``, will materialize parameters and
              buffers that are on meta device, and will move module to ``device``. Note that
              this only applies if `device.type != "meta"``. Default: `False`.
  
  
      Example::
  
          @torch.no_grad()
          def init_weights(m):
              if isinstance(m, nn.Linear):
                  m.weight.fill_(1.0)
              elif isinstance(m, EmbeddingBagCollection):
                  for param in m.parameters():
                      init.kaiming_normal_(param)
  
          m = MyModel(device='meta')
          m = shard(m)
          assert isinstance(m.embedding_bag_collection, ShardedEmbeddingBagCollection)
      """
  
      torch._C._log_api_usage_once("torchrec.distributed.shard_modules")
      return _shard_modules(
          module, env, device, plan, sharders, init_params
      )
  
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 4 for calling a global function with example usage: shard_modules(module, env, device, plan, sharders, init_params) and returns: _shard_modules