Function _emit_conv_transpose_weight

backends/mlx/ops.py:2241–2284

Get conv_transpose weight in MLX format, handling grouped convolutions. PyTorch conv_transpose weight shape: [C_in, C_out/G, *K] MLX expects: [C_out, *K, C_in/G] For groups=1, a simple permute suffices (C_in==C_in/G, C_out/G==C_out). For groups>1, we need re

(
    P: MLXProgramBuilder, w_node: Node, groups: int, ndim: int
)

def _emit_conv_transpose_weight(
    P: MLXProgramBuilder, w_node: Node, groups: int, ndim: int
) -> Slot:
    """Get conv_transpose weight in MLX format, handling grouped convolutions.

    PyTorch conv_transpose weight shape: [C_in, C_out/G, *K]
    MLX expects:                         [C_out, *K, C_in/G]

    For groups=1, a simple permute suffices (C_in==C_in/G, C_out/G==C_out).
    For groups>1, we need reshape-permute-reshape to rearrange the group dim:
      [C_in, C_out/G, *K] -> [G, C_in/G, C_out/G, *K]
                           -> [G, C_out/G, *K, C_in/G]
                           -> [C_out, *K, C_in/G]
    """
    if groups == 1:
        # Simple permute: [C_in, C_out, *K] -> [C_out, *K, C_in]
        # e.g. 1D: [1, 2, 0], 2D: [1, 2, 3, 0], 3D: [1, 2, 3, 4, 0]
        perm = list(range(1, ndim + 2)) + [0]
        return _emit_channel_last_weight(P, w_node, perm)

    # Grouped: need reshape-permute-reshape at compile time
    if w_node.op != "placeholder":
        raise ValueError(
            f"conv_transpose with groups > 1 requires static weights, "
            f"got dynamic weight from {w_node.op}"
        )

    w_target, w_tensor = P.get_placeholder_target_and_tensor(w_node)
    c_in = w_tensor.shape[0]
    c_out_per_g = w_tensor.shape[1]
    kernel_shape = list(w_tensor.shape[2:])
    c_in_per_g = c_in // groups

    # [C_in, C_out/G, *K] -> [G, C_in/G, C_out/G, *K]
    w = w_tensor.reshape([groups, c_in_per_g, c_out_per_g] + kernel_shape)
    # [G, C_in/G, C_out/G, *K] -> [G, C_out/G, *K, C_in/G]
    # perm: [0, 2, 3, ..., ndim+1, 1]
    perm = [0, 2] + list(range(3, ndim + 3)) + [1]
    w = w.permute(perm).contiguous()
    # [G, C_out/G, *K, C_in/G] -> [C_out, *K, C_in/G]
    c_out = groups * c_out_per_g
    w = w.reshape([c_out] + kernel_shape + [c_in_per_g])

    return P.make_or_get_constant(f"{w_target}_channel_last", w)


def _emit_conv_bias(