[Relax][Frontend][ONNX] Add support for Pad mode="wrap" for opset 19

[napronald]

Summary

The ONNX Pad operator introduced mode="wrap" (circular padding) in opset 19. Currently, the Relax ONNX frontend has no support for opset 19, which raises

OpAttributeInvalid(tvm.error.OpAttributeInvalid: Value wrap in attribute "mode" is invalid for operator Pad.

Changes

Add opset 19 handling to the Pad converter that dispatches mode="wrap" to topi.nn.circular_pad, which already implements circular padding but was never wired up to the ONNX frontend. Existing behavior for earlier Pad opsets is unchanged.

Reproduce

import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper

import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

def make_model():
    x = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 4])
    y = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 3, 8])

    pads = numpy_helper.from_array(
        np.array([0, 0, 2, 0, 0, 2], dtype=np.int64),
        name="pads",
    )

    node = helper.make_node(
        "Pad",
        inputs=["input", "pads"],
        outputs=["output"],
        mode="wrap",
    )

    graph = helper.make_graph([node], "pad_wrap_graph", [x], [y], initializer=[pads])
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 19)])
    onnx.checker.check_model(model)
    return model

def run_tvm(model, x_np):
    mod = from_onnx(model, shape_dict={"input": list(x_np.shape)})

    target = tvm.target.Target("llvm")
    dev = tvm.cpu(0)

    with tvm.transform.PassContext(opt_level=3):
        ex = relax.build(mod, target)

    vm = relax.VirtualMachine(ex, dev)
    out = vm["main"](tvm.runtime.tensor(x_np, dev))
    return out.numpy() if hasattr(out, "numpy") else out.asnumpy()

x_np = np.array(
    [[[1, 2, 3, 4],
      [5, 6, 7, 8],
      [9, 10, 11, 12]]],
    dtype=np.float32,
)

expected = np.pad(x_np, [[0, 0], [0, 0], [2, 2]], mode="wrap")
actual = run_tvm(make_model(), x_np)

print("Expected:")
print(expected[0])
print("Actual:")
print(actual[0])
print("Matches expected:", np.allclose(actual, expected))

[gemini-code-assist]

Code Review

This pull request adds support for ONNX Pad operator version 19 (_impl_v19) in the Relax frontend, enabling the "wrap" padding mode (circular padding). Unit tests are updated to include a test case for this new mode using opset=19. The feedback suggests refactoring _impl_v19 to delegate non-wrap modes to _impl_v11 to avoid substantial code duplication and improve maintainability.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The implementation of _impl_v19 is almost identical to _impl_v11, introducing significant code duplication. Since the only difference is the support for mode="wrap", we can simplify _impl_v19 by handling the "wrap" mode directly and delegating all other modes to _impl_v11. This reduces duplication and improves maintainability.

    @classmethod
    def _impl_v19(cls, bb, inputs, attr, params):
        pad_mode = attr.get("mode", b"constant").decode("utf-8")
        if pad_mode == "wrap":
            pads = get_constant(inputs[1], params)
            if isinstance(pads, relax.Constant):
                pad_before, pad_after = _np.split(pads.data.numpy(), 2)
                pad_before = _np.ndarray.tolist(pad_before)
                pad_after = _np.ndarray.tolist(pad_after)
            else:
                raise ValueError("Dynamic pads are not supported yet.")
            return bb.emit_te(topi.nn.circular_pad, inputs[0], pad_before, pad_after)

        return cls._impl_v11(bb, inputs, attr, params)

[tlopex]

Sorry for the late reply

I did a search and found that Pad opset 19 supports the optional 4th input axes. When axes is provided, pads has shape [2 * len(axes)], not [2 * rank(data)].

For example, this is valid ONNX:

input rank = 3
axes = [2]
pads = [2, 2]
mode = "wrap"

It means “pad only axis 2 by 2 before and 2 after”. But _impl_v19 currently ignores inputs[3], splits pads directly, and passes:

pad_before = [2]
pad_after = [2]

to topi.nn.circular_pad, which expects pad lists matching the full input rank. So this valid Pad-19 model will either fail with a pad length mismatch or be handled incorrectly.

Could we either support axes by expanding the pads to full-rank form, e.g.

pad_before = [0, 0, 2]
pad_after = [0, 0, 2]

or explicitly reject axes with a clear error? Please also add a test for mode="wrap" with axes so this path is covered.

[napronald]

Updated.

_impl_v19 now handles the optional axes input. When axes is provided, it expands pad_before/pad_after before calling topi.nn.circular_pad. If axes is dynamic, it raises a clear error.

I added a mode="wrap" test case with axes=[2] and pads=[2, 2], alongside a multi-axis case.

[tlopex]

LGTM. Thanks for the contribution!

[tlopex]

Could you please resolve the conflict and change the test to structural check? Thanks!

[napronald]

Resolved and updated the test to use structural checks.