Skip to content

GitLab

Commit 82c3e836 authored by zhuyue's avatar zhuyue
Browse files

Issue/714 - feat(rearrange): add test cases for edge conditions.

parent e1c836b8
Show whitespace changes
Inline Side-by-side
Showing with 5 additions and 0 deletions
test/infiniop/rearrange.py
View file @ 82c3e836
...@@ -76,6 +76,9 @@ _TEST_CASES = [ ...@@ -76,6 +76,9 @@ _TEST_CASES = [
column_major_strides((3, 4, 50, 50, 5, 7)), # y_stride column_major_strides((3, 4, 50, 50, 5, 7)), # y_stride
), ),
((15, 10752), (0, 1), (10752, 1)), ((15, 10752), (0, 1), (10752, 1)),
((2, 2, 2, 2, 2, 2), (4, 8, 16, 32, 64, 128), (64, 32, 16, 8, 4, 2)), # shape # x_stride # y_stride
((8, 4, 20, 64), (5120, 64, 256, 1), None), # shape # x_stride # y_stride
((8, 4, 20, 64), (5120, 64, 256, 1), (1048576, 262144, 64, 1)), # shape # x_stride # y_stride
] ]
# Data types used for testing # Data types used for testing
...@@ -94,6 +97,8 @@ NUM_ITERATIONS = 1000 ...@@ -94,6 +97,8 @@ NUM_ITERATIONS = 1000
def rearrange_torch(y, x, x_shape, y_stride): def rearrange_torch(y, x, x_shape, y_stride):
if y_stride is None:
y_stride = row_major_strides(x_shape)
y.set_(y.untyped_storage(), 0, x_shape, y_stride) y.set_(y.untyped_storage(), 0, x_shape, y_stride)
y.copy_(x.expand_as(y)) y.copy_(x.expand_as(y))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment