Skip to content

Commit b0e4e5c

Browse files
committed
support negtive dims in flip
1 parent fc9960a commit b0e4e5c

File tree

2 files changed

+5
-0
lines changed

2 files changed

+5
-0
lines changed

py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py

+1
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,7 @@ def flip(
241241

242242
shape = input.shape
243243
rank = len(shape)
244+
dims = get_positive_dim(dims, rank)
244245

245246
for i in range(rank):
246247
if i in dims:

tests/py/dynamo/conversion/test_flip_aten.py

+4
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,16 @@ class TestFlipConverter(DispatchTestCase):
1010
@parameterized.expand(
1111
[
1212
((3,), [0]),
13+
((3,), [-1]),
1314
((3, 3), [0, 1]),
15+
((3, 3), [-2, 1]),
1416
((2, 3, 4), [0]),
1517
((3, 3, 3), (0, 1)),
1618
((2, 3, 4), [0, 1, 2]),
19+
((2, 3, 4), [-3, -2, -1]),
1720
((3, 3, 3, 3), [0]),
1821
((2, 3, 4, 5), [0, 1, 2, 3]),
22+
((2, 3, 4, 5), [-4, 1, -2, 3]),
1923
]
2024
)
2125
def test_flip(self, shape, dims):

0 commit comments

Comments
 (0)