code: def construct_module_sharding_plan(
      module: nn.Module,
      per_param_sharding: Dict[str, ParameterShardingGenerator],
      sharder: Optional[ModuleSharder[nn.Module]] = None,
      local_size: Optional[int] = None,
      world_size: Optional[int] = None,
      device_type: Optional[str] = None,
  ) -> EmbeddingModuleShardingPlan:
      """
      Helper function to create module sharding plans (EmbeddingModuleShardingPlan) for an module
      Args:
          module (nn.Module): module to create plan for.
          per_param_sharding: Dict[str, Callable[[nn.Parameter, int, int, str], ParameterSharding]]: A mapping of parameter names to a generator function
          that takes in [parameter, local_size, world_size, device_type] and returns a ParameterSharding. We recommend using one of the predefined generator functions
          e.g. table_wise_sharding, row_wise_sharding, etc,
          sharder: Optional[ModuleSharder[nn.Module]]: Sharder that we are creating a plan for. If set to none, we will try to derive it from the module. We recommend setting this to None.
          local_size: Optional[int] = None: process group local size
          world_size: Optional[int] = None: process_group world_size
          device_type: str : Torch device type,
  
      Example::
  
          ebc = EmbeddingBagCollection(...)
          plan = construct_module_sharding_plan(
              ebc,
              {
                  "table_0": data_parallel(),
                  "table_1": row_wise(),
                  "table_2": column_wise(),
                  "table_3": column_wise(ranks=[0,1,2]),
                  "table_4": table_row_wise(host_index=2),
              },
          )
      """
      if device_type is None:
          device_type = "cuda" if torch.cuda.is_available() else "cpu"
      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)}"
  
      assert isinstance(
          module, sharder.module_type
      ), f"Incorrect sharder for module type {type(module)}"
      shardable_parameters = sharder.shardable_parameters(module)
      assert (
          shardable_parameters.keys() == per_param_sharding.keys()
      ), "per_param_sharding_config doesn't match the shardable parameters of the module"
  
      local_size = local_size or get_local_size()
      world_size = world_size or dist.get_world_size()
  
      per_parameter_sharding = EmbeddingModuleShardingPlan()
      for (
          table_name,
          sharding_plan_generator,
      ) in per_param_sharding.items():
          param = shardable_parameters[table_name]
          per_parameter_sharding[table_name] = sharding_plan_generator(
              param, local_size, world_size, device_type, sharder
          )
  
      return per_parameter_sharding
  
imports: import math
  import warnings
  from typing import Callable, cast, Dict, List, Optional, Tuple, Type
  import torch
  from torch import distributed, nn
  from torchrec.distributed.comm import get_local_size
  from torchrec.distributed.embedding import EmbeddingCollectionSharder
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder
  from torchrec.distributed.fp_embeddingbag import FeatureProcessedEmbeddingBagCollectionSharder
  from torchrec.distributed.fused_embeddingbag import FusedEmbeddingBagCollectionSharder
  from torchrec.distributed.mc_embedding import ManagedCollisionEmbeddingCollectionSharder
  from torchrec.distributed.mc_embeddingbag import ManagedCollisionEmbeddingBagCollectionSharder
  from torchrec.distributed.planner.constants import MIN_CW_DIM
  from torchrec.distributed.quant_embedding import QuantEmbeddingCollectionSharder
  from torchrec.distributed.quant_embeddingbag import QuantEmbeddingBagCollectionSharder
  from torchrec.distributed.types import EmbeddingModuleShardingPlan, EnumerableShardingSpec, ModuleSharder, ParameterSharding, ShardingType, ShardMetadata
  from torchrec.distributed.utils import none_throws
label_desc: Use image node_id 17 for calling a global function with example usage: construct_module_sharding_plan(module, per_param_sharding, sharder, local_size, world_size, device_type) and returns: per_parameter_sharding