Skip to content

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

Open
napronald wants to merge 1 commit into
apache:mainfrom
napronald:relax-onnx-pad-wrap
Open

[Relax][Frontend][ONNX] Add support for Pad mode="wrap" for opset 19#19827
napronald wants to merge 1 commit into
apache:mainfrom
napronald:relax-onnx-pad-wrap

Conversation

@napronald

Copy link
Copy Markdown

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 gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +2762 to +2793
@classmethod
def _impl_v19(cls, bb, inputs, attr, params):
pads = get_constant(inputs[1], params)
constant_value = get_constant(inputs[2], params)
if constant_value is not None:
constant_value = constant_value.data.numpy().item()
else:
constant_value = 0.0

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.")

pad_mode = attr.get("mode", b"constant").decode("utf-8")
if pad_mode not in ["constant", "edge", "reflect", "wrap"]:
raise tvm.error.OpAttributeInvalid(
"Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.'
)

if pad_mode == "constant":
return bb.emit_te(topi.nn.pad, inputs[0], pad_before, pad_after, constant_value)
elif pad_mode == "reflect":
return bb.emit_te(topi.nn.mirror_pad, inputs[0], pad_before, pad_after, "REFLECT")
elif pad_mode == "wrap":
return bb.emit_te(topi.nn.circular_pad, inputs[0], pad_before, pad_after)
else:
# edge mode - replicate border values
return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before, pad_after)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant