class_docstr: Stats for a sharding planner execution.
class_docstr_tok: ["Stats","for","a","sharding","planner","execution","."]
docstr: Gets ranks, input sizes, and output sizes per shard.
  Input size is a function of pooling factor.
  Output size is a function of embedding dimension * number of features.
  
  Returns:
      ranks: list of ranks.
      input_sizes: input size per iter in MB across ranks for given shard.
      output_sizes: output size per iter in MB across ranks for given shard.
docstr_tok: ["Gets","ranks",",","input","sizes",",","and","output","sizes","per","shard",".","Input","size","is","a","function","of","pooling","factor",".","Output","size","is","a","function","of","embedding","dimension","*","number","of","features",".","Returns",":","ranks",":","list","of","ranks",".","input_sizes",":","input","size","per","iter","in","MB","across","ranks","for","given","shard",".","output_sizes",":","output","size","per","iter","in","MB","across","ranks","for","given","shard","."]
code: def _get_shard_stats(
      self,
      shard: ParameterSharding,
      sharding_option: ShardingOption,
      world_size: int,
      local_world_size: int,
      constraints: Optional[Dict[str, ParameterConstraints]] = None,
  ) -> Tuple[List[int], List[float], List[float]]:
      """
      Gets ranks, input sizes, and output sizes per shard.
      Input size is a function of pooling factor.
      Output size is a function of embedding dimension * number of features.
  
      Returns:
          ranks: list of ranks.
          input_sizes: input size per iter in MB across ranks for given shard.
          output_sizes: output size per iter in MB across ranks for given shard.
      """
      assert shard.ranks
      ranks = shard.ranks
  
      num_poolings = (
          cast(
              List[float],
              constraints[sharding_option.name].num_poolings,
          )
          if constraints
          and constraints.get(sharding_option.name)
          and constraints[sharding_option.name].num_poolings
          else [1.0] * sharding_option.num_inputs
      )
      batch_sizes = (
          cast(List[int], constraints[sharding_option.name].batch_sizes)
          if constraints
          and constraints.get(sharding_option.name)
          and constraints[sharding_option.name].batch_sizes
          else [sharding_option.batch_size] * sharding_option.num_inputs
      )
      input_data_type_size = BIGINT_DTYPE
      output_data_type_size = sharding_option.tensor.element_size()
  
      input_sizes, output_sizes = _calculate_shard_io_sizes(
          sharding_type=sharding_option.sharding_type,
          batch_sizes=batch_sizes,
          world_size=world_size,
          local_world_size=local_world_size,
          input_lengths=sharding_option.input_lengths,
          emb_dim=sharding_option.tensor.shape[1],
          shard_sizes=[shard.size for shard in sharding_option.shards],
          input_data_type_size=input_data_type_size,
          output_data_type_size=output_data_type_size,
          num_poolings=num_poolings,
          is_pooled=sharding_option.is_pooled,
      )
  
      input_sizes = [
          bytes_to_mb(input_size) for input_size in input_sizes
      ]
      output_sizes = [
          bytes_to_mb(output_size) for output_size in output_sizes
      ]
  
      return ranks, input_sizes, output_sizes
  
imports: import logging
  from collections import defaultdict
  from typing import Any, cast, Dict, List, Optional, Tuple, Union
  from torch import nn
  from torchrec.distributed.planner.constants import BIGINT_DTYPE, NUM_POOLINGS
  from torchrec.distributed.planner.shard_estimators import _calculate_shard_io_sizes
  from torchrec.distributed.planner.storage_reservations import FixedPercentageStorageReservation, HeuristicalStorageReservation, InferenceStorageReservation
  from torchrec.distributed.planner.types import ParameterConstraints, Perf, ShardingOption, Stats, Storage, StorageReservation, Topology
  from torchrec.distributed.planner.utils import _find_imbalance_tables, bytes_to_gb, bytes_to_mb, sharder_name
  from torchrec.distributed.types import ModuleSharder, ParameterSharding, ShardingPlan, ShardingType
label_desc: Use image node_id 3 for calling the EmbeddingStats obj's underlying member method code with example usage: obj._get_shard_stats(shard, sharding_option, world_size, local_world_size, constraints) and returns: ranks,  input_sizes,  output_sizes