(self, in_features: int, out_features: int, bias: bool = False)
| 58 | weight_scale: torch.Tensor |
| 59 | |
| 60 | def __init__(self, in_features: int, out_features: int, bias: bool = False): |
| 61 | super().__init__() |
| 62 | self.in_features = in_features |
| 63 | self.out_features = out_features |
| 64 | |
| 65 | self.weight = torch.nn.Parameter(torch.zeros(out_features, in_features//4, dtype=torch.int8), requires_grad=False) |
| 66 | self.weight_scale = torch.nn.Parameter(torch.zeros(4, dtype=torch.bfloat16), requires_grad=False) |
| 67 | |
| 68 | @torch.compile |
| 69 | def quant_input(self, input): |