[Relax][Frontend][ONNX] Add Pad wrap mode

Author: napronald

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -2759,6 +2759,63 @@ 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_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."""

tests/python/relax/test_frontend_onnx.py

@@ -3656,23 +3656,37 @@ def test_pad(dynamic):
     if dynamic:
         pytest.skip("Dynamic pad not supported")
 
-    def verify_pad(input_shape, pads, mode="constant", value=0.0):
+    def verify_pad(input_shape, pads, mode="constant", value=0.0, opset=14, axes=None):
         indata = np.random.normal(size=input_shape).astype(np.float32)
         #  numpy expect result
         len_dim = len(pads) // 2
         np_pads = [(pads[i], pads[i + len_dim]) for i in range(len_dim)]
         pads = np.array(pads)
         #  onnx graph
-        if mode in ["edge", "reflect"]:
+        if mode in ["edge", "reflect", "wrap"]:
+            if axes is not None:
+                rank = len(input_shape)
+                full_pads = [(0, 0)] * rank
+                for i, ax in enumerate(axes):
+                    full_pads[ax if ax >= 0 else ax + rank] = np_pads[i]
+                np_pads = full_pads
+
             outdata = np.pad(indata, pad_width=np_pads, mode=mode)
-            node = helper.make_node("Pad", inputs=["input", "pads"], outputs=["output"], mode=mode)
+            node_inputs = ["input", "pads"] if axes is None else ["input", "pads", "", "axes"]
+            node = helper.make_node("Pad", inputs=node_inputs, outputs=["output"], mode=mode)
+            initializer = [helper.make_tensor("pads", TensorProto.INT64, (len(pads),), pads)]
+            if axes is not None:
+                axes_arr = np.array(axes, dtype=np.int64)
+                initializer.append(
+                    helper.make_tensor("axes", TensorProto.INT64, (len(axes_arr),), axes_arr)
+                )
             graph = helper.make_graph(
                 [node],
                 "pad_test",
                 inputs=[
                     helper.make_tensor_value_info("input", TensorProto.FLOAT, list(indata.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))
                 ],
@@ -3700,14 +3714,17 @@ def verify_pad(input_shape, pads, mode="constant", value=0.0):
                 ],
             )
         model = helper.make_model(graph, producer_name="pad_test")
-        check_correctness(model)
+        check_correctness(model, opset=opset)
 
     verify_pad((2, 2), [0, 1, 0, 0], "constant", 0.0)
     verify_pad((2, 3), [1, 0, 0, 1], "constant", 0.0)
     verify_pad((3, 2), [0, 0, 1, 0], "constant", 5.0)
     verify_pad((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "reflect")
     verify_pad((2, 3), [1, 1, 1, 1], "edge")
     verify_pad((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "edge")
+    verify_pad((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "wrap", opset=19)
+    verify_pad((1, 3, 4), [2, 2], "wrap", opset=19, axes=[2])
+    verify_pad((1, 3, 4, 5), [1, 1, 1, 1], "wrap", opset=19, axes=[1, 3])
 
 
 @pytest.mark.parametrize("dynamic", [True, False])