[Bug] `'IntImm' object is not iterable` in hybrid.parser.visit_Subscript

[Cnly]

Expected behavior

topi.scatter is implemented with @hybrid.script and its body has something that looks like

for i in range(data.shape[0]):

which should not be a problem when a module using scatter is built.

Actual behavior

However, when topi.scatter gets interpreted by te/hybrid/parser.py, there'll be an error 'IntImm' object is not iterable from this line

https://github.com/apache/tvm/blob/211291fb9a59b5ca6aee6e83b7cd700f98463032/python/tvm/te/hybrid/parser.py#L379

If I apply this patch, it seems to resolve the problem partially (I'm not sure how to patch the if isinstance(node.ctx, ast.Load) case below):

diff --git a/python/tvm/te/hybrid/parser.py b/python/tvm/te/hybrid/parser.py
index 442aeb6f1..255e930ae 100644
--- a/python/tvm/te/hybrid/parser.py
+++ b/python/tvm/te/hybrid/parser.py
@@ -376,6 +376,8 @@ class HybridParser(ast.NodeVisitor):
         args = self.visit(node.slice)
         arr = self.visit(node.value)
         if isinstance(arr, Array):
+            if isinstance(args, _expr.IntImm):
+                args = [args]
             for i in args:
                 if isinstance(i, numbers.Integral):
                     arr = arr[i]

Environment

OS: Ubuntu 18.04.6 LTS TVM: Commit 211291fb9

Steps to reproduce

import numpy as np
import tvm
import tvm.relay as relay

idxes = np.array([[0, 1],
                  [2, 3],
                  [4, 5]])

x = relay.zeros((3, 10), 'int32')
y = relay.var("y", shape=idxes.shape, dtype='int32')
z = relay.scatter(x, y, relay.full_like(y, relay.const(1)), -1)

f = relay.Function([y], z)
m = tvm.IRModule.from_expr(f)

# Either line triggers the error (EDITED):
relay.build(m, target='llvm')
relay.create_executor('graph', mod=m).evaluate()(idxes)

[masahi]

Doesn't this mean that the rhs of the slice assignment arr is an array but the lhs is a scalar? I'm assuming node.slice is the lhs of assignment, but it may be the index of subscripting.

[Cnly]

@masahi Sorry I don't think I follow you... but I tried adding print statements in scatter() and visit_Subscript() to see the arguments and here's what I get:

scatter():
data=Tensor(shape=[3, 10], op.name=placeholder)
indices=Tensor(shape=[3, 2], op.name=placeholder)

visit_Subscript():
node=<ast.Subscript object at 0x7f17fea956a0>
node.slice=<ast.Constant object at 0x7f17fea955e0>
args=0
node.value=<ast.Attribute object at 0x7f17fea95670>
arr=[3, 10]

Do you think this makes things clearer?

[masahi]

Yeah sorry I was confused with something. I don't know what args = 0 corresponds to in your scatter op workload.

[Cnly]

I think it's from this line

https://github.com/apache/tvm/blob/31de5bc40f22ba761361464921b9b73af51bb214/python/tvm/topi/scatter.py#L36

If I changed the 0 to 99 then it prints args=99 before the exception.

[masahi]

I see, so does this mean the parser expects subscript index (named args) to be an array? It's a bit surprising that we didn't hit this error before, since I think scalar indexing like shape[0] is used all over the places.

[Cnly]

I guess this only affects operators implemented with @hybrid.script? A quick grep shows not much usage in tvm. To see if I'm using the scatter operator wrong and to confirm my guess, I tried with relay.unique and saw the same exception. Reproduction snippet:

import numpy as np
import tvm
import tvm.relay as relay

idxes = np.array([[0, 1],
                  [2, 3],
                  [4, 5]])

x = relay.zeros((3, 10), 'int32')
y = relay.var("y", shape=idxes.shape, dtype='int32')
z = relay.unique(relay.sum(x, axis=0))[0]

f = relay.Function([], z)
m = tvm.IRModule.from_expr(f)

# Either line triggers the error:
relay.build(m, target='llvm')
relay.create_executor('graph', mod=m).evaluate()()

The line pointed to by traceback:

https://github.com/apache/tvm/blob/31de5bc40f22ba761361464921b9b73af51bb214/python/tvm/topi/unique.py#L273

[masahi]

hmm interesting... I believe unique op should be well tested so I don't know why it fails for that example. Can you send a PR with your proposed solution, along with test cases for scatter and unique?

[Hzfengsy]

What's your Python version? Hybrid Script is not supported by Python3.9+ (only from my experience)

[Cnly]

Oops, it's 3.9.9. Switching back to 3.6.9 did seem to fix the problem. Why is that?

[Hzfengsy]

Hybrid Script is based on Python AST which may change between different Python versions.

[Cnly]

I see, thanks! Any suggestion on what to do about this issue? Should we close it? Or rename it?

[masahi]

There is already an issue for 3.9 https://github.com/apache/tvm/issues/8577, we can close this.