code: def __init__(
      self,
      bits,
      groupsize,
      infeatures,
      outfeatures,
      bias,
      tp_size=1,
      tp_rank=0,
      row_split=False,
  ):
      super().__init__()
      if bits not in [2, 4, 8]:
          raise NotImplementedError("Only 2,4,8 bits are supported.")
      self.infeatures = infeatures
      self.outfeatures = outfeatures
      self.bits = bits
      self.maxq = 2**self.bits - 1
      self.groupsize = groupsize if groupsize != -1 else infeatures
  
      self.register_buffer(
          "qweight",
          torch.zeros(
              (infeatures // 32 * self.bits, outfeatures),
              dtype=torch.int32,
          ),
      )
      self.register_buffer(
          "qzeros",
          torch.zeros(
              (
                  math.ceil(infeatures / self.groupsize),
                  outfeatures // 32 * self.bits,
              ),
              dtype=torch.int32,
          ),
      )
      self.register_buffer(
          "scales",
          torch.zeros(
              (math.ceil(infeatures / self.groupsize), outfeatures),
              dtype=torch.float16,
          ),
      )
      if row_split:
          self.register_buffer(
              "g_idx",
              torch.tensor(
                  [
                      (i + (tp_rank * self.infeatures))
                      // self.groupsize
                      for i in range(infeatures)
                  ],
                  dtype=torch.int32,
              ),
          )
      else:
          self.register_buffer(
              "g_idx",
              torch.tensor(
                  [i // self.groupsize for i in range(infeatures)],
                  dtype=torch.int32,
              ),
          )
  
      if bias:
          self.register_buffer(
              "bias", torch.zeros((outfeatures), dtype=torch.float16)
          )
      else:
          self.bias = None
  
      self.gptq_linear = CaiGPTQLinearOp(groupsize, bits)
  
      self.q4 = None
      self.empty_tensor = torch.empty((1, 1), device="meta")
      self.tp_size = tp_size
      self.tp_rank = tp_rank
      self.row_split = row_split
  
imports: import math
  import warnings
  from typing import List, Union
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed import ProcessGroup
  from colossalai.lazy import LazyInitContext
  from colossalai.shardformer.layer import ParallelModule
  from .gptq_op import CaiGPTQLinearOp
label_desc: Use image node_id 1 to create a new CaiQuantLinear object from inherited base classes: nn with example: obj = CaiQuantLinear(bits, groupsize, infeatures, outfeatures, bias, tp_size, tp_rank, row_split)