From 9fae34cd7360a9f8808a615741f6775303b0ea20 Mon Sep 17 00:00:00 2001 From: wzh99 Date: Thu, 25 Aug 2022 18:29:26 +0800 Subject: [PATCH 1/3] Fix empty axis of `squeeze` in TOPI. --- include/tvm/topi/transform.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h index 81935dd72dda..7accbf86912d 100644 --- a/include/tvm/topi/transform.h +++ b/include/tvm/topi/transform.h @@ -396,7 +396,7 @@ inline Tensor unravel_index(const Tensor& x, const Tensor& shape, std::string na * The removed dimensions must have a constant size of 1. * * \param x The input tensor - * \param axis Indices of the dimensions to remove. If this is empty, + * \param axis Indices of the dimensions to remove. If this is None, * all entries with a constant size of 1 will be removed. * \param atleast1d Whether the output need to be atleast1d. * \param name The name of the operation @@ -408,7 +408,7 @@ inline Tensor squeeze(const Tensor& x, Array axis, bool atleast1d = fal std::string name = "T_squeeze", std::string tag = kInjective) { auto ndim = x->shape.size(); std::vector axis_val; - if (!axis.defined() || axis.size() == 0) { + if (!axis.defined()) { for (size_t i = 0; i < ndim; ++i) { if (IsConstInt(x->shape[i]) && GetConstInt(x->shape[i]) == 1) { axis_val.push_back(static_cast(i)); From 239cf2a9984c18d2ca4fdab892c09f69e5e6cfa6 Mon Sep 17 00:00:00 2001 From: wzh99 Date: Fri, 26 Aug 2022 09:33:17 +0800 Subject: [PATCH 2/3] Add test case for `squeeze` with empty `axis`. --- tests/python/topi/python/test_topi_transform.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/python/topi/python/test_topi_transform.py b/tests/python/topi/python/test_topi_transform.py index c3155c948a8d..266fb1531759 100644 --- a/tests/python/topi/python/test_topi_transform.py +++ b/tests/python/topi/python/test_topi_transform.py @@ -946,6 +946,7 @@ def test_squeeze(): verify_squeeze((1, 2, 1, 4), None) verify_squeeze((1, 1, 1, 4), (1, 2)) verify_squeeze((1, 1, 1, 1), None) + verify_squeeze((1, 1, 1, 1), ()) # a special case to trigger inline let expression A = te.placeholder((2,), "float32", "A") From 76086e1ec56fcf8caa1153db4efa92051885d272 Mon Sep 17 00:00:00 2001 From: wzh99 Date: Fri, 26 Aug 2022 09:40:14 +0800 Subject: [PATCH 3/3] Add LLVM target for `test_squeeze`. --- tests/python/topi/python/test_topi_transform.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/topi/python/test_topi_transform.py b/tests/python/topi/python/test_topi_transform.py index 266fb1531759..dd5ad1b11926 100644 --- a/tests/python/topi/python/test_topi_transform.py +++ b/tests/python/topi/python/test_topi_transform.py @@ -940,7 +940,7 @@ def test_where(): verify_where((1, 2, 3, 4)) -@tvm.testing.requires_gpu +@tvm.testing.uses_gpu def test_squeeze(): verify_squeeze((1, 2, 3, 4), 0) verify_squeeze((1, 2, 1, 4), None) @@ -952,7 +952,7 @@ def test_squeeze(): A = te.placeholder((2,), "float32", "A") E = topi.squeeze(A) C = te.compute((1,), lambda i: E[(2 * A[0] - 1).astype("int32")]) - for target in ["cuda", "opencl"]: + for target in ["llvm", "cuda", "opencl"]: dev = tvm.device(target, 0) if tvm.testing.device_enabled(target): with tvm.target.Target(target):