code: def _get_rw_sharding_perf(
      batch_sizes: List[int],
      world_size: int,
      local_world_size: int,
      input_lengths: List[float],
      emb_dim: int,
      input_data_type_size: float,
      table_data_type_size: float,
      fwd_a2a_comm_data_type_size: float,
      bwd_a2a_comm_data_type_size: float,
      fwd_sr_comm_data_type_size: float,
      bwd_sr_comm_data_type_size: float,
      num_poolings: List[float],
      ddr_mem_bw: float,
      device_bw: float,
      inter_host_bw: float,
      intra_host_bw: float,
      bwd_compute_multiplier: float,
      is_pooled: bool,
      is_weighted: bool = False,
      expected_cache_fetches: float = 0,
  ) -> Perf:
      batch_inputs = (
          sum(
              [
                  x * y * z
                  for x, y, z in zip(
                      input_lengths, num_poolings, batch_sizes
                  )
              ]
          )
          / world_size
      )
      batch_outputs = (
          sum([x * y for x, y in zip(num_poolings, batch_sizes)])
          if is_pooled
          else batch_inputs
      )
  
      input_read_size = math.ceil(
          batch_inputs * world_size * input_data_type_size
      )
      if is_weighted:
          input_read_size *= 2
  
      embedding_lookup_size = (
          batch_inputs * world_size * emb_dim * table_data_type_size
      )
  
      fwd_output_write_size = (
          batch_outputs
          * world_size
          * emb_dim
          * fwd_sr_comm_data_type_size
          if is_pooled
          else batch_outputs
          * world_size
          * emb_dim
          * fwd_a2a_comm_data_type_size
      )
      bwd_output_write_size = (
          batch_outputs
          * world_size
          * emb_dim
          * bwd_sr_comm_data_type_size
          if is_pooled
          else batch_outputs
          * world_size
          * emb_dim
          * bwd_a2a_comm_data_type_size
      )
  
      comms_bw = (
          inter_host_bw
          if world_size > local_world_size
          else intra_host_bw
      )
      fwd_comms = fwd_output_write_size / comms_bw
  
      fwd_compute = (
          input_read_size
          + embedding_lookup_size
          + fwd_output_write_size
      ) / device_bw
  
      bwd_comms = bwd_output_write_size / comms_bw
  
      bwd_batched_copy = (
          bwd_output_write_size * BATCHED_COPY_PERF_FACTOR / device_bw
      )
  
      bwd_grad_indice_weights_kernel = (
          fwd_compute * WEIGHTED_KERNEL_MULTIPLIER if is_weighted else 0
      )
  
      bwd_compute = fwd_compute * bwd_compute_multiplier
  
      # for row-wise, expected_cache_fetches per shard is / world_size
      prefetch_compute = _get_expected_cache_prefetch_time(
          ddr_mem_bw,
          expected_cache_fetches / world_size,
          emb_dim,
          table_data_type_size,
      )
  
      return Perf(
          fwd_compute=fwd_compute,
          fwd_comms=fwd_comms,
          bwd_compute=bwd_compute + bwd_grad_indice_weights_kernel,
          bwd_comms=bwd_comms + bwd_batched_copy,
          prefetch_compute=prefetch_compute,
      )
  
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 4 for calling a global function with example usage: _get_rw_sharding_perf(batch_sizes, world_size, local_world_size, input_lengths, emb_dim, input_data_type_size, table_data_type_size, fwd_a2a_comm_data_type_size, bwd_a2a_comm_data_type_size, fwd_sr_comm_data_type_size, bwd_sr_comm_data_type_size, num_poolings, ddr_mem_bw, device_bw, inter_host_bw, intra_host_bw, bwd_compute_multiplier, is_pooled, is_weighted, expected_cache_fetches) and returns: Perf