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 forward(
      self,
      indices: torch.Tensor,
      offsets: torch.Tensor = None,
      per_sample_weights=None,
      shape_hook=None,
      already_split_along_rank=True,
  ):
      if not already_split_along_rank:
          # not recommanded. it takes time.
          batch_size = (offsets.shape[0]) // self.global_tables_num
          (
              local_indices,
              local_offsets,
              local_per_sample_weights,
          ) = self.split_along_rank(
              batch_size, indices, offsets, per_sample_weights
          )
      else:
          # recommanded.
          batch_size = (offsets.shape[0]) // len(
              self.assigned_table_list
          )
          local_indices, local_offsets, local_per_sample_weights = (
              indices,
              offsets,
              per_sample_weights,
          )
      if self.cache_op:
          with torch.no_grad():
              indices = self.cache_weight_mgr.prepare_ids(local_indices)
      local_output = F.embedding_bag(
          indices.cuda(),
          self.cache_weight_mgr.cuda_cached_weight,
          local_offsets,
          self.max_norm,
          self.norm_type,
          self.scale_grad_by_freq,
          self.mode,
          self.sparse,
          local_per_sample_weights,
          self.include_last_offset,
          self.padding_idx,
      )
      local_output = torch.cat(local_output.split(batch_size), 1)
      remains = batch_size % self.world_size
      scatter_strides = [
          batch_size // self.world_size + int(i < remains)
          for i in range(self.world_size)
      ]
      output_full = dual_all_to_all_tablewise(
          local_output,
          self.pg,
          scatter_strides,
          self.embedding_dim_per_rank,
      )
      if shape_hook is not None:
          output_full = shape_hook(output_full)
      return output_full
  
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 2 for calling the ParallelCachedEmbeddingBagTablewise obj's underlying member method code with example usage: obj.forward(indices, offsets, per_sample_weights, shape_hook, already_split_along_rank) and returns: output_full