(
cls, ep: ExportedProgram, head: Node
)
| 970 | |
| 971 | @classmethod |
| 972 | def maybe_create( |
| 973 | cls, ep: ExportedProgram, head: Node |
| 974 | ) -> Optional["QuantizedEmbeddingHandler"]: |
| 975 | embedding_node = head |
| 976 | if not match_target(embedding_node, torch.ops.aten.embedding.default): |
| 977 | return None |
| 978 | |
| 979 | w, x = embedding_node.args[0:2] |
| 980 | |
| 981 | dequant_node = w |
| 982 | if not match_target(dequant_node, torch.ops.torchao.dequantize_affine.default): |
| 983 | return None |
| 984 | if not has_single_user(dequant_node): |
| 985 | return None |
| 986 | |
| 987 | parsed = parse_dequant_node(dequant_node) |
| 988 | if parsed is None: |
| 989 | return None |
| 990 | qdata, scale, zero_point, group_size, bits, out_dtype, _quantized_dim = parsed |
| 991 | out_dtype = scale.meta["val"].dtype if out_dtype is None else out_dtype |
| 992 | |
| 993 | head = embedding_node |
| 994 | body = [dequant_node] |
| 995 | return QuantizedEmbeddingHandler( |
| 996 | head, |
| 997 | body, |
| 998 | qdata=qdata, |
| 999 | scale=scale, |
| 1000 | zero_point=zero_point, |
| 1001 | group_size=group_size, |
| 1002 | bits=bits, |
| 1003 | out_dtype=out_dtype, |
| 1004 | ) |
| 1005 | |
| 1006 | def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: |
| 1007 | assert n == self.head |
nothing calls this directly
no test coverage detected