[Bug] [Relay] Problems observed working with `relay.squeeze`

[padreofthegame]

Thanks for participating in the TVM community! We use https://discuss.tvm.ai for any general usage questions and discussions. The issue tracker is used for actionable items such as feature proposals discussion, roadmaps, and bug tracking. You are always welcomed to post on the forum first 😸

Issues that are inactive for a period of time may get closed. We adopt this policy so that we won't lose track of actionable issues that may fall at the bottom of the pile. Feel free to reopen a new one if you feel there is an additional problem that needs attention when an old one gets closed.

Hello everyone,

I am new here so I am not completely familiar with the way of communication here, but I will try to summarize my observations related to relay.squeeze functionality. This report may be a more comprehensive version of #11697, but also it contain additional scenario which is not related to that one.

relay.squeeze with argument axis=[]

Scenario 1.

Only relay.squeeze(axis=[])

def @main(%x: Tensor[(1, 1, 1, 1, 1, 2, 1, 3, 4, 5), float32]) { squeeze(%x, axis=[]) }

Expected behavior

Program should compile without errors.

Actual behavior

Program compile and the shape of the result tensor is the same as the input tensor shape.

Input tensor shape: (1, 1, 1, 1, 1, 2, 1, 3, 4, 5)
Output tensor shape: (1, 1, 1, 1, 1, 2, 1, 3, 4, 5)

Looks like parameter axis=[] does not change the shape of the relay.squeeze input tensor. (this is the assumption that may make problems later)

Scenario 2.

relay.squeeze(axis=[]) + relay.squeeze(axis=[...])

def @main(%x: Tensor[(1, 1, 1, 1, 1, 2, 1, 3, 4, 5), float32]) { %0 = squeeze(%x, axis=[]); squeeze(%0, axis=[0]) }

Expected behavior

Program should compile without errors as long as we send the appropriate parameters for the axis in second squeeze call. (for example using axis=[5] should be an error because dimension which is not equal to 1 could not be squeezed).

Actual behavior

Program does not compile, and there is 2 types of unexpected errors:
First one is obtained for parameter axis=[0], and the error message is:

TVMError:
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
Check failed: GetConstInt(x->shape[val]) == 1 (2 vs. 1) : Dimension 0 must have size 1

Error of the same type is obtained using any of numbers 1, 2, 3 instead of 0 for the axis in previous example.

Second one is obtained for parameter axis=[4], and the error message is:

TVMError:
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
Check failed: (0 <= i && i < p->size_) is false: IndexError: indexing 4 on an array of size 4

Error of the same type is also obtained using number 6 instead of 0 for the axis in previous example.

For other possible options for axis (which are 5, 7, 8 and 9) we get expected error since in input tensor those dimensions are not equal to 1.

From these error descriptions it looks like that using parameter axis=[] actually changes the shape of tensor internally (in CPP) in a same way as it will be when we send axis=None, but Python for some reason does not see this.

Scenario 3.

relay.squeeze(axis=[]) + relay.nn.bias_add

def @main(%x: Tensor[(1, 1, 1, 1, 5, 2, 1, 3, 4, 5), float32], %b: Tensor[(5), float32]) { %0 = squeeze(%x, axis=[]); nn.bias_add(%0, %b, axis=4) }

Expected behavior

Program should compile without errors as long as we send the appropriate parameters for the axis in bias_add call. (for example using axis=0 should be an error because corresponding dimension must be equal to the dimension of the tensor b which should be added).

Actual behavior

Program does not compile, and there is additional 2 types of unexpected errors:
First one is obtained for parameter axis=4, and the error message is:

TVMError:
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
Check failed: ret == 0 (-1 vs. 0) : Assert fail: (5 == tir.tvm_struct_get(arg2, 0, 4)), arg2.ndim is expected to equal 5

Second one is obtained for parameter axis=9, and the error message is:

TVMError:
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
Check failed: (num_newaxis >= 0) is false: expand_dims only accepts num_newaxis >= 0, but got num_newaxis = -5

For other possible options for axis (which are 0, 1, 2, 3, 5, 6, 7 and 8) we get expected error since in input tensor those dimensions are not equal to the shape of tensor b.

From these error descriptions it looks like that using parameter axis=[] actually changes the input tensor as it is described in Scenario 2.

relay.squeeze with argument axis=Constant

Scenario 1.

Using axis = Constant([]) instead of axis = [] will result with same errors described above, when used in the same scenario.

Scenario 2.

Using axis = Constant(2) where the argument of the Constant is tvm.nd.array created only using integer (in this example that integer is 2).

Expected behavior

Following code should pass without errors:

import numpy as np
import tvm
from tvm import relay, transform, cpu, IRModule

x_shape = (1, 1, 1, 1, 1, 2, 1, 3, 4, 5)
x = relay.var('x', shape=x_shape)

arr = np.array(0)
ndarr = tvm.nd.array(arr)
const = relay.Constant(ndarr)

y = relay.squeeze(x, axis=const)
mod = IRModule.from_expr(y)

Actual behavior

Program does not compile with the following error:

Traceback (most recent call last):
File "/home/padre/TVM/Issues/Relay/Bugs/issue_with_relay_squeeze/relay_squeeze_test_constant.py", line 104, in <module>
y = relay.squeeze(x, axis=const)

File "/home/padre/TVM_full_repo/tvm/python/tvm/relay/op/transform.py", line 221, in squeeze
axis = list(axis.data.numpy())

TypeError: iteration over a 0-d array

Since the functionality of the tvm.nd.array is based on np.array, which could have dimension () and which in that case is not iterable, this scenario should be extra covered.

Environment

Ubuntu 20.04, TVM commit f5f5a75 compiled with LLVM support.

Steps to reproduce

Code may be the same as in #11697, just relay definition and corresponding GraphModule.run function should be changed according to the mentioned scenario.

Solution

All mentioned errors could be solved implementing few small changes in squeeze function defined inside python/tvm/relay/op/transform.py file.

1. change

Instead of:

if isinstance(axis, Constant):
axis = list(axis.data.numpy())

should be:

if isinstance(axis, Constant):
if axis.data.shape:
axis = list(axis.data.numpy())
else:
axis = [axis.data.numpy().item()]

Internal if statement could also be changed with try except block, depending on the speed.

2. change

Depending on the wanted functionality in front of the line:

return _make.squeeze(data, axis)

should stay (if using axis=[] should be same as using axis=None):

if not axis:
axis = None

or (if using axis=[] should not change the input tensor data):

if axis == []:
return data

[padreofthegame]

@ganler @jroesch, or maybe you know someone who may take a look?

[ganler]

Thanks @padreofthegame. Would you like to provide minimal reproducible for us to check the issues you mentioned?

[ganler]

Had a closer look. It seems those are not bugs as you might misunderstood the axis attribute. According to relay.squeeze doc, axis should be None or a list of integer or an expression.

axis (None or List[int] or Expr) – The set of axes to remove.

If axis is [] it means none of the axies will be squeezed so that it is fairly correct to have identical output shapes. If you meant to have dim-1 axies automatically squeezed, it should be axis=None.

TypeError: iteration over a 0-d array

I also don't think it is related to #11697 as that happens at runtime and the cases here happens even before compile time.

For axis=relay.Constant(tvm.nd.array(np.array(0))), the np.array(0) is a scalar instead of an integer list as is prompted by the error message.

[padreofthegame]

Thanks @ganler for your quick and detailed response. As I already said I am new here and I am really interested in the work you did, and you are doing at the moment, and I would like to help if it is possible. So I decided to chose some package to look at, try some features and comment problems that I came up working with it (actually problems may appear due to misunderstanding of some parameters).

But here I found some features really interesting:

If axis is [] it means none of the axies will be squeezed so that it is fairly correct to have identical output shapes.

I also thought logically about that and in a same way, but then I came up with errors described in a Scenario 2. and 3. of the first Test. Looking at Scenario 2. basically if we think that axis=[] does nothing to input tensor (as it is shown from result in Scenario 1.) there should be no problems for example to squeeze that tensor again after first squeeze. But here things becomes to be interesting. From first error in Scenario 2. we see that it pass pre-compile check in python (since 0 is valid dimension to squeeze) but later raise an error that corresponding dimension is not equal to 1. The same thing happens if instead of 0 we use 1, 2 or 3. Also, if we use axis to be 4 or 6, which also pass the pre-compile check, we got the error described second in Scenario 2. that we are indexing out of bounds of the tensor.

From this, the logical explanation I came up with is that internally axis=[] changes shape of the input tensor as it was None, because if we assume that output tensor of first squeeze would be `[2, 3, 4, 5] all errors described above should be expected. And this is the main reason I decided to write here.

On the other hand I also wrote about feature I found working with parameter axis of type Constant because it occur in the same function. In my opinion case Constant(0) should be covered because Constant is a valid axis type, and 0 should be valid axis dimension. Also looking deeply in transform package I also find that for example full function allow similar types for parameter shape as squeeze for axis, but also cover the case when the shape is int for example. What is your opinion about that since you are much longer in this community? If needed, I could go through similar functions in the package, and try to identify potentially similar features.

Maybe, also adding possibility for axis to be of type int should be good thing, since very often only one dimension should be squeezed?

I also provided and tested solution I mentioned above, and it looks like it eliminate errors mentioned. I would be pleased to push the solution to repository, but at first I wanted to comment with someone relevant about this.

[wzh99]

@padreofthegame @ganler Hey guys, I am the author of #11697. I appreciate your efforts in solving problems related to squeeze. Here I mainly want to share my analysis for the axis=[] case. After checking the source code in Relay and TOPI, I think the root cause lies in the difference between how Relay and TOPI handle empty axis.

In Relay type inference of squeeze, axis=[] does nothing to the input tensor:

tvm/src/relay/op/tensor/transform.cc

Lines 2301 to 2311 in 3983a47

	if (!param->axis.defined()) {
	for (const auto& e : data->shape) {
	if (!e.as<IntImmNode>()) {
	LOG(FATAL) << "axis needs to be defined for dynamic input.";
	}
	const int64_t* axis_ptr = tir::as_const_int(e);
	ICHECK(axis_ptr != nullptr) << "the axes attribute must be concrete";
	if (*axis_ptr != 1) {
	result_shape.push_back(e);
	}
	}

In TOPI implementation of squeeze, however, axis=[] is the same as axis=None, which squeezes all possible axes:

tvm/include/tvm/topi/transform.h

Lines 411 to 416 in 3983a47

	if (!axis.defined() || axis.size() == 0) {
	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));
	}
	}

That explains why the tensor shape in compiled module is not consistent with the one in Relay.

To fix this, I think we should modify the C++ code in either Relay or TOPI. It should be decided whether axis=[] means doing nothing, or squeezing all possible axes. Maybe we can consult @ganler and other TVM developers for determination. Personally I do not completely agree with the solution proposed by @padreofthegame, as this is just a workaround at Python side and does not completely solve this problem.