class_docstr: A manager class to manipulate the tensors in chunks.
  
  Args:
      chunk_configuration (Dict[int, Dict]): the configuration dictionary of this chunk manager.
      init_device (torch.device): optional, the device on which the chunk is initialized. The default is None.
class_docstr_tok: ["A","manager","class","to","manipulate","the","tensors","in","chunks",".","Args",":","chunk_configuration","(","Dict","[","int",",","Dict","]",")",":","the","configuration","dictionary","of","this","chunk","manager",".","init_device","(","torch.device",")",":","optional",",","the","device","on","which","the","chunk","is","initialized",".","The","default","is","None","."]
docstr: Register a tensor to the chunk manager.
  Then, the tensor should be accessed by `get_chunks`.
  
  Args:
      tensor: the tensor appended to the chunk
      group_type: the data type of the group.
      config_key: the key of the group's name, the size of the dp world
      cpu_offload: if True, the chunk will be closed on CPU
      pin_memory: whether the chunk is pinned in the cpu memory
docstr_tok: ["Register","a","tensor","to","the","chunk","manager",".","Then",",","the","tensor","should","be","accessed","by","`","get_chunks","`",".","Args",":","tensor",":","the","tensor","appended","to","the","chunk","group_type",":","the","data","type","of","the","group",".","config_key",":","the","key","of","the","group","'s","name",",","the","size","of","the","dp","world","cpu_offload",":","if","True",",","the","chunk","will","be","closed","on","CPU","pin_memory",":","whether","the","chunk","is","pinned","in","the","cpu","memory"]
code: def register_tensor(
      self,
      tensor: torch.Tensor,
      group_type: str,
      config_key: int,
      zero_group: ProcessGroup,
      extra_dp_group: ProcessGroup = None,
      cpu_offload: bool = False,
      pin_memory: bool = False,
  ) -> None:
      """
      Register a tensor to the chunk manager.
      Then, the tensor should be accessed by `get_chunks`.
  
      Args:
          tensor: the tensor appended to the chunk
          group_type: the data type of the group.
          config_key: the key of the group's name, the size of the dp world
          cpu_offload: if True, the chunk will be closed on CPU
          pin_memory: whether the chunk is pinned in the cpu memory
      """
      assert tensor not in self.tensor_chunk_map
      assert isinstance(
          tensor, torch.Tensor
      ), "Please feed Tensor to this ChunkManager"
      assert config_key in self.dp_degree_chunk_size_dict
  
      chunk_size = self.dp_degree_chunk_size_dict[config_key]
      chunk_kwargs = self.kwargs_config[config_key]
      group_name = "{}_{}".format(group_type, config_key)
      chunk_group = self.__get_chunk_group(group_name)
  
      try:
          # append the tensor to the last chunk
          chunk_group[-1].append_tensor(tensor)
      except (IndexError, ChunkFullError):
          # the except statement will be triggered when there is no chunk or
          # the last chunk in the chunk group is full
          # this will create a new chunk and allocate this chunk to its corresponding process
          if chunk_group:
              # the chunk group is not empty
              # close the last chunk
              self.__close_one_chunk(chunk_group[-1])
  
          if tensor.numel() > chunk_size:
              chunk_size = tensor.numel()
              dp_size = dist.get_world_size(zero_group)
              chunk_size = chunk_size + (-chunk_size % dp_size)
  
          chunk = Chunk(
              chunk_size=chunk_size,
              zero_group=zero_group,
              dtype=tensor.dtype,
              cpu_shard_init=cpu_offload,
              pin_memory=pin_memory,
              extra_dp_group=extra_dp_group,
              **chunk_kwargs,
          )
  
          chunk_group.append(chunk)
          chunk.append_tensor(tensor)
          self.__add_memory_usage(chunk.memory_usage)
  
      self.tensor_chunk_map[tensor] = chunk_group[-1]
  
imports: from collections import deque
  from typing import Deque, Dict, Iterable, List, Optional, Set, Tuple
  import torch
  import torch.distributed
  from torch.distributed import ProcessGroup
  from colossalai.utils import free_storage, get_current_device
  from .chunk import Chunk, ChunkFullError, TensorState
label_desc: Use image node_id 2 for calling the ChunkManager obj's underlying member method code with example usage: obj.register_tensor(tensor, group_type, config_key, zero_group, extra_dp_group, cpu_offload, pin_memory) without return types