code: def should_do_dim_bucketing(
      embedding_tables: List[ShardedEmbeddingTable],
  ) -> bool:
      """
      When embedding memory offloading with caching is enabled, we prefer to
      do dim bucketing for better utilization of cache space. Only applied to
      "prefetch-sparse-dist" training pipeline.
  
      Currently using the compute kernel to deduct caching is enabled.
      """
      table_pipeline_count = 0
      for table in embedding_tables:
          if (
              table.fused_params is not None
              and "prefetch_pipeline" in table.fused_params
              and table.fused_params["prefetch_pipeline"]
          ):
              table_pipeline_count += 1
  
      if table_pipeline_count > 0 and table_pipeline_count != len(
          embedding_tables
      ):
          AssertionError(
              f"Only {table_pipeline_count} tables have prefetch-sparse-dist pipeline. It should be all {len(embedding_tables)} tables."
          )
  
      for table in embedding_tables:
          if (
              table.compute_kernel
              == EmbeddingComputeKernel.FUSED_UVM_CACHING
              and table_pipeline_count
          ):
              return True
      return False
  
file_path: torchrec/torchrec/distributed/embedding_dim_bucketer.py
imports: from enum import Enum, unique
  from typing import Dict, List
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel, ShardedEmbeddingTable
  from torchrec.modules.embedding_configs import DATA_TYPE_NUM_BITS, DataType
label_desc: Use image node_id 1 for calling a global function with example usage: should_do_dim_bucketing(embedding_tables) and returns: False, True