class_docstr: Buckets embedding dimensions into different groups based on their sizes. This is intended to be leveraged
  once planning is done, and at the sharding stage, per rank.
  
  The rationale to use bucketization is
  
   - When UVM_CACHING is used: FBGEMM Table Batched Embedding Operator supports a software managed cache for the embeddings placed on UVM (Host memory).
     However, the cache uses maximum embedding dim of all the tables batched in the operator as its unit of allocation. This results in waisted HBM memory,
     and higher miss rate, hence lower performance. Bucketizing can address this issue, allowing for higher effective cache size and better performace.
  
  - When all tables are placed on HBM: When tables with widely different embedding dimension are batched together, the register allocation in GPU will
    be mainly decided by the table with largest embedding dimension. This can lead to worse performance due to lower number of threads and lower occupancy.
  
  Note that Column wise sharding also to some extent addresses this problem, but has its own limitations.
  
  
  Generally, we expect the CACHELINE_BUCKETS policy perform better than ALL_BUCKETS, as it addresses the main issues and limits the number of buckets.
  
  
  Args:
      embedding_tables (List[ShardedEmbeddingTable]): list of sharded embedding
      cfg (EmbDimBucketerPolicy): Bucketing policy
  
  returns:
      emb_dim_buckets (Dict[int, int]): Mapping from embedding dim to bucket id
  
  
  Example:
      emb_dim_bucketer = EmbDimBucketer(embedding_tables, EmbDimBucketerPolicy.SINGLE_BUCKET)
      ...
      bucket = emb_dim_bucketer.get_bucket(embedding_tables[0], embedding_tables[0].data_type) # bucket table 0 is assigned to.
code: def all_emb_buckets(
      self,
      embedding_tables: List[ShardedEmbeddingTable],
  ) -> Dict[int, int]:
      buckets: Dict[int, int] = {}
      bucket_id = -1
  
      for table in embedding_tables:
          dim_in_bytes = self.dim_in_bytes(
              table.embedding_dim, table.data_type
          )
          if dim_in_bytes not in buckets.keys():
              bucket_id += 1
              buckets[dim_in_bytes] = bucket_id
  
      self.num_buckets = bucket_id + 1  # id starts from 0
  
      return buckets
  
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 5 for calling the EmbDimBucketer obj's underlying member method code with example usage: obj.all_emb_buckets(embedding_tables) and returns: buckets