class_docstr: all tables assigned to this class instance are managed by a single CachedEmbeddingBag.
  Those parameters in TablewiseEmbeddingBagConfig are ignored: cuda_row_num, buffer_size, initial_weight.
class_docstr_tok: ["all","tables","assigned","to","this","class","instance","are","managed","by","a","single","CachedEmbeddingBag",".","Those","parameters","in","TablewiseEmbeddingBagConfig","are","ignored",":","cuda_row_num",",","buffer_size",",","initial_weight","."]
code: def __init__(
      self,
      embedding_bag_config_list: List[TablewiseEmbeddingBagConfig],
      embedding_dim: int,
      padding_idx=None,
      max_norm=None,
      norm_type=2.0,
      scale_grad_by_freq=False,
      sparse=False,
      _weight=None,
      mode="mean",
      include_last_offset=False,
      dtype=None,
      device=None,
      cache_ratio=0.01,
      warmup_ratio=0.7,
      buffer_size=50_000,
      pin_weight=False,
      evict_strategy: EvictionStrategy = EvictionStrategy.LFU,
  ):
      self.rank = dist.get_rank()
      self.world_size = dist.get_world_size()
      self.rank_of_tables = [
          config.assigned_rank for config in embedding_bag_config_list
      ]
      self.global_table_num_embeddings_list = [
          config.num_embeddings for config in embedding_bag_config_list
      ]
      self.global_tables_num = len(embedding_bag_config_list)
      self.global_tables_offsets = torch.cumsum(
          torch.tensor([0] + self.global_table_num_embeddings_list), 0
      ).cuda()
      self.assigned_table_list: List[int] = []
      self.pg = ProcessGroup(tp_degree=self.world_size)
      self.num_embeddings = 0
      for i, rank in enumerate(self.rank_of_tables):
          if rank == self.rank:
              self.assigned_table_list.append(i)
              self.num_embeddings += (
                  self.global_table_num_embeddings_list[i]
              )
      self.include_last_offset = include_last_offset
  
      ids_freq_mapping = []
      for config in embedding_bag_config_list:
          if config.assigned_rank == self.rank:
              if config.ids_freq_mapping != None:
                  ids_freq_mapping.extend(config.ids_freq_mapping)
              else:
                  ids_freq_mapping = None
                  break
      self.cache_ratio = cache_ratio
      # table-associate cache
      int(cache_ratio * self.num_embeddings)
      super(ParallelCachedEmbeddingBagTablewise, self).__init__(
          self.num_embeddings,
          embedding_dim,
          padding_idx,
          max_norm,
          norm_type,
          scale_grad_by_freq,
          sparse,
          _weight,
          mode,
          include_last_offset,
          dtype,
          device,
          cache_ratio,
          ids_freq_mapping,
          warmup_ratio,
          buffer_size,
          pin_weight,
          evict_strategy,
      )
  
      # for assigned tables reconnection:
      self.idx_offset_list = []
      offset_cumsum = 0
      for table_i, table_num_embeddings in enumerate(
          self.global_table_num_embeddings_list
      ):
          if self.rank_of_tables[table_i] == self.rank:
              self.idx_offset_list.append(offset_cumsum)
          else:
              offset_cumsum += table_num_embeddings
  
      # prepare list shape for all_to_all output
      self.embedding_dim_per_rank = [0 for i in range(self.world_size)]
      for rank in self.rank_of_tables:
          self.embedding_dim_per_rank[rank] += embedding_dim
  
      self.cache_op = True
  
file_path: colossalai/colossalai/legacy/nn/parallel/layers/cache_embedding/parallel_cached_embedding_tablewise.py
imports: from typing import List
  import torch
  import torch.distributed
  import torch.nn.functional
  from colossalai.legacy.nn._ops._utils import dual_all_to_all_tablewise
  from colossalai.legacy.tensor import ProcessGroup
  from .cache_mgr import EvictionStrategy
  from .cached_embedding import CachedEmbeddingBag
  from .embedding_config import TablewiseEmbeddingBagConfig
label_desc: Use image node_id 1 to create a new ParallelCachedEmbeddingBagTablewise object from inherited base classes: CachedEmbeddingBag with example: obj = ParallelCachedEmbeddingBagTablewise(embedding_bag_config_list, embedding_dim, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse, _weight, mode, include_last_offset, dtype, device, cache_ratio, warmup_ratio, buffer_size, pin_weight, evict_strategy)