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:
      """
      See shard
      """
      if sharder is None:
          sharder = get_module_to_default_sharders().get(
              type(module), None
          )
      assert (
          sharder is not None
      ), f"Could not find a valid sharder type for {type(module)}"
  
      if env is None:
          pg = dist.GroupMember.WORLD
          assert pg is not None, "Process group is not initialized"
          env = ShardingEnv.from_process_group(pg)
  
      if device is None:
          if torch.cuda.is_available():
              device = torch.device(torch.cuda.current_device())
          else:
              device = torch.device("cpu")
  
      if isinstance(plan, ModuleShardingPlan):
          return sharder.shard(module, plan, env, device)
  
      # Run sharding generators.
      shardable_parameters = sharder.shardable_parameters(module)
      if isinstance(plan, Callable):
          gen = plan
          plan = {}
          for table_name, param in shardable_parameters.items():
              plan[table_name] = gen(
                  param,
                  get_local_size(env.world_size),
                  env.world_size,
                  device.type,
                  sharder,
              )
      else:
          for table_name, sharding in plan.items():
              if isinstance(sharding, Callable):
                  param = shardable_parameters[table_name]
                  # pyre-fixme[6]: For 2nd argument expected `(Parameter, int, int,
                  #  str, ModuleSharder[Module]) -> ParameterSharding` but got
                  #  `ParameterSharding`.
                  plan[table_name] = sharding(
                      param,
                      get_local_size(env.world_size),
                      env.world_size,
                      device.type,
                      sharder,
                  )
  
      return sharder.shard(module, plan, env, device)
  
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 3 for calling a global function with example usage: _shard(module, plan, env, device, sharder) and returns: sharder, sharder