From 411455d36a3748f7576b7d04f9a87acdbac2a184 Mon Sep 17 00:00:00 2001
From: FFChopon <3483776996@qq.com>
Date: Tue, 18 Aug 2026 04:53:56 +0800
Subject: [PATCH] [Relax][Frontend][ONNX] Support Pad mode="wrap" and axes
 input for opset 18

The ONNX Pad spec adds mode="wrap" (circular padding) and the optional
axes input in opset 18. The Relax ONNX frontend only handled these in
_impl_v19 (added by #19827), which dispatches for opset >= 19; models
with opset 18 still resolve to _impl_v11, which rejects wrap with
OpAttributeInvalid and ignores axes (padding the full rank instead),
so legal opset-18 Pad models fail to import.

Add Pad._impl_v18 with the same wrap/axes handling as _impl_v19: expand
the axes input into full-rank pads and dispatch mode="wrap" to
topi.nn.circular_pad, which already implements circular padding.

Validated by differential testing against onnxruntime over 81 legal
Pad models (3 input shapes x all modes x positive/negative pads, plus
axes cases): the 22 previously-rejected cases (wrap x19, constant+axes
x3) now all import and match onnxruntime; no regression on the 59
already-supported cases.

Co-Authored-By: Claude <noreply@anthropic.com>
---
 .../tvm/relax/frontend/onnx/onnx_frontend.py  | 61 +++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 65bd5bfe1a2f..71f10018e0c5 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2974,6 +2974,67 @@ def _impl_v11(cls, bb, inputs, attr, params):
             # edge mode - replicate border values
             return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before, pad_after)
 
+    @classmethod
+    def _impl_v18(cls, bb, inputs, attr, params):
+        # ONNX Pad-18 introduces mode="wrap" and the optional axes input (v13 and
+        # earlier have neither). _impl_v19 (#19827) added wrap/axes support for
+        # opset >= 19, but opset-18 models still resolve to _impl_v11, which
+        # rejects wrap and ignores axes. This method covers opset 18.
+        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)
+
     @classmethod
     def _impl_v19(cls, bb, inputs, attr, params):
         pads = get_constant(inputs[1], params)