class_docstr: Manage tensors stored in a compressed format.
class_docstr_tok: ["Manage","tensors","stored","in","a","compressed","format","."]
docstr: Compress a torch.Tensor. Round up the shape to group boundary.
docstr_tok: ["Compress","a","torch.Tensor",".","Round","up","the","shape","to","group","boundary","."]
code: def compress(self, tensor, comp_config):
      """Compress a torch.Tensor. Round up the shape to group boundary."""
      group_size, num_bits, group_dim, symmetric = (
          comp_config.group_size,
          comp_config.num_bits,
          comp_config.group_dim,
          comp_config.symmetric,
      )
      assert num_bits == 4 and group_size % 2 == 0 and not symmetric
  
      if tensor.device.type == "cpu" and tensor.dtype == torch.float16:
          tensor = tensor.float()
  
      shape = tensor.shape
      num_groups = (shape[group_dim] + group_size - 1) // group_size
  
      # Pad
      new_shape = (
          shape[:group_dim]
          + (num_groups, group_size)
          + shape[group_dim + 1 :]
      )
      pad_len = (
          group_size - shape[group_dim] % group_size
      ) % group_size
      if pad_len != 0:
          pad_shape = (
              shape[:group_dim] + (pad_len,) + shape[group_dim + 1 :]
          )
          tensor = torch.cat(
              [
                  tensor,
                  torch.zeros(
                      pad_shape,
                      dtype=tensor.dtype,
                      device=tensor.device,
                  ),
              ],
              dim=group_dim,
          )
      data = tensor.view(new_shape)
  
      # Quantize
      B = 2**num_bits - 1
      mn = torch.min(data, dim=group_dim + 1, keepdim=True)[0]
      mx = torch.max(data, dim=group_dim + 1, keepdim=True)[0]
  
      scale = B / (mx - mn)
      data = data - mn
      data.mul_(scale)
      data = data.clamp_(0, B).round_().to(torch.uint8)
  
      # Pack
      left_indices = tuple(
          slice(0, x) for x in data.shape[: group_dim + 1]
      ) + (
          slice(0, data.shape[group_dim + 1], 2),
      )
      right_indices = tuple(
          slice(0, x) for x in data.shape[: group_dim + 1]
      ) + (
          slice(1, data.shape[group_dim + 1], 2),
      )
      data = torch.bitwise_or(
          data[left_indices].bitwise_left_shift(4), data[right_indices]
      )
  
      # Reshape
      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 = data.view(data_shape)
      scale = torch.cat([scale, mn], dim=group_dim + 1).view(
          scale_shape
      )
  
      data = TorchTensor.create_from_torch(data, self.base_device)
      scale = TorchTensor.create_from_torch(scale, self.base_device)
  
      return TorchTensor(
          shape, tensor.dtype, (data, scale, comp_config), self
      )
  
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 5 for calling the TorchCompressedDevice obj's underlying member method code with example usage: obj.compress(tensor, comp_config) and returns: TorchTensor