code: def __init__(
      self,
      num_embeddings,
      embedding_dim,
      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,
      ids_freq_mapping=None,
      warmup_ratio=0.7,
      buffer_size=50_000,
      pin_weight=False,
      evict_strategy: EvictionStrategy = EvictionStrategy.DATASET,
  ):
      self.rank = torch.distributed.get_rank()
      self.world_size = torch.distributed.get_world_size()
  
      (
          self.partition_start_index,
          self.partition_end_index,
          divisible,
      ) = get_partition(embedding_dim, self.rank, self.world_size)
      self.embedding_dim_per_partition = (
          self.partition_end_index - self.partition_start_index
      )
  
      super(ParallelCachedEmbeddingBag, self).__init__(
          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,
      )
      self.cache_op = True
  
file_path: ColossalAI/colossalai/legacy/nn/parallel/layers/cache_embedding/parallel_cached_embedding.py
imports: from typing import List, Optional, Tuple
  import torch
  import torch.nn.functional
  from colossalai.legacy.nn._ops._utils import dual_all_to_all
  from colossalai.legacy.tensor import ColoTensorSpec, ComputePattern, ProcessGroup, ShardSpec
  from colossalai.tensor import ColoTensor
  from .cache_mgr import EvictionStrategy
  from .cached_embedding import CachedEmbeddingBag
label_desc: Use image node_id 1 to create a new ParallelCachedEmbeddingBag object from inherited base classes: CachedEmbeddingBag with example: obj = ParallelCachedEmbeddingBag(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)