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","."]
docstr: if input indices and offsets haven't been splitted along assigned rank, this function will do it.
  it takes time. please consider splitting data during batch loading.
docstr_tok: ["if","input","indices","and","offsets","have","n't","been","splitted","along","assigned","rank",",","this","function","will","do","it",".","it","takes","time",".","please","consider","splitting","data","during","batch","loading","."]
code: def split_along_rank(
      self,
      batch_size,
      indices: torch.Tensor,
      offsets: torch.Tensor = None,
      per_sample_weights=None,
  ):
      """
      if input indices and offsets haven't been splitted along assigned rank, this function will do it.
      it takes time. please consider splitting data during batch loading.
      """
      local_indices_list: List(torch.Tensor) = []
      local_offsets_list: List(torch.Tensor) = []
      if per_sample_weights != None:
          local_per_sample_weights_list: List(torch.Tensor) = []
  
      offset_pre_end = 0  # local_offsets trick
      for i, handle_table in enumerate(self.assigned_table_list):
          indices_start_position = offsets[batch_size * handle_table]
          if (not self.include_last_offset) and (
              batch_size * (handle_table + 1) >= indices.shape[0]
          ):
              # till-the-end special case
              indices_end_position = indices.shape[0]
          else:
              indices_end_position = offsets[
                  batch_size * (handle_table + 1)
              ]
          # alternative approach: reduce malloc
          """
              # 1. local_indices_list:
              local_indices = indices.narrow(0, indices_start_position, indices_end_position - indices_start_position)
              torch.sub(local_indices, self.idx_offset_list[i], out=local_indices)
              local_indices_list.append(local_indices)
              # 2. local_offsets_list:
              if i + 1 == len(self.assigned_table_list):
                  # till-the-end special case
                  if not self.include_last_offset:
                      local_offsets = offsets.narrow(0, batch_size * handle_table, batch_size)
                  else:
                      local_offsets = offsets.narrow(0, batch_size * handle_table, batch_size + 1)
                  torch.add(local_offsets, offset_pre_end - offsets[batch_size * handle_table], out=local_offsets)
                  local_offsets_list.append(local_offsets)
              else:
                  temp_holder = offsets[batch_size * handle_table].item()
                  local_offsets = offsets.narrow(0, batch_size * handle_table, batch_size)
                  torch.add(local_offsets, offset_pre_end - offsets[batch_size * handle_table], out=local_offsets)
                  offset_pre_end = offsets[batch_size * (handle_table + 1)] + offset_pre_end - temp_holder
                  local_offsets_list.append(local_offsets)
              """
          # 1. local_indices_list:
          local_indices_list.append(
              indices.narrow(
                  0,
                  indices_start_position,
                  indices_end_position - indices_start_position,
              ).sub(self.idx_offset_list[i])
          )
          # 2. local_offsets_list:
          if i + 1 == len(self.assigned_table_list):
              # till-the-end special case
              if not self.include_last_offset:
                  local_offsets = offsets.narrow(
                      0, batch_size * handle_table, batch_size
                  ).add(
                      offset_pre_end
                      - offsets[batch_size * (handle_table)]
                  )
              else:
                  local_offsets = offsets.narrow(
                      0, batch_size * handle_table, batch_size + 1
                  ).add(
                      offset_pre_end
                      - offsets[batch_size * (handle_table)]
                  )
              local_offsets_list.append(local_offsets)
          else:
              local_offsets = offsets.narrow(
                  0, batch_size * handle_table, batch_size + 1
              ).add(
                  offset_pre_end - offsets[batch_size * (handle_table)]
              )
              offset_pre_end = local_offsets[-1]
              local_offsets_list.append(local_offsets[:-1])
          # 3. local_per_sample_weights_list:
          if per_sample_weights != None:
              local_per_sample_weights_list.append(
                  per_sample_weights[
                      indices_start_position:indices_end_position
                  ]
              )
      local_indices = torch.cat(local_indices_list, 0)
      local_offsets = torch.cat(local_offsets_list, 0)
      local_per_sample_weights = None
      if per_sample_weights != None:
          local_per_sample_weights = torch.cat(
              local_per_sample_weights_list, 0
          )
      return local_indices, local_offsets, local_per_sample_weights
  
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 3 for calling the ParallelCachedEmbeddingBagTablewise obj's underlying member method code with example usage: obj.split_along_rank(batch_size, indices, offsets, per_sample_weights) and returns: local_indices,  local_offsets,  local_per_sample_weights