How to get the shape of output tensor returned from tvm.relay.expr.call?

Troubleshooting
#1

in tensorflow, we can simply do

output = a + b
dims = len(output.shape.as_list())

in tvm/relay, how can we get the shape of the returned output from a callnode?

Thanks

#2

At the TVM DSL level, you can directly say output.shape to access the shape of any Tensor.

tvm.compute is converted to ComputeOpNode at the AST level.

You can access the shape info by accessing the ComputeOpNode's ``axis` attribute.

However, i think ComputeOpNode itself is not directly accessible. In my understanding the way you can access the ComputeOpNode is through the Provide or Store or perhaps Call node for a particular tensor, by accessing the FuncRef attribute of the aforementioned nodes.

Adding to it, the information in the axis attribute may not be the most updated after the InferBound pass.

#3

Yes. at tvm level, we can get the shape easily, like

n = tvm.var(“n”)
A = tvm.placeholder((n,), name=‘A’)
B = tvm.compute(A.shape, lambda i: A[i] + 1.0, name=“B”)

print(A.shape)
print(B.shape)

however, at relay level, there is no way to get it.

x = relay.var(“x”, shape=(3, 2))
y = relay.var(“y”, shape=(2, 3))
z = relay.add(x, y)

print(x.shape) <—

File “/home/yinma/baidu/tvm/xmir/python/tvm/_ffi/_ctypes/node.py”, line 75, in getattr
“’%s’ object has no attribute ‘%s’” % (str(type(self)), name))

AttributeError: ‘<class ‘tvm.relay.expr.Call’>’ object has no attribute ‘shape’

I am wondering if there is a way to get the shape information at relay level?

#4

I use “get_const_tuple(x.type_annotation.shape)” to get the shape.

1 Like
#5 
from tvm import relay
x = relay.var('x', shape=(10, 10))
y = relay.var('y', shape=(10, 10))
z = relay.add(x, y)
f = relay.Function([x, y], z)
m = relay.Module()
m['main'] = f
m = relay.transform.InferType()(m)
print(m['main'].ret_type)

p.s. Your example won’t work for the above code due to shape mismatch. You’ll see an error when assigning the function to the module.

1 Like