Function _constant_pad_nd_handler

backends/mlx/ops.py:2815–2877

Handle aten.constant_pad_nd - pad with a constant value. PyTorch pad format: [left_0, right_0, left_1, right_1, ...] MLX pad_width format: [(before_0, after_0), (before_1, after_1), ...] Note: PyTorch pads in reverse order (last dimensions first).

(P: MLXProgramBuilder, n: Node)

@REGISTRY.register(target=[torch.ops.aten.constant_pad_nd.default])
def _constant_pad_nd_handler(P: MLXProgramBuilder, n: Node) -> Slot:
    """Handle aten.constant_pad_nd - pad with a constant value.

    PyTorch pad format: [left_0, right_0, left_1, right_1, ...]
    MLX pad_width format: [(before_0, after_0), (before_1, after_1), ...]

    Note: PyTorch pads in reverse order (last dimensions first).
    """
    args = P.args(n)
    require_args(args, 2, 3, "aten.constant_pad_nd")
    require_kwargs(P.kwargs(n), set(), "aten.constant_pad_nd")
    x_node, pad = args[0], args[1]
    value = args[2] if len(args) > 2 else 0

    if not isinstance(value, (int, float)):
        raise ValueError(
            f"aten.constant_pad_nd: constant value must be a scalar, got {type(value)}"
        )

    # Convert PyTorch pad format to MLX pad_width format
    # PyTorch: [left_D, right_D, left_D-1, right_D-1, ...]
    # MLX: [(before_0, after_0), (before_1, after_1), ..., (before_D, after_D)]
    if len(pad) % 2 != 0:
        raise ValueError(
            f"aten.constant_pad_nd: pad length must be even, got {len(pad)}"
        )

    x = P.slot_map([x_node])[0]
    x_meta = n.args[0].meta.get("val")
    if x_meta is None:
        raise ValueError("Input tensor metadata not found for constant_pad_nd")

    ndim = len(x_meta.shape)
    num_pad_dims = len(pad) // 2

    if num_pad_dims > ndim:
        raise ValueError(
            f"aten.constant_pad_nd: trying to pad {num_pad_dims} dimensions "
            f"but input has only {ndim} dimensions"
        )

    # Build MLX pad_width: start with zeros for non-padded dims
    pad_width = []
    for _ in range(ndim - num_pad_dims):
        pad_width.extend([0, 0])  # No padding for these dimensions

    # Add padding for the padded dimensions (reverse order)
    for i in range(num_pad_dims - 1, -1, -1):
        left = pad[i * 2]
        right = pad[i * 2 + 1]
        pad_width.extend([left, right])

    out = P.make_or_get_slot(n)
    P.emit(
        PadNode(
            x=P.slot_to_tid(x),
            out=P.slot_to_tid(out),
            pad_width=[P.to_int_or_vid(v) for v in pad_width],