Adds support for reflect mode in the pad operation

Author: pranavm-nvidia

tripy/nvtripy/frontend/ops/pad.py

@@ -21,7 +21,7 @@
 from nvtripy.common.exception import raise_error
 from nvtripy.frontend.ops import utils as op_utils
 from nvtripy.trace.ops.shape import Shape
-from nvtripy.trace.ops.slice import SliceFill
+from nvtripy.trace.ops.slice import SliceFill, SliceReflect
 from nvtripy.types import IntLike
 from nvtripy.utils import wrappers
 
@@ -46,22 +46,36 @@ def pad(
             of ``input``. Each element of ``pad`` is a tuple of integers or :class:`DimensionSize` s ``(low, high)``,
             which represents the padding sizes before the lowest index and after the highest index at
             the corresponding dimension.
-        mode: The padding mode. Only "constant" is supported.
-        value: The padding value for "constant" mode.
+        mode: The padding mode. Must be one of:
+
+            - ``"constant"``: Pads with a constant value (default).
+            - ``"reflect"``: Pads by reflecting the input tensor.
+        value: The padding value for "constant" mode. Has no effect for other modes.
 
     Returns:
         The padded tensor.
 
     .. code-block:: python
         :linenos:
-        :caption: Constant padding.
+        :caption: Constant Padding
 
         input = tp.reshape(tp.arange(6, dtype=tp.float32), (2, 3))
         output = tp.pad(input, [(1, 0), (0, 1)])
 
         input_np = np.arange(6, dtype=np.float32).reshape((2, 3)) # doc: omit
         expected = np.pad(input_np, ((1, 0), (0, 1))) # doc: omit
         assert np.array_equal(cp.from_dlpack(output).get(), expected)
+
+    .. code-block:: python
+        :linenos:
+        :caption: Reflect Padding
+
+        input = tp.reshape(tp.arange(6, dtype=tp.float32), (2, 3))
+        output = tp.pad(input, [(1, 1), (1, 1)], mode="reflect")
+
+        input_np = np.arange(6, dtype=np.float32).reshape((2, 3)) # doc: omit
+        expected = np.pad(input_np, ((1, 1), (1, 1)), mode="reflect") # doc: omit
+        assert np.array_equal(cp.from_dlpack(output).get(), expected)
     """
     from nvtripy.frontend.ops.cast import cast
     from nvtripy.frontend.tensor import Tensor
@@ -72,7 +86,7 @@ def pad(
             [f"Got pad={pad}, ", f" input's rank={input.rank}"],
         )
 
-    supported_modes = {"constant"}
+    supported_modes = {"constant", "reflect"}
     if mode not in supported_modes:
         raise_error(
             "Unsupported padding mode.",
@@ -89,4 +103,9 @@ def pad(
     sizes = input_shape + padding_lows + padding_highs
     steps = op_utils.tensor_from_shape_like([1] * input.rank)
 
-    return op_utils.create_op(SliceFill, [input, starts, sizes, steps, cast(Tensor(value), dtype=input.dtype)])
+    inputs = [input, starts, sizes, steps]
+    if mode == "constant":
+        inputs.append(cast(Tensor(value), dtype=input.dtype))
+
+    OpType = {"constant": SliceFill, "reflect": SliceReflect}[mode]
+    return op_utils.create_op(OpType, inputs)

tripy/nvtripy/trace/ops/slice.py

@@ -22,40 +22,33 @@
 from nvtripy.trace.ops.base import TraceOp
 
 
-@dataclass(repr=False)
-class Slice(TraceOp):
+def make_slice_op(name, mode):
+    @dataclass(repr=False)
+    class SliceOp(TraceOp):
 
-    infer_rank = op_utils.InferRankPolicies.same_as_input()
+        infer_rank = op_utils.InferRankPolicies.same_as_input()
 
-    def infer_dtypes(self):
-        self.outputs[0].dtype = self.inputs[0].dtype
+        def infer_dtypes(self):
+            self.outputs[0].dtype = self.inputs[0].dtype
 
-    def to_mlir(self, inputs, outputs):
-        assert len(inputs) == 4, "Slice operation must have exactly 4 inputs."
+        def to_mlir(self, inputs, outputs):
+            mode_attr = tensorrt.SliceModeAttr.get(mode)
 
-        return [tensorrt.slice(inputs[0], start=inputs[1], size=inputs[2], stride=inputs[3])]
+            return [
+                tensorrt.slice(
+                    inputs[0],
+                    start=inputs[1],
+                    size=inputs[2],
+                    stride=inputs[3],
+                    fill=inputs[4] if len(inputs) > 4 else None,
+                    mode=mode_attr,
+                )
+            ]
 
+    SliceOp.__name__ = name
+    return SliceOp
 
-@dataclass(repr=False)
-class SliceFill(TraceOp):
 
-    infer_rank = op_utils.InferRankPolicies.same_as_input()
-
-    def infer_dtypes(self):
-        self.outputs[0].dtype = self.inputs[0].dtype
-
-    def to_mlir(self, inputs, outputs):
-        assert len(inputs) == 5, "SliceFill operation must have exactly 5 inputs."
-
-        mode_attr = tensorrt.SliceModeAttr.get("kFILL")
-
-        return [
-            tensorrt.slice(
-                inputs[0],
-                start=inputs[1],
-                size=inputs[2],
-                stride=inputs[3],
-                fill=inputs[4],
-                mode=mode_attr,
-            )
-        ]
+Slice = make_slice_op("Slice", "kDEFAULT")
+SliceFill = make_slice_op("SliceFill", "kFILL")
+SliceReflect = make_slice_op("SliceReflect", "kREFLECT")

tripy/tests/integration/test_pad.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -45,3 +45,12 @@ def test_pad_tensor(self, eager_or_compiled):
         expected = np.pad(inp, ((0, 2), (3, 0)))
 
         assert np.array_equal(cp.from_dlpack(out).get(), expected)
+
+    @pytest.mark.parametrize("pad", [((1, 0), (0, 1)), ((2, 1), (1, 2))])
+    def test_pad_reflect(self, pad, eager_or_compiled):
+        inp = np.arange(6, dtype=np.float32).reshape((2, 3))
+
+        out = eager_or_compiled(tp.pad, tp.Tensor(inp), pad, mode="reflect")
+        expected = np.pad(inp, pad, mode="reflect")
+
+        assert np.array_equal(cp.from_dlpack(out).get(), expected)