class_docstr: Manage tensors stored in a compressed format.
class_docstr_tok: ["Manage","tensors","stored","in","a","compressed","format","."]
docstr: Allocate a compressed TorchTensor. Round up the shape to group boundary.
docstr_tok: ["Allocate","a","compressed","TorchTensor",".","Round","up","the","shape","to","group","boundary","."]
code: def allocate(
      self, shape, dtype, comp_config, pin_memory=None, name=None
  ):
      """Allocate a compressed TorchTensor. Round up the shape to group boundary."""
      assert comp_config.num_bits == 4 and dtype == np.float16
  
      group_size, group_dim = (
          comp_config.group_size,
          comp_config.group_dim,
      )
  
      # Round up
      num_groups = (shape[group_dim] + group_size - 1) // group_size
      data_shape = (
          shape[:group_dim]
          + (num_groups * (group_size // 2),)
          + shape[group_dim + 1 :]
      )
      scale_shape = (
          shape[:group_dim] + (num_groups, 2) + shape[group_dim + 1 :]
      )
  
      data = self.base_device.allocate(
          data_shape, np.uint8, pin_memory=pin_memory
      )
      scale = self.base_device.allocate(
          scale_shape, np.float16, pin_memory=pin_memory
      )
  
      return TorchTensor(
          shape,
          np_dtype_to_torch_dtype[dtype],
          (data, scale, comp_config),
          self,
          name=name,
      )
  
imports: import dataclasses
  import torch
  import numpy
  from flexgen.pytorch_backend import TorchTensor, TorchDevice, DeviceType, general_copy, fix_recursive_import
  from flexgen.utils import np_dtype_to_torch_dtype
label_desc: Use image node_id 2 for calling the TorchCompressedDevice obj's underlying member method code with example usage: obj.allocate(shape, dtype, comp_config, pin_memory, name) and returns: TorchTensor