How to implement aten::item in relay.from_pytorch

1

I have question that is it possible for relay op to return a scalar? I have tried some methods and both of them failed… My pytorch code is as follows:

import torch

@torch.jit.script
def max_item(x, y):
if x.item() > y.item():
r = x
else:
r = y
return r

def bar(x, y, z):
return max_item(x, y) + z

a = torch.randn([1])
b = torch.randn([1])
c = torch.randn([1])
print(“a.item():”, a.item())
print(“b.item():”, b.item())
print(“c.item():”, c.item())
traced_graph = torch.jit.trace(bar, (a, b, c))

print(traced_graph.graph_for(a, b, c))

JIT Graph:
graph(%x : Tensor,
%y.1 : Tensor,
%z : Tensor):
%3 : int = prim::Constantvalue=1 # torch_script.py:13:0
%4 : Scalar = aten::item(%x) # torch_script.py:5:7
%5 : Scalar = aten::item(%y.1) # torch_script.py:5:18
%6 : bool = aten::gt(%4, %5) # torch_script.py:5:7
%y : Tensor = prim::If(%6) # torch_script.py:5:4
block0():
→ (%x)
block1():
→ (%y.1)
%12 : Tensor = aten::add(%y, %z, %3) # torch_script.py:13:0
return (%12)

2

I think it should be possible, how about op.take(scalar_tensor, 0)?

3

Thanks for quick response.
If I used your suggested implementation:
The error will be:

File “/xtcl-master/baidu/xpu/xmir/include/tvm/runtime/packed_func.h”, line 687 TVMError: In function relay.op._make.take: error while converting argument 1: [02:09:14] /xtcl-master/baidu/xpu/xmir/include/tvm/runtime/packed_func.h:1584:

An internal invariant was violated during the execution of TVM. Please read TVM’s error reporting guidelines. More details can be found here: https://discuss.tvm.ai/t/error-reporting/7793.

Check failed: type_code_ == kTVMObjectHandle (0 vs. 8) : expected Object but got int

If I use implementation same as tensortonum:

def tensortonum(self, inputs, input_types):  
    return inputs[0]

It will report error like following:

tensor type bool has 0 dimensions, while Tensor[(1), bool] has 1 dimensions
The Relay type checker is unable to show the following types match.
In particular Tensor[(1), bool] does not match bool
note: run with TVM_BACKTRACE=1 environment variable to display a backtrace.