Function _dequantize_affine_handler

backends/mlx/ops.py:3770–3867

Handle standalone torchao.dequantize_affine (not fused with linear/embedding). MLX's dequantize always operates along the last axis. When the quantized dimension is not last (e.g. Conv2d with block_size=[1,32,1,1]), we permute the constant weight/scale/zero_point tensors at compile tim

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.torchao.dequantize_affine.default])
def _dequantize_affine_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """Handle standalone torchao.dequantize_affine (not fused with linear/embedding).

    MLX's dequantize always operates along the last axis.  When the quantized
    dimension is not last (e.g. Conv2d with block_size=[1,32,1,1]), we permute
    the constant weight/scale/zero_point tensors at compile time so the
    quantized dim becomes last, emit the DequantizeNode, then emit a
    TransposeNode with the inverse permutation to restore the original layout.
    """
    parsed = parse_dequant_node(n)
    if parsed is None:
        raise NotImplementedError(
            f"dequantize_affine: unsupported quantization config at {n}"
        )
    (
        qdata_node,
        scale_node,
        zero_point_node,
        group_size,
        bits,
        out_dtype,
        quantized_dim,
    ) = parsed

    qdata_target, qdata = P.get_placeholder_target_and_tensor(qdata_node)
    zero_point_target, zero_point = P.get_placeholder_target_and_tensor(zero_point_node)
    scale_target, scale = P.get_placeholder_target_and_tensor(scale_node)

    if out_dtype is None:
        out_dtype = scale_node.meta["val"].dtype
    out_scalar_type = torch_dtype_to_scalar_type(out_dtype)

    ndim = qdata.ndim
    needs_permute = quantized_dim != ndim - 1

    if needs_permute:
        perm = list(range(ndim))
        perm.remove(quantized_dim)
        perm.append(quantized_dim)
        qdata = qdata.permute(perm).contiguous()
        scale = scale.permute(perm).contiguous()
        zero_point = zero_point.permute(perm).contiguous()

    # to_mlx_qparams expects 2D tensors; flatten N-D to 2D for packing,
    # then restore the (possibly permuted) leading dimensions afterward.
    permuted_shape = qdata.shape
    qdata_2d = qdata.reshape(-1, qdata.shape[-1])
    scale_2d = scale.reshape(-1, scale.shape[-1])
    zero_point_2d = zero_point.reshape(-1, zero_point.shape[-1])

    Q, B = to_mlx_qparams(qdata_2d, scale_2d, zero_point_2d, bits)

    leading_dims = permuted_shape[:-1]
    Q = Q.reshape(*leading_dims, Q.shape[-1])
    scale_nd = scale_2d.reshape(*leading_dims, scale_2d.shape[-1])
    if B is not None:
        B = B.reshape(*leading_dims, B.shape[-1])