Skip to content

[TOPI][Bugfix] Make semantics of empty axis in squeeze consistent with Relay #12596

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
masahi merged 3 commits into from
Aug 26, 2022
Merged

[TOPI][Bugfix] Make semantics of empty axis in squeeze consistent with Relay #12596

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9fae34c
Fix empty axis of `squeeze` in TOPI.
wzh99 Aug 25, 2022
239cf2a
Add test case for `squeeze` with empty `axis`.
wzh99 Aug 26, 2022
76086e1
Add LLVM target for `test_squeeze`.
wzh99 Aug 26, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions include/tvm/topi/transform.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -408,7 +408,7 @@ inline Tensor squeeze(const Tensor& x, Array<Integer> axis, bool atleast1d = fal
std::string name = "T_squeeze", std::string tag = kInjective) {
auto ndim = x->shape.size();
std::vector<int> 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<int>(i));
Expand Down
5 changes: 3 additions & 2 deletions tests/python/topi/python/test_topi_transform.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -940,18 +940,19 @@ 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)
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")
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):
Expand Down