class_docstr: every table assigned to this class instance is managed by a CachedEmbeddingBag.
class_docstr_tok: ["every","table","assigned","to","this","class","instance","is","managed","by","a","CachedEmbeddingBag","."]
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,
      mode="mean",
      include_last_offset=False,
      dtype=None,
      device=None,
      warmup_ratio=0.7,
      pin_weight=False,
      evict_strategy: EvictionStrategy = EvictionStrategy.LFU,
  ):
      super(
          ParallelCachedEmbeddingBagTablewiseSpiltCache, self
      ).__init__()
      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] = []
      for i, rank in enumerate(self.rank_of_tables):
          if rank == self.rank:
              self.assigned_table_list.append(i)
      self.include_last_offset = include_last_offset
      self.pg = ProcessGroup(tp_degree=self.world_size)
  
      # prepare CachedEmbeddingBag list
  
      self.cached_embedding_bag_list: nn.ModuleList = nn.ModuleList()
      for config in embedding_bag_config_list:
          if config.assigned_rank != self.rank:
              continue
          self.cached_embedding_bag_list.append(
              CachedEmbeddingBag(
                  num_embeddings=config.num_embeddings,
                  embedding_dim=embedding_dim,
                  padding_idx=padding_idx,
                  max_norm=max_norm,
                  norm_type=norm_type,
                  scale_grad_by_freq=scale_grad_by_freq,
                  sparse=sparse,
                  _weight=config.initial_weight,
                  mode=mode,
                  include_last_offset=include_last_offset,
                  dtype=dtype,
                  device=device,
                  cuda_row_num=config.cuda_row_num,
                  ids_freq_mapping=config.ids_freq_mapping,
                  warmup_ratio=warmup_ratio,
                  buffer_size=config.buffer_size,
                  pin_weight=pin_weight,
                  evict_strategy=evict_strategy,
              )
          )
  
      # 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
  
file_path: ColossalAI/colossalai/legacy/nn/parallel/layers/cache_embedding/parallel_cached_embedding_tablewise_split_cache.py
imports: import abc
  from typing import List
  import torch
  import torch.distributed
  import torch.nn
  from torch.profiler import record_function
  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 ParallelCachedEmbeddingBagTablewiseSpiltCache object from inherited base classes: abc, nn with example: obj = ParallelCachedEmbeddingBagTablewiseSpiltCache(embedding_bag_config_list, embedding_dim, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse, mode, include_last_offset, dtype, device, warmup_ratio, pin_weight, evict_strategy)