(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: torch.Tensor, weight_bit_width)
| 9 | class W8A16Linear(torch.autograd.Function): |
| 10 | @staticmethod |
| 11 | def forward(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: torch.Tensor, weight_bit_width): |
| 12 | ctx.inp_shape = inp.size() |
| 13 | ctx.weight_shape = quant_w.size() |
| 14 | ctx.weight_bit_width = weight_bit_width |
| 15 | out_features = quant_w.size(0) |
| 16 | inp = inp.contiguous().view(-1, inp.size(-1)) |
| 17 | weight = extract_weight_to_half(quant_w, scale_w, weight_bit_width) |
| 18 | output = inp.mm(weight.t()) |
| 19 | ctx.save_for_backward(inp, quant_w, scale_w) |
| 20 | return output.view(*(ctx.inp_shape[:-1] + (out_features,))) |
| 21 | |
| 22 | @staticmethod |
| 23 | def backward(ctx, grad_output: torch.Tensor): |
nothing calls this directly
no test coverage detected