This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d5c6f2d484 [Relax][Frontend][ONNX] Add support for Pad mode="wrap" for
opset 19 (#19827)
d5c6f2d484 is described below
commit d5c6f2d484264fed2ba172693a39f1b144243be6
Author: Ronald Nap <[email protected]>
AuthorDate: Thu Jul 9 15:23:26 2026 -0700
[Relax][Frontend][ONNX] Add support for Pad mode="wrap" for opset 19
(#19827)
## 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
```text
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
```python
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))
```
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 57 ++++++++++++
tests/python/relax/test_frontend_onnx.py | 119 ++++++++++++++++++++----
2 files changed, 160 insertions(+), 16 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index f460418315..e9b951e0f9 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2752,6 +2752,63 @@ class Pad(OnnxOpConverter):
# edge mode - replicate border values
return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before,
pad_after)
+ @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.")
+
+ axes_input = inputs[3] if len(inputs) > 3 else None
+ if axes_input is not None:
+ axes_const = get_constant(axes_input, params)
+ if not isinstance(axes_const, relax.Constant):
+ raise ValueError("Dynamic axes are not supported for Pad yet.")
+
+ axes = axes_const.data.numpy().tolist()
+ if len(pad_before) != len(axes):
+ raise ValueError(
+ f"Pad expects pads length 2 * len(axes), got "
+ f"{len(pad_before) + len(pad_after)} pads and {len(axes)}
axes."
+ )
+
+ rank = _get_known_tensor_rank(inputs[0])
+ if rank is None:
+ raise ValueError("Pad with axes requires a statically known
input rank.")
+
+ axes = _normalize_constant_axes([int(a) for a in axes], rank,
"Pad")
+ full_before = [0] * rank
+ full_after = [0] * rank
+ for i, ax in enumerate(axes):
+ full_before[ax] = pad_before[i]
+ full_after[ax] = pad_after[i]
+ pad_before, pad_after = full_before, full_after
+
+ 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)
+
class Tile(OnnxOpConverter):
"""Converts an onnx Tile node into an equivalent Relax expression."""
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 9058cdf70a..f96984e423 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -6290,9 +6290,18 @@ def test_attention(dynamic):
)
-def _make_pad_expected_ir(input_shape, pads, mode="constant", value=0.0,
opset=14):
+def _make_pad_expected_ir(input_shape, pads, mode="constant", value=0.0,
opset=14, axes=None):
len_dim = len(pads) // 2
np_pads = [(pads[i], pads[i + len_dim]) for i in range(len_dim)]
+
+ if axes is not None:
+ rank = len(input_shape)
+ full_pads = [(0, 0)] * rank
+ for i, axis in enumerate(axes):
+ axis = axis if axis >= 0 else axis + rank
+ full_pads[axis] = np_pads[i]
+ np_pads = full_pads
+
if mode == "constant":
out_shape = np.pad(
np.empty(input_shape, dtype=np.float32),
@@ -6307,6 +6316,7 @@ def _make_pad_expected_ir(input_shape, pads,
mode="constant", value=0.0, opset=1
input_shape = tuple(input_shape)
out_shape = tuple(out_shape)
pads_shape = (len(pads),)
+ axes_shape = None if axes is None else (len(axes),)
if mode == "constant" and opset >= 11:
@@ -6468,6 +6478,60 @@ def _make_pad_expected_ir(input_shape, pads,
mode="constant", value=0.0, opset=1
return ExpectedPadEdgeAttrs
+ if mode == "wrap" and opset >= 19:
+ if axes is None:
+
+ @I.ir_module
+ class ExpectedPadWrapWithInputs:
+ @T.prim_func(private=True, s_tir=True)
+ def circular_pad(input: T.handle, CircularPadInput: T.handle):
+ T.evaluate(0)
+
+ @R.function
+ def main(
+ input: R.Tensor(input_shape, dtype="float32"),
+ pads: R.Tensor(pads_shape, dtype="int64"),
+ ) -> R.Tensor(out_shape, dtype="float32"):
+ R.func_attr({"num_input": 1})
+ cls = ExpectedPadWrapWithInputs
+ with R.dataflow():
+ lv = R.call_tir(
+ cls.circular_pad,
+ (input,),
+ out_ty=R.Tensor(out_shape, dtype="float32"),
+ )
+ gv: R.Tensor(out_shape, dtype="float32") = lv
+ R.output(gv)
+ return gv
+
+ return ExpectedPadWrapWithInputs
+
+ @I.ir_module
+ class ExpectedPadWrapWithAxes:
+ @T.prim_func(private=True, s_tir=True)
+ def circular_pad(input: T.handle, CircularPadInput: T.handle):
+ T.evaluate(0)
+
+ @R.function
+ def main(
+ input: R.Tensor(input_shape, dtype="float32"),
+ pads: R.Tensor(pads_shape, dtype="int64"),
+ axes: R.Tensor(axes_shape, dtype="int64"),
+ ) -> R.Tensor(out_shape, dtype="float32"):
+ R.func_attr({"num_input": 1})
+ cls = ExpectedPadWrapWithAxes
+ with R.dataflow():
+ lv = R.call_tir(
+ cls.circular_pad,
+ (input,),
+ out_ty=R.Tensor(out_shape, dtype="float32"),
+ )
+ gv: R.Tensor(out_shape, dtype="float32") = lv
+ R.output(gv)
+ return gv
+
+ return ExpectedPadWrapWithAxes
+
raise AssertionError(f"No Pad expected IR for mode={mode}, opset={opset}")
@@ -6476,21 +6540,39 @@ def test_pad(dynamic):
if dynamic:
pytest.skip("Dynamic pad not supported")
- def verify_pad(input_shape, pads, expected, mode="constant", value=0.0):
+ def verify_pad(input_shape, pads, expected, mode="constant", value=0.0,
opset=14, axes=None):
len_dim = len(pads) // 2
np_pads = [(pads[i], pads[i + len_dim]) for i in range(len_dim)]
- pads = np.array(pads)
+
+ if axes is not None:
+ rank = len(input_shape)
+ full_pads = [(0, 0)] * rank
+ for i, axis in enumerate(axes):
+ axis = axis if axis >= 0 else axis + rank
+ full_pads[axis] = np_pads[i]
+ np_pads = full_pads
+
+ pads = np.array(pads, dtype=np.int64)
# onnx graph
- if mode in ["edge", "reflect"]:
+ if mode in ["edge", "reflect", "wrap"]:
outdata = np.pad(np.empty(input_shape, dtype=np.float32),
pad_width=np_pads, mode=mode)
- node = helper.make_node("Pad", inputs=["input", "pads"],
outputs=["output"], mode=mode)
+
+ node_inputs = ["input", "pads"]
+ initializer = [helper.make_tensor("pads", TensorProto.INT64,
(len(pads),), pads)]
+
+ if axes is not None:
+ axes = np.array(axes, dtype=np.int64)
+ node_inputs = ["input", "pads", "", "axes"]
+ initializer.append(helper.make_tensor("axes",
TensorProto.INT64, (len(axes),), axes))
+
+ node = helper.make_node("Pad", inputs=node_inputs,
outputs=["output"], mode=mode)
graph = helper.make_graph(
[node],
"pad_test",
inputs=[
helper.make_tensor_value_info("input", TensorProto.FLOAT,
list(input_shape))
],
- initializer=[helper.make_tensor("pads", TensorProto.INT64,
(len(pads),), pads)],
+ initializer=initializer,
outputs=[
helper.make_tensor_value_info("output", TensorProto.FLOAT,
list(outdata.shape))
],
@@ -6523,8 +6605,8 @@ def test_pad(dynamic):
],
)
model = helper.make_model(graph, producer_name="pad_test")
- model.opset_import[0].version = 14
- tvm_model = from_onnx(model, opset=14, keep_params_in_input=True)
+ model.opset_import[0].version = opset
+ tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
tvm_model["main"] = tvm_model["main"].without_attr("params")
expected = tvm.IRModule(expected.functions)
for gv in expected.get_global_vars():
@@ -6532,20 +6614,25 @@ def test_pad(dynamic):
expected.update_func(gv, tvm_model[gv.name_hint])
tvm.ir.assert_structural_equal(tvm_model, expected)
- for input_shape, pads, mode, value in [
- ((2, 2), [0, 1, 0, 0], "constant", 0.0),
- ((2, 3), [1, 0, 0, 1], "constant", 0.0),
- ((3, 2), [0, 0, 1, 0], "constant", 5.0),
- ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "reflect", 0.0),
- ((2, 3), [1, 1, 1, 1], "edge", 0.0),
- ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "edge", 0.0),
+ for input_shape, pads, mode, value, opset, axes in [
+ ((2, 2), [0, 1, 0, 0], "constant", 0.0, 14, None),
+ ((2, 3), [1, 0, 0, 1], "constant", 0.0, 14, None),
+ ((3, 2), [0, 0, 1, 0], "constant", 5.0, 14, None),
+ ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "reflect", 0.0, 14, None),
+ ((2, 3), [1, 1, 1, 1], "edge", 0.0, 14, None),
+ ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "edge", 0.0, 14, None),
+ ((1, 3, 4), [0, 0, 2, 0, 0, 2], "wrap", 0.0, 19, None),
+ ((1, 3, 4), [2, 2], "wrap", 0.0, 19, [2]),
+ ((1, 3, 4), [1, 2, 1, 2], "wrap", 0.0, 19, [1, 2]),
]:
verify_pad(
input_shape,
pads,
- _make_pad_expected_ir(input_shape, pads, mode=mode, value=value,
opset=14),
+ _make_pad_expected_ir(input_shape, pads, mode=mode, value=value,
opset=opset, axes=axes),
mode,
value,
+ opset,
+ axes,
)
