Function _emit_conv

backends/mlx/ops.py:2316–2410

Shared logic for regular convolution emission. Handles weight transform, input/output transposition, bias, and node emission for all spatial dimensions (1D, 2D, 3D). Weight: [C_out, C_in/G, *K] -> [C_out, *K, C_in/G] Input: (N, C, *spatial) -> (N, *spatial, C) Output: (N, *spat

(
    P: MLXProgramBuilder,
    n: Node,
    x_node: Node,
    w_node: Node,
    bias_node,
    stride: list,
    padding: list,
    dilation: list,
    groups: int,
    ndim: int,
)

def _emit_conv(
    P: MLXProgramBuilder,
    n: Node,
    x_node: Node,
    w_node: Node,
    bias_node,
    stride: list,
    padding: list,
    dilation: list,
    groups: int,
    ndim: int,
) -> Slot:
    """Shared logic for regular convolution emission.

    Handles weight transform, input/output transposition, bias, and node emission
    for all spatial dimensions (1D, 2D, 3D).

    Weight: [C_out, C_in/G, *K] -> [C_out, *K, C_in/G]
    Input: (N, C, *spatial) -> (N, *spatial, C)
    Output: (N, *spatial, C) -> (N, C, *spatial)
    """
    if ndim == 3 and groups != 1:
        raise ValueError(
            "conv3d with groups != 1 is not supported by MLX. " f"Got groups={groups}."
        )

    # Permutation: channels-first [N, C, *spatial] <-> channels-last [N, *spatial, C]
    ch_first_to_last = [0] + list(range(2, ndim + 2)) + [1]
    ch_last_to_first = [0, ndim + 1] + list(range(1, ndim + 1))

    # Weight: [C_out, C_in/G, *K] -> [C_out, *K, C_in/G] (same permutation)
    w = _emit_channel_last_weight(P, w_node, ch_first_to_last)

    x, bias = P.slot_map([x_node, bias_node])

    _, tmp = P.make_tmp_slot()
    P.emit(
        TransposeNode(x=P.slot_to_tid(x), out=P.slot_to_tid(tmp), perm=ch_first_to_last)
    )

    if ndim == 1:
        P.emit(
            Conv1DNode(
                x=P.slot_to_tid(tmp),
                w=P.slot_to_tid(w),
                out=P.slot_to_tid(tmp),
                stride=stride[0],
                padding=padding[0],
                dilation=dilation[0],
                groups=groups,
            )
        )
    elif ndim == 2:
        P.emit(
            Conv2DNode(
                x=P.slot_to_tid(tmp),
                w=P.slot_to_tid(w),
                out=P.slot_to_tid(tmp),