Parse a torchao.dequantize_affine node. Accepts N-dimensional block_size with a single non-1 element identifying the quantized dimension and group_size. For example: - Linear weights (2D): block_size=[1, 32] → quantized_dim=1 - Conv2d weights (4D): block_size=[1, 32, 1,
(
node: Node,
)
| 310 | |
| 311 | |
| 312 | def parse_dequant_node( |
| 313 | node: Node, |
| 314 | ) -> Optional[Tuple[Node, Node, Node, int, int, Optional[torch.dtype], int]]: |
| 315 | """Parse a torchao.dequantize_affine node. |
| 316 | |
| 317 | Accepts N-dimensional block_size with a single non-1 element identifying |
| 318 | the quantized dimension and group_size. For example: |
| 319 | - Linear weights (2D): block_size=[1, 32] → quantized_dim=1 |
| 320 | - Conv2d weights (4D): block_size=[1, 32, 1, 1] → quantized_dim=1 |
| 321 | |
| 322 | Returns (qdata, scale, zero_point, group_size, bits, out_dtype, quantized_dim) |
| 323 | or None if unsupported. |
| 324 | """ |
| 325 | qdata, block_size, scale, zero_point, dtype, qmin, qmax = node.args[0:7] |
| 326 | out_dtype = ( |
| 327 | node.args[7] if len(node.args) > 7 else node.kwargs.get("output_dtype", None) |
| 328 | ) |
| 329 | if dtype != torch.int8: |
| 330 | return None |
| 331 | if len(block_size) < 2: |
| 332 | return None |
| 333 | non_one = [(i, d) for i, d in enumerate(block_size) if d != 1] |
| 334 | if len(non_one) != 1: |
| 335 | return None |
| 336 | quantized_dim, group_size = non_one[0] |
| 337 | if group_size not in [32, 64, 128]: |
| 338 | return None |
| 339 | |
| 340 | # TODO: MLX supports 3, 5, and 7, but we need to figure out the |
| 341 | # packing story in to_mlx_qparams to use them |
| 342 | bits = (qmax - qmin + 1).bit_length() - 1 |
| 343 | if bits not in [2, 4, 8]: |
| 344 | return None |
| 345 | return qdata, scale, zero_point, group_size, bits, out_dtype, quantized_dim |
| 346 | |
| 347 | |
| 348 | # Mapping from torch dtype to ET ScalarType int value |
no test coverage detected