This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 04c6d7a45f [Unity][Fix] Infer Layout must support negative axes
(#14365)
04c6d7a45f is described below
commit 04c6d7a45f6d06fb0b4a8e34127856583aac1df4
Author: Prakalp Srivastava <[email protected]>
AuthorDate: Wed Mar 22 03:07:52 2023 -0400
[Unity][Fix] Infer Layout must support negative axes (#14365)
This PR fixes a bug in Infer Layout for reduction ops in when `axis` has
negative indices.
---
src/relax/op/tensor/statistical.cc | 2 +-
.../python/relax/test_transform_convert_layout.py | 41 ++++++++++++++++++++++
2 files changed, 42 insertions(+), 1 deletion(-)
diff --git a/src/relax/op/tensor/statistical.cc
b/src/relax/op/tensor/statistical.cc
index 4de8e3dd63..57a31dbb44 100644
--- a/src/relax/op/tensor/statistical.cc
+++ b/src/relax/op/tensor/statistical.cc
@@ -107,7 +107,7 @@ InferLayoutOutput InferLayoutStatistical(const Call& call,
std::string axis_str(ndim, '0');
for (const auto& iter : axis) {
- axis_str[iter->value] = '1';
+ axis_str[(iter->value + ndim) % ndim] = '1';
}
for (int i = 0, j = 0; i < ndim; ++i) {
if (axis_str[i] != '1') {
diff --git a/tests/python/relax/test_transform_convert_layout.py
b/tests/python/relax/test_transform_convert_layout.py
index 1b65c0812a..78a1b166dd 100644
--- a/tests/python/relax/test_transform_convert_layout.py
+++ b/tests/python/relax/test_transform_convert_layout.py
@@ -624,6 +624,47 @@ def test_conv2d_sum_keepdim():
verify(Input, Expected)
+def test_conv2d_sum_negative_dims():
+ @I.ir_module
+ class Input:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3, 28, 28), "float32"), w: R.Tensor((4, 3, 3, 3),
"float32")
+ ) -> R.Tensor(None, "float32", ndim=4):
+ with R.dataflow():
+ gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w,
out_dtype="float32")
+ gv2: R.Tensor((2, 4), "float32") = R.sum(gv, axis=[-2, -1])
+ R.output(gv2)
+ return gv2
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((4, 3,
3, 3), dtype="float32")
+ ) -> R.Tensor(None, dtype="float32", ndim=4):
+ with R.dataflow():
+ lv: R.Tensor((2, 28, 28, 3), dtype="float32") =
R.permute_dims(x, axes=[0, 2, 3, 1])
+ lv1: R.Tensor((4, 3, 3, 3), dtype="float32") =
R.permute_dims(w, axes=[0, 2, 3, 1])
+ gv: R.Tensor((2, 26, 26, 4), dtype="float32") = R.nn.conv2d(
+ lv,
+ lv1,
+ strides=[1, 1],
+ padding=[0, 0, 0, 0],
+ dilation=[1, 1],
+ groups=1,
+ data_layout="NHWC",
+ kernel_layout="OHWI",
+ out_layout="NHWC",
+ out_dtype="float32",
+ )
+ gv2: R.Tensor((2, 4), dtype="float32") = R.sum(gv, axis=[1, 2])
+ R.output(gv2)
+ return gv2
+
+ verify(Input, Expected)
+
+
def test_conv2d_transpose():
@I.ir_module
class Input:
