Add support for the reflect padding mode in Conv2D operation (#53)

tenstorrent · Dec 17, 2024 · 05cd0bc · 05cd0bc
1 parent 38d2522
commit 05cd0bc
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/python/tvm/relay/op/contrib/forge/forge_passes.py b/python/tvm/relay/op/contrib/forge/forge_passes.py
@@ -156,9 +156,16 @@ def callback(self, pre, post, node_map):
         pad_width = pad.attrs.pad_width
         top_pad, bottom_pad = pad_width[-2]
         left_pad, right_pad = pad_width[-1]
+
+        pad_mode = pad.attrs.pad_mode
+
+        if pad_mode == "constant":
+            padding = [top_pad, left_pad, bottom_pad, right_pad]
 
-        padding = [top_pad, left_pad, bottom_pad, right_pad]
-
+        if pad_mode == "reflect":
+            act = tvm.relay.op.nn.pad(act, pad_width, pad_mode="reflect")
+            padding = [0,0,0,0]
+
         op_attrs = {**conv_pool.attrs}
         op_attrs["padding"] = padding