code: def _calculate_shard_io_sizes(
      sharding_type: str,
      batch_sizes: List[int],
      world_size: int,
      local_world_size: int,
      input_lengths: List[float],
      emb_dim: int,
      shard_sizes: List[List[int]],
      input_data_type_size: int,
      output_data_type_size: int,
      num_poolings: List[float],
      is_pooled: bool,
  ) -> Tuple[List[int], List[int]]:
      if sharding_type == ShardingType.DATA_PARALLEL.value:
          return _calculate_dp_shard_io_sizes(
              batch_sizes=batch_sizes,
              input_lengths=input_lengths,
              emb_dim=emb_dim,
              num_shards=len(shard_sizes),
              input_data_type_size=input_data_type_size,
              output_data_type_size=output_data_type_size,
              num_poolings=num_poolings,
              is_pooled=is_pooled,
          )
      elif sharding_type == ShardingType.TABLE_WISE.value:
          return _calculate_tw_shard_io_sizes(
              batch_sizes=batch_sizes,
              world_size=world_size,
              input_lengths=input_lengths,
              emb_dim=emb_dim,
              input_data_type_size=input_data_type_size,
              output_data_type_size=output_data_type_size,
              num_poolings=num_poolings,
              is_pooled=is_pooled,
          )
      elif sharding_type in {
          ShardingType.COLUMN_WISE.value,
          ShardingType.TABLE_COLUMN_WISE.value,
      }:
          return _calculate_cw_shard_io_sizes(
              batch_sizes=batch_sizes,
              world_size=world_size,
              input_lengths=input_lengths,
              shard_sizes=shard_sizes,
              input_data_type_size=input_data_type_size,
              output_data_type_size=output_data_type_size,
              num_poolings=num_poolings,
              is_pooled=is_pooled,
          )
      elif sharding_type == ShardingType.ROW_WISE.value:
          return _calculate_rw_shard_io_sizes(
              batch_sizes=batch_sizes,
              world_size=world_size,
              input_lengths=input_lengths,
              shard_sizes=shard_sizes,
              input_data_type_size=input_data_type_size,
              output_data_type_size=output_data_type_size,
              num_poolings=num_poolings,
              is_pooled=is_pooled,
          )
      elif sharding_type == ShardingType.TABLE_ROW_WISE.value:
          return _calculate_twrw_shard_io_sizes(
              batch_sizes=batch_sizes,
              world_size=world_size,
              local_world_size=local_world_size,
              input_lengths=input_lengths,
              shard_sizes=shard_sizes,
              input_data_type_size=input_data_type_size,
              output_data_type_size=output_data_type_size,
              num_poolings=num_poolings,
              is_pooled=is_pooled,
          )
      else:
          raise ValueError(
              f"Unrecognized or unsupported sharding type provided: {sharding_type}"
          )
  
imports: import logging
  import math
  from typing import cast, Dict, List, Optional, Tuple, Type
  import torch
  import torchrec.optim
  from torch import nn
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel
  from torchrec.distributed.planner.constants import BATCHED_COPY_PERF_FACTOR, BIGINT_DTYPE, DP_ELEMENTWISE_KERNELS_PERF_FACTOR, FULL_BLOCK_EMB_DIM, HALF_BLOCK_PENALTY, kernel_bw_lookup, QUARTER_BLOCK_PENALTY, UVM_CACHING_RATIO, WEIGHTED_KERNEL_MULTIPLIER
  from torchrec.distributed.planner.types import ParameterConstraints, Perf, PlannerError, ShardEstimator, ShardingOption, Storage, Topology
  from torchrec.distributed.planner.utils import prod, sharder_name
  from torchrec.distributed.types import CacheStatistics, CommOp, ModuleSharder, ShardingType
  from torchrec.modules.embedding_modules import EmbeddingBagCollectionInterface
label_desc: Use image node_id 9 for calling a global function with example usage: _calculate_shard_io_sizes(sharding_type, batch_sizes, world_size, local_world_size, input_lengths, emb_dim, shard_sizes, input_data_type_size, output_data_type_size, num_poolings, is_pooled) and returns: _calculate_dp_shard_io_sizes, _calculate_tw_shard_io_sizes, _calculate_cw_shard_io_sizes, _calculate_rw_shard_io_sizes, _calculate_twrw_shard_io_sizes