Repository: apache/tvm

Issue:
[Bug] Building relay ops defined using hybrid_script gives ValueError: All indices are supposed to be constants
### Expected behavior

Be able to build a relay containing relay.vision.non_max_suppression or relay.vision.multibox_transform_loc.

### Actual behavior

Obtained error:

```
Traceback (most recent call last):
  File "test.py", line 45, in <module>
    mod = relay.build(mod, executor=EXECUTOR, target=TARGET,runtime=RUNTIME)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/relay/build_module.py", line 364, in build
    graph_json, runtime_mod, params = bld_mod.build(
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/relay/build_module.py", line 161, in build
    self._build(
  File "tvm/_ffi/_cython/./packed_func.pxi", line 331, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 276, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./base.pxi", line 181, in tvm._ffi._cy3.core.CHECK_CALL
ValueError: Traceback (most recent call last):
  41: TVMFuncCall
  40: tvm::relay::backend::RelayBuildModule::GetFunction(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, tvm::runtime::ObjectPtr<tvm::runtime::Object> const&)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#3}::operator()(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) const
  39: tvm::relay::backend::RelayBuildModule::BuildRelay(tvm::IRModule, tvm::runtime::String const&)
  38: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::relay::backend::AOTExecutorCodegenModule::GetFunction(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, tvm::runtime::ObjectPtr<tvm::runtime::Object> const&)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#2}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  37: tvm::relay::backend::AOTExecutorCodegen::Codegen(tvm::IRModule, tvm::relay::Function, tvm::runtime::String)
  36: tvm::transform::Pass::operator()(tvm::IRModule) const
  35: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  34: tvm::transform::SequentialNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  33: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  32: tvm::transform::ModulePassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  31: _ZN3tvm7runtime13PackedFuncObj9ExtractorINS0_16PackedFuncSubObjIZNS0_15TypedPackedFuncIFNS_8IRModuleES5_NS_9transform11PassContextEEE17AssignTypedLambdaIZNS_5relay3tec7LowerTEENS0_6StringENS_17CompilationConfigESt8functionIFvNS_8BaseFuncEEEEUlS5_S7_E_EEvT_EUlRKNS0_7TVMArgsEPNS0_11TVMRetValueEE_EEE4CallEPKS1_SL_SP_
  30: tvm::relay::tec::LowerTE(tvm::IRModule const&, tvm::runtime::String const&, std::function<void (tvm::BaseFunc)>, tvm::CompilationConfig)
  29: tvm::transform::Pass::operator()(tvm::IRModule) const
  28: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  27: tvm::relay::transform::FunctionPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  26: _ZN3tvm7runtime13PackedFuncObj9ExtractorINS0_16PackedFuncSubObjIZNS0_15TypedPackedFuncIFNS_5relay8FunctionES6_NS_8IRModuleENS_9transform11PassContextEEE17AssignTypedLambdaIZNS5_3tec15LowerTensorExprENSD_10TECompilerESt8functionIFvNS_8BaseFuncEEENS_17CompilationConfigEEUlS6_S7_S9_E_EEvT_EUlRKNS0_7TVMArgsEPNS0_11TVMRetValueEE_EEE4CallEPKS1_SM_SQ_
  25: tvm::relay::ExprMutator::VisitExpr(tvm::RelayExpr const&)
  24: _ZZN3tvm5relay11ExprFunctorIFNS_9RelayExprERKS2_EE10InitVTableEvENUlRKNS_7r
  23: tvm::relay::transform::DeviceAwareExprMutator::VisitExpr_(tvm::relay::FunctionNode const*)
  22: tvm::relay::tec::LowerTensorExprMutator::DeviceAwareVisitExpr_(tvm::relay::FunctionNode const*)
  21: _ZN3tvm5relay9transform22DeviceAwareExprMutator21DeviceAwareVisit
  20: tvm::relay::ExprMutator::VisitExpr_(tvm::relay::FunctionNode const*)
  19: tvm::relay::ExprMutator::VisitExpr(tvm::RelayExpr const&)
  18: _ZZN3tvm5relay11ExprFunctorIFNS_9RelayExprERKS2_EE10InitVTableEvENUlRKNS_7r
  17: tvm::relay::transform::DeviceAwareExprMutator::VisitExpr_(tvm::relay::LetNode const*)
  16: tvm::relay::tec::LowerTensorExprMutator::PreVisitLetBinding_(tvm::relay::Var const&, tvm::RelayExpr const&)
  15: tvm::relay::ExprMutator::VisitExpr(tvm::RelayExpr const&)
  14: _ZZN3tvm5relay11ExprFunctorIFNS_9RelayExprERKS2_EE10InitVTableEvENUlRKNS_7r
  13: tvm::relay::transform::DeviceAwareExprMutator::VisitExpr_(tvm::relay::CallNode const*)
  12: tvm::relay::ExprMutator::VisitExpr(tvm::RelayExpr const&)
  11: _ZZN3tvm5relay11ExprFunctorIFNS_9RelayExprERKS2_EE10InitVTableEvENUlRKNS_7r
  10: tvm::relay::transform::DeviceAwareExprMutator::VisitExpr_(tvm::relay::CallNode const*)
  9: tvm::relay::tec::LowerTensorExprMutator::DeviceAwareVisitExpr_(tvm::relay::CallNode const*)
  8: tvm::relay::tec::TECompilerImpl::Lower(tvm::relay::tec::CCacheKey const&)
  7: tvm::relay::tec::TECompilerImpl::LowerInternal(tvm::relay::tec::CCacheKey const&, tvm::GlobalVarSupply)
  6: tvm::relay::tec::PrimFuncFor(tvm::relay::Function const&, tvm::Target const&, tvm::GlobalVarSupply)
  5: tvm::relay::tec::ScheduleBuilder::Create(tvm::relay::Function const&, tvm::GlobalVarSupply)
  4: tvm::relay::tec::LowerToTECompute::Lower(tvm::relay::Function const&)
  3: tvm::relay::backend::MemoizedExprTranslator<tvm::runtime::Array<tvm::te::Tensor, void> >::VisitExpr(tvm::RelayExpr const&)
  2: _ZZN3tvm5relay11ExprFunctorIFNS_7runtime5ArrayINS_2te6TensorEvEERKNS_9Relay
  1: tvm::relay::tec::LowerToTECompute::VisitExpr_(tvm::relay::CallNode const*)
  0: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<TVMFuncCreateFromCFunc::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#2}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) [clone .cold]
  File "tvm/_ffi/_cython/./packed_func.pxi", line 56, in tvm._ffi._cy3.core.tvm_callback
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/relay/backend/te_compiler.py", line 317, in lower_call
    best_impl, outputs = select_implementation(op, call.attrs, inputs, ret_type, target)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/relay/backend/te_compiler.py", line 207, in select_implementation
    outs = impl.compute(attrs, inputs, out_type)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/relay/op/op.py", line 126, in compute
    return _OpImplementationCompute(self, attrs, inputs, out_type)
  File "tvm/_ffi/_cython/./packed_func.pxi", line 331, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 276, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./base.pxi", line 181, in tvm._ffi._cy3.core.CHECK_CALL
  3: TVMFuncCall
  2: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::relay::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#4}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  1: tvm::relay::OpImplementation::Compute(tvm::Attrs const&, tvm::runtime::Array<tvm::te::Tensor, void> const&, tvm::Type const&)
  0: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<TVMFuncCreateFromCFunc::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#2}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) [clone .cold]
  File "tvm/_ffi/_cython/./packed_func.pxi", line 56, in tvm._ffi._cy3.core.tvm_callback
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/relay/op/strategy/generic.py", line 1262, in _compute_nms
    return topi_compute(
  File "/local_disk/local_sw/tvm/tvm/lib/python3.8/site-packages/decorator.py", line 232, in fun
    return caller(func, *(extras + args), **kw)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/target/generic_func.py", line 287, in dispatch_func
    return func(*args, **kwargs)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/topi/vision/nms.py", line 577, in non_max_suppression
    out, box_indices = hybrid_nms(
  File "/local_disk/local_sw/tvm/tvm/lib/python3.8/site-packages/decorator.py", line 232, in fun
    return caller(func, *(extras + args), **kw)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/__init__.py", line 60, in wrapped_func
    return source_to_op(src, args, func.__globals__, closure_vars)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/parser.py", line 645, in source_to_op
    parser = parse_python(src, args, symbols, closure_vars)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/parser.py", line 615, in parse_python
    parser.parsed_body = parser.visit(root)
  File "/usr/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/parser.py", line 229, in visit_Module
    return self.visit(node.body[0])
  File "/usr/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/parser.py", line 242, in visit_FunctionDef
    res = visit_list_to_block(self.visit, node.body)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/parser.py", line 58, in visit_list_to_block
    lst = [visit(stmt) for stmt in lst if not utils.is_docstring(stmt)]
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/parser.py", line 58, in <listcomp>
    lst = [visit(stmt) for stmt in lst if not utils.is_docstring(stmt)]
  File "/usr/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/parser.py", line 303, in visit_Assign
    rhs = self.visit(node.value)
  File "/usr/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/parser.py", line 388, in visit_Subscript
    _internal_assert(
  File "/local_disk/local_sw/tvm_main/tvm/python/tvm/te/hybrid/utils.py", line 43, in _internal_assert
    raise ValueError(err)
ValueError: All indices are supposed to be constants
```

### Environment

Branch main, commit 9b17f344a31a13226458a2d48dcb7b55ce282274

### Steps to reproduce

Code to reproduce:

```python
from tvm.relay.op.vision import non_max_suppression
from tvm import te
import numpy as np
from tvm import topi
import tvm
from tvm import relay

nms = False

if not nms:
    x0 = relay.var("x0", relay.ty.TensorType((1, 6, 5), "float32"))
    x1 = relay.var("x1", relay.ty.TensorType((1, 5 * 4), "int32"))
    x2 = relay.var("x1", relay.ty.TensorType((1, 5, 4), "int32"))
    z = relay.vision.multibox_transform_loc(
        x0, x1, x2, clip=False, threshold=-float("inf"), variances=(1.0, 1.0, 1.0, 1.0)
    )
    z = z.astuple()
    func = relay.Function([x0, x1, x2], z)
else:
    x0 = relay.var("x0", relay.ty.TensorType((1, 100, 6), "float32"))
    x1 = relay.var("x1", relay.ty.TensorType((1,), "int32"))
    x2 = relay.var("x2", relay.ty.TensorType((1, 100), "int32"))
    x3 = relay.var("x3", relay.ty.TensorType((), "int32"))
    z = relay.vision.non_max_suppression(
        x0,
        x1,
        x2,
        x3,
        iou_threshold=0.5,
        force_suppress=True,
        top_k=4,
        return_indices=True,
        invalid_to_bottom=False,
    )
    z = z.astuple()
    func = relay.Function([x0, x1, x2, x3], z)
    
mod = tvm.IRModule()
mod["main"] = func

print(mod["main"])
TARGET = tvm.target.target.Target({"kind": "c"})
mod = relay.build(mod, target=TARGET)

print(mod)
```

I started to try to understand what the problem is, but my experience with hybrid_script is very low. What I noticed is that function visit_Index in [parser.py](https://github.com/apache/tvm/blob/main/python/tvm/te/hybrid/parser.py) is returning a list, which then breaks function visit_Subscript (same file). But I was not able to understand the reason why visit_Index is now returning a list. A change was probably done before, that is making visit_Index return a list.



Discussion:

Resolving patch:
diff --git a/python/tvm/te/hybrid/parser.py b/python/tvm/te/hybrid/parser.py
--- a/python/tvm/te/hybrid/parser.py
+++ b/python/tvm/te/hybrid/parser.py
@@ -374,7 +374,7 @@ def visit_Attribute(self, node):
 
     def visit_Subscript(self, node):
         args = self.visit(node.slice)
-        if sys.version_info > (3, 8):
+        if sys.version_info >= (3, 9):
             if not isinstance(node.slice, ast.Tuple):
                 args = [args]
 


Test patch: