class_docstr: Provides an optimized sharding plan for a given module with shardable parameters
  according to the provided sharders, topology, and constraints.
class_docstr_tok: ["Provides","an","optimized","sharding","plan","for","a","given","module","with","shardable","parameters","according","to","the","provided","sharders",",","topology",",","and","constraints","."]
docstr: Call self.plan(...) on rank 0 and broadcast
docstr_tok: ["Call","self.plan","(","...",")","on","rank","0","and","broadcast"]
code: def collective_plan(
      self,
      module: nn.Module,
      sharders: Optional[List[ModuleSharder[nn.Module]]] = None,
      pg: Optional[dist.ProcessGroup] = dist.GroupMember.WORLD,
  ) -> ShardingPlan:
      """
      Call self.plan(...) on rank 0 and broadcast
      """
      assert pg is not None, "Process group is not initialized"
      if sharders is None:
          sharders = get_default_sharders()
      return invoke_on_rank_and_broadcast_result(
          pg,
          0,
          self.plan,
          module,
          sharders,
      )
  
imports: import copy
  from functools import reduce
  from time import perf_counter
  from typing import cast, Dict, List, Optional, Tuple, Union
  import torch
  import torch.distributed
  from torch import nn
  from torchrec.distributed.collective_utils import invoke_on_rank_and_broadcast_result
  from torchrec.distributed.comm import get_local_size
  from torchrec.distributed.planner.constants import BATCH_SIZE, MAX_SIZE
  from torchrec.distributed.planner.enumerators import EmbeddingEnumerator
  from torchrec.distributed.planner.partitioners import GreedyPerfPartitioner
  from torchrec.distributed.planner.perf_models import NoopPerfModel
  from torchrec.distributed.planner.proposers import GreedyProposer, GridSearchProposer, UniformProposer
  from torchrec.distributed.planner.stats import EmbeddingStats
  from torchrec.distributed.planner.storage_reservations import HeuristicalStorageReservation
  from torchrec.distributed.planner.types import Enumerator, ParameterConstraints, Partitioner, PerfModel, PlannerError, PlannerErrorType, Proposer, ShardingOption, Stats, Storage, StorageReservation, Topology
  from torchrec.distributed.planner.utils import bytes_to_gb, reset_shard_rank, storage_repr_in_gb
  from torchrec.distributed.sharding_plan import get_default_sharders, placement
  from torchrec.distributed.types import EmbeddingModuleShardingPlan, EnumerableShardingSpec, ModuleSharder, ParameterSharding, ShardingPlan, ShardingPlanner, ShardingType, ShardMetadata
label_desc: Use image node_id 2 for calling the EmbeddingShardingPlanner obj's underlying member method code with example usage: obj.collective_plan(module, sharders, pg) and returns: invoke_on_rank_and_broadcast_result