Hi all, I found there’s a is_dynamic_shape bool variable in relay/op/tensor/reduce.cc(https://github.com/apache/tvm/blob/main/src/relay/op/tensor/reduce.cc#L276-L301).
template <typename AttrsType>
inline std::vector<IndexExpr> ReduceShapeImpl(const std::vector<IndexExpr>& in_shape,
const AttrsType* param,
const TypeReporter& reporter) {
uint32_t indim = in_shape.size();
auto r_axes = GetReduceAxes(indim, param->axis, param->exclude);
if (!r_axes.size()) {
return in_shape;
}
auto max_shape = tir::make_const(DataType::Int(64), 1);
bool is_dynamic_input = false;
for (int64_t axis : r_axes) {
if (in_shape[axis].as<IntImmNode>()) {
max_shape *= in_shape[axis];
} else {
is_dynamic_input = true;
break;
}
}
if (is_dynamic_input) {
ICHECK(reporter->Assert(
max_shape < tir::make_const(DataType::Int(64), std::numeric_limits<int32_t>::max())))
<< "The maximum possible index of reduced shape cannot be more than int32 max.";
}
This is_dynamic_shape is used to check the max_shape < int64_MAX, but there’s a problem: a break in r_axes for-loop, which make max_shape less than real reduced shape when meeting dynamic shape?
Thus, I think this break in r_axes for-loop need to remove. Besides, if break in r_axes for-loop removed, I think this assert is not enough safe, either.
if (is_dynamic_input) {
ICHECK(reporter->Assert(
max_shape < tir::make_const(DataType::Int(64), std::numeric_limits<int32_t>::max())))
<< "The maximum possible index of reduced shape cannot be more than int32 max.";
}