(
cls, ep: ExportedProgram, head: Node
)
| 851 | |
| 852 | @classmethod |
| 853 | def maybe_create( |
| 854 | cls, ep: ExportedProgram, head: Node |
| 855 | ) -> Optional["QuantizedLinearHandler"]: |
| 856 | linear_node = head |
| 857 | if not match_target(linear_node, torch.ops.aten.linear.default): |
| 858 | return None |
| 859 | |
| 860 | x, w = linear_node.args[0:2] |
| 861 | dequant_node = w |
| 862 | if not match_target(dequant_node, torch.ops.torchao.dequantize_affine.default): |
| 863 | return None |
| 864 | if not has_single_user(dequant_node): |
| 865 | return None |
| 866 | |
| 867 | parsed = parse_dequant_node(dequant_node) |
| 868 | if parsed is None: |
| 869 | return None |
| 870 | qdata, scale, zero_point, group_size, bits, out_dtype, _quantized_dim = parsed |
| 871 | out_dtype = x.meta["val"].dtype if out_dtype is None else out_dtype |
| 872 | |
| 873 | head = linear_node |
| 874 | body = [dequant_node] |
| 875 | return QuantizedLinearHandler( |
| 876 | head, |
| 877 | body, |
| 878 | qdata=qdata, |
| 879 | scale=scale, |
| 880 | zero_point=zero_point, |
| 881 | group_size=group_size, |
| 882 | bits=bits, |
| 883 | out_dtype=out_dtype, |
| 884 | ) |
| 885 | |
| 886 | def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: |
| 887 | assert n == self.head |
nothing calls this directly
no test coverage detected