code: def perf_func_emb_wall_time(
      shard_sizes: List[List[int]],
      compute_kernel: str,
      compute_device: str,
      sharding_type: str,
      batch_sizes: List[int],
      world_size: int,
      local_world_size: int,
      input_lengths: List[float],
      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],
      hbm_mem_bw: float,
      ddr_mem_bw: float,
      intra_host_bw: float,
      inter_host_bw: float,
      bwd_compute_multiplier: float,
      is_pooled: bool,
      is_weighted: bool = False,
      caching_ratio: Optional[float] = None,
      is_inference: bool = False,
      prefetch_pipeline: bool = False,
      expected_cache_fetches: float = 0,
  ) -> List[Perf]:
      """
      Attempts to model perfs as a function of relative wall times.
  
      Args:
          shard_sizes (List[List[int]]): the list of (local_rows, local_cols) of each
              shard.
          compute_kernel (str): compute kernel.
          compute_device (str): compute device.
          sharding_type (str): tw, rw, cw, twrw, dp.
          batch_sizes (List[int]): batch size for each input feature.
          world_size (int): the number of devices for all hosts.
          local_world_size (int): the number of the device for each host.
          input_lengths (List[float]): the list of the average number of lookups of each
              input query feature.
          input_data_type_size (float): the data type size of the distributed
              data_parallel input.
          table_data_type_size (float): the data type size of the table.
          fwd_comm_data_type_size (float): the data type size of the distributed
              data_parallel input during forward communication.
          bwd_comm_data_type_size (float): the data type size of the distributed
              data_parallel input during backward communication.
          num_poolings (List[float]): number of poolings per sample, typically 1.0.
          hbm_mem_bw (float): the bandwidth of the device HBM.
          ddr_mem_bw (float): the bandwidth of the system DDR memory.
          intra_host_bw (float): the bandwidth within a single host like multiple threads.
          inter_host_bw (float): the bandwidth between two hosts like multiple machines.
          is_pooled (bool): True if embedding output is pooled (ie. `EmbeddingBag`), False
              if unpooled/sequential (ie. `Embedding`).
          is_weighted (bool = False): if the module is an EBC and is weighted, typically
              signifying an id score list feature.
          is_inference (bool = False): if planning for inference.
          caching_ratio (Optional[float] = None): cache ratio to determine the bandwidth
              of device.
          prefetch_pipeline (bool = False): whether prefetch pipeline is enabled.
          expected_cache_fetches (float): number of expected cache fetches across global batch
  
      Returns:
          List[float]: the list of perf for each shard.
      """
  
      shard_perfs = []
      device_bw = kernel_bw_lookup(
          compute_device,
          compute_kernel,
          hbm_mem_bw,
          ddr_mem_bw,
          caching_ratio,
          prefetch_pipeline,
      )
      if device_bw is None:
          raise PlannerError(
              f"No kernel bandwidth exists for this combo of compute device: {compute_device}, compute kernel: {compute_kernel}"
          )
  
      for hash_size, emb_dim in shard_sizes:
          if (
              sharding_type == ShardingType.TABLE_WISE.value
              or sharding_type == ShardingType.COLUMN_WISE.value
              or sharding_type == ShardingType.TABLE_COLUMN_WISE.value
          ):
              shard_perf = _get_tw_sharding_perf(
                  batch_sizes=batch_sizes,
                  world_size=world_size,
                  local_world_size=local_world_size,
                  input_lengths=input_lengths,
                  emb_dim=emb_dim,
                  input_data_type_size=input_data_type_size,
                  table_data_type_size=table_data_type_size,
                  fwd_a2a_comm_data_type_size=fwd_a2a_comm_data_type_size,
                  bwd_a2a_comm_data_type_size=bwd_a2a_comm_data_type_size,
                  num_poolings=num_poolings,
                  ddr_mem_bw=ddr_mem_bw,
                  device_bw=device_bw,
                  inter_host_bw=inter_host_bw,
                  intra_host_bw=intra_host_bw,
                  bwd_compute_multiplier=bwd_compute_multiplier,
                  is_pooled=is_pooled,
                  is_weighted=is_weighted,
                  is_inference=is_inference,
                  expected_cache_fetches=expected_cache_fetches,
              )
          elif sharding_type == ShardingType.ROW_WISE.value:
              shard_perf = _get_rw_sharding_perf(
                  batch_sizes=batch_sizes,
                  world_size=world_size,
                  local_world_size=local_world_size,
                  input_lengths=input_lengths,
                  emb_dim=emb_dim,
                  input_data_type_size=input_data_type_size,
                  table_data_type_size=table_data_type_size,
                  fwd_a2a_comm_data_type_size=fwd_a2a_comm_data_type_size,
                  bwd_a2a_comm_data_type_size=bwd_a2a_comm_data_type_size,
                  fwd_sr_comm_data_type_size=fwd_sr_comm_data_type_size,
                  bwd_sr_comm_data_type_size=bwd_sr_comm_data_type_size,
                  num_poolings=num_poolings,
                  ddr_mem_bw=ddr_mem_bw,
                  device_bw=device_bw,
                  inter_host_bw=inter_host_bw,
                  intra_host_bw=intra_host_bw,
                  bwd_compute_multiplier=bwd_compute_multiplier,
                  is_pooled=is_pooled,
                  is_weighted=is_weighted,
                  expected_cache_fetches=expected_cache_fetches,
              )
          elif sharding_type == ShardingType.TABLE_ROW_WISE.value:
              shard_perf = _get_twrw_sharding_perf(
                  batch_sizes=batch_sizes,
                  world_size=world_size,
                  local_world_size=local_world_size,
                  input_lengths=input_lengths,
                  emb_dim=emb_dim,
                  input_data_type_size=input_data_type_size,
                  table_data_type_size=table_data_type_size,
                  fwd_a2a_comm_data_type_size=fwd_a2a_comm_data_type_size,
                  bwd_a2a_comm_data_type_size=bwd_a2a_comm_data_type_size,
                  fwd_sr_comm_data_type_size=fwd_sr_comm_data_type_size,
                  bwd_sr_comm_data_type_size=bwd_sr_comm_data_type_size,
                  num_poolings=num_poolings,
                  ddr_mem_bw=ddr_mem_bw,
                  device_bw=device_bw,
                  inter_host_bw=inter_host_bw,
                  intra_host_bw=intra_host_bw,
                  bwd_compute_multiplier=bwd_compute_multiplier,
                  is_pooled=is_pooled,
                  is_weighted=is_weighted,
                  expected_cache_fetches=expected_cache_fetches,
              )
          elif sharding_type == ShardingType.DATA_PARALLEL.value:
              shard_perf = _get_dp_sharding_perf(
                  batch_sizes=batch_sizes,
                  world_size=world_size,
                  local_world_size=local_world_size,
                  input_lengths=input_lengths,
                  grad_num_elem=hash_size * emb_dim,
                  emb_dim=emb_dim,
                  input_data_type_size=input_data_type_size,
                  table_data_type_size=table_data_type_size,
                  num_poolings=num_poolings,
                  device_bw=device_bw,
                  inter_host_bw=inter_host_bw,
                  bwd_compute_multiplier=bwd_compute_multiplier,
                  is_pooled=is_pooled,
                  is_weighted=is_weighted,
              )
          else:
              raise ValueError(
                  f"Unrecognized or unsupported sharding type provided: {sharding_type}"
              )
          shard_perfs.append(shard_perf)
  
      return shard_perfs
  
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 1 for calling a global function with example usage: perf_func_emb_wall_time(shard_sizes, compute_kernel, compute_device, sharding_type, batch_sizes, world_size, local_world_size, input_lengths, 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, hbm_mem_bw, ddr_mem_bw, intra_host_bw, inter_host_bw, bwd_compute_multiplier, is_pooled, is_weighted, caching_ratio, is_inference, prefetch_pipeline, expected_cache_fetches) and returns: shard_perfs