[Bug] Type inference of `nn.softmax` does not reject invalid `axis`

[wzh99]

Expected behavior

The following Relay program should NOT pass type inference:

#[version = "0.0.5"]
def @main(%x: Tensor[(4), float32]) {
  nn.softmax(%x, axis=1)
}

The input tensor %x of nn.softmax has only one dimension. The valid range of axis is [-1, 1). axis=1 is obviously invalid in this case.

Actual behavior

This program passes type inference of Relay.

Environment

macOS 12.4. Compiled using Clang 13.1.6 with LLVM support. TVM commit 0df6961.

Steps to reproduce

from tvm import relay, IRModule

x = relay.var('x', shape=(4,))
y = relay.nn.softmax(x, axis=1)
mod = IRModule.from_expr(y)
mod = relay.transform.InferType()(mod)

Possible fix

tvm/src/relay/op/nn/nn.cc

Lines 409 to 423 in 0df6961

	RELAY_REGISTER_OP("nn.softmax")
	.describe(R"code(Softmax layer.
	.. math:: \text{softmax}(x)_i = \frac{exp(x_i)}{\sum_j exp(x_j)}
	.. note::
	This operator can be optimized away for inference.
	- **data**: The input data
	)code" TVM_ADD_FILELINE)
	.set_attrs_type<SoftmaxAttrs>()
	.set_num_inputs(1)
	.add_argument("data", "Tensor", "The input tensor.")
	.set_support_level(1)
	.add_type_rel("Identity", IdentityRel);

In operator registration of nn.softmax, its type relation is set to be IdentityRel. However, nn.softmax has an attribute axis that is not checked by IdentityRel.

A possible fix is to implement a new type relation function that checks axis attribute in SoftmaxAttrs. This type relation also applies to nn.fast_softmax and nn.log_softmax.

[lileidev]

didn't reproduce any error on my linux platform.
env:
Python 3.8.10 (default, Mar 15 2022, 12:22:08)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.

import tvm
print(tvm.version)
0.9.dev0

test code:
from tvm import relay, IRModule

x = relay.var('x', shape=(4, ))
y = relay.nn.softmax(x, axis=1)
mod = IRModule.from_expr(y)
mod = relay.transform.InferType()(mod)
print(mod)

output:
def @main(%x: Tensor[(4), float32] /* ty=Tensor[(4), float32] /) -> Tensor[(4), float32] {
nn.softmax(%x, axis=1) / ty=Tensor[(4), float32] */
}

[wzh99]

@lileidev Thanks for your participation in this issue. This test case is invalid, so an error is actually expected. However, the type inference in Relay does not report any error. This is somehow different from common bug issues, where we typically expect no errors.

To actually reproduce an error, you can try the following code:

from tvm import relay, transform, IRModule

x = relay.var('x', shape=(4,))
y = relay.nn.softmax(x, axis=1)
mod = IRModule.from_expr(y)
mod = relay.transform.InferType()(mod)
with transform.PassContext(opt_level=0):
    lib = relay.build(mod, target='llvm')

And you may see the following error report:

Traceback (most recent call last):
  File "/Users/wzh/tvm-bug/bug_softmax_axis.py", line 10, in <module>
    lib = relay.build(mod, target='llvm')
  File "/Users/wzh/tvm-dev/python/tvm/relay/build_module.py", line 416, in build
    graph_json, runtime_mod, params = bld_mod.build(
  File "/Users/wzh/tvm-dev/python/tvm/relay/build_module.py", line 154, in build
    self._build(mod, raw_targets, executor, runtime, workspace_memory_pools, mod_name)
  File "/Users/wzh/tvm-dev/python/tvm/_ffi/_ctypes/packed_func.py", line 237, in __call__
    raise get_last_ffi_error()
IndexError: Traceback (most recent call last):
  [bt] (8) 9   libtvm.dylib                        0x000000011443e8fa tvm::relay::tec::ScheduleBuilder::Create(tvm::relay::Function const&, std::__1::function<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > (std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> >)>) + 218
  [bt] (7) 8   libtvm.dylib                        0x0000000114445ae3 tvm::relay::tec::LowerToTECompute::Lower(tvm::relay::Function const&, std::__1::function<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > (std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> >)>) + 1123
  [bt] (6) 7   libtvm.dylib                        0x000000011444bef2 tvm::relay::backend::MemoizedExprTranslator<tvm::runtime::Array<tvm::te::Tensor, void> >::VisitExpr(tvm::RelayExpr const&) + 226
  [bt] (5) 6   libtvm.dylib                        0x000000011444c08e tvm::relay::ExprFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::RelayExpr const&)>::VisitExpr(tvm::RelayExpr const&) + 222
  [bt] (4) 5   libtvm.dylib                        0x000000011444c3a8 tvm::NodeFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::RelayExpr const&)>*)>::operator()(tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::RelayExpr const&)>*) const + 312
  [bt] (3) 4   libtvm.dylib                        0x000000011444df88 tvm::relay::ExprFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::RelayExpr const&)>::InitVTable()::'lambda4'(tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::RelayExpr const&)>*)::__invoke(tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::RelayExpr const&)>*) + 24
  [bt] (2) 3   libtvm.dylib                        0x0000000114451d98 tvm::relay::tec::LowerToTECompute::VisitExpr_(tvm::relay::CallNode const*) + 2232
  [bt] (1) 2   libtvm.dylib                        0x00000001146e5cfd tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<TVMFuncCreateFromCFunc::$_2> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) + 109
  [bt] (0) 1   libtvm.dylib                        0x0000000114700518 tvm::runtime::Backtrace() + 24
  File "/Users/wzh/tvm-dev/python/tvm/_ffi/_ctypes/packed_func.py", line 81, in cfun
    rv = local_pyfunc(*pyargs)
  File "/Users/wzh/tvm-dev/python/tvm/relay/backend/te_compiler.py", line 313, in lower_call
    best_impl, outputs = select_implementation(op, call.attrs, inputs, ret_type, target)
  File "/Users/wzh/tvm-dev/python/tvm/relay/backend/te_compiler.py", line 203, in select_implementation
    outs = impl.compute(attrs, inputs, out_type)
  File "/Users/wzh/tvm-dev/python/tvm/relay/op/op.py", line 126, in compute
    return _OpImplementationCompute(self, attrs, inputs, out_type)
  File "/Users/wzh/tvm-dev/python/tvm/_ffi/_ctypes/packed_func.py", line 237, in __call__
    raise get_last_ffi_error()
  [bt] (6) 7   ???                                 0x00007ff7be438320 0x0 + 140702025745184
  [bt] (5) 6   _ctypes.cpython-38-darwin.so        0x0000000101ce3fb7 ffi_call_unix64 + 79
  [bt] (4) 5   libtvm.dylib                        0x00000001146e343e TVMFuncCall + 62
  [bt] (3) 4   libtvm.dylib                        0x0000000114518ebe tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::relay::$_3> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) + 190
  [bt] (2) 3   libtvm.dylib                        0x000000011451754a tvm::relay::OpImplementation::Compute(tvm::Attrs const&, tvm::runtime::Array<tvm::te::Tensor, void> const&, tvm::Type const&) + 202
  [bt] (1) 2   libtvm.dylib                        0x00000001146e5cfd tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<TVMFuncCreateFromCFunc::$_2> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) + 109
  [bt] (0) 1   libtvm.dylib                        0x0000000114700518 tvm::runtime::Backtrace() + 24
  File "/Users/wzh/tvm-dev/python/tvm/_ffi/_ctypes/packed_func.py", line 81, in cfun
    rv = local_pyfunc(*pyargs)
  File "/Users/wzh/tvm-dev/python/tvm/relay/op/strategy/generic.py", line 153, in _compute_softmax
    return [topi_compute(inputs[0], axis)]
  File "<decorator-gen-77>", line 2, in softmax
  File "/Users/wzh/tvm-dev/python/tvm/te/tag.py", line 56, in tagged_fdecl
    return func(*args, **kwargs)
  File "/Users/wzh/tvm-dev/python/tvm/topi/nn/softmax.py", line 41, in softmax
    return softmax_common(x, axis, False)
  File "/Users/wzh/tvm-dev/python/tvm/topi/nn/softmax.py", line 73, in softmax_common
    k1 = te.reduce_axis((0, shape[axis]), name="k")
  File "/Users/wzh/tvm-dev/python/tvm/ir/container.py", line 36, in __getitem__
    return getitem_helper(self, _ffi_api.ArrayGetItem, len(self), idx)
  File "/Users/wzh/tvm-dev/python/tvm/runtime/container.py", line 57, in getitem_helper
    raise IndexError("Index out of range. size: {}, got index {}".format(length, idx))
IndexError: Index out of range. size: 1, got index 1

This cause of this error is the invalid axis=1 in the call to nn.softmax. However, this error message does not clearly tell that axis is invalid. I think the invalid axis should be reported immediately in Relay type inference, instead of elsewhere after it in the compilation process.