Function _convolution_handler

backends/mlx/ops.py:2573–2617

Handle aten.convolution.default — the unified convolution op. Args layout: convolution(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups) This op appears when PyTorch doesn't decompose to specific conv ops (e.g. grouped conv

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten.convolution.default])
def _convolution_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """Handle aten.convolution.default — the unified convolution op.

    Args layout: convolution(input, weight, bias, stride, padding, dilation,
                             transposed, output_padding, groups)

    This op appears when PyTorch doesn't decompose to specific conv ops
    (e.g. grouped conv_transpose).
    """
    raw_args = n.args
    x_node, w_node = raw_args[0], raw_args[1]
    bias_node = raw_args[2] if len(raw_args) > 2 else None
    transposed = raw_args[6] if len(raw_args) > 6 else False
    groups = raw_args[8] if len(raw_args) > 8 else 1

    if not transposed:
        raise ValueError(
            "aten.convolution with transposed=False: use aten.conv{1,2,3}d instead"
        )

    x_meta = x_node.meta.get("val")
    if x_meta is None:
        raise ValueError("aten.convolution: input shape metadata required")
    ndim = len(x_meta.shape) - 2

    stride = _normalize_conv_param(raw_args[3] if len(raw_args) > 3 else 1, ndim, 1)
    padding = _normalize_conv_param(raw_args[4] if len(raw_args) > 4 else 0, ndim, 0)
    dilation = _normalize_conv_param(raw_args[5] if len(raw_args) > 5 else 1, ndim, 1)
    output_padding = _normalize_conv_param(
        raw_args[7] if len(raw_args) > 7 else 0, ndim, 0
    )

    return _emit_conv_transpose(
        P,
        n,
        x_node,
        w_node,
        bias_node,
        stride,
        padding,
        dilation,
        output_padding,
        groups,
        ndim,
    )


@REGISTRY.register(target=[torch.ops.aten.conv_transpose1d.default])