Hi all,
I am trying to construct Relay expressions using the Python API. I am using the following code
‘’’ import tvm from tvm import relay from tvm.relay import var, Var, Function, Let, TensorType, Module from tvm.relay.prelude import Prelude
from network import Network
class GatherNDTest():
def __init__(self, name="f", **kwargs):
self.mod = Module()
self.p = Prelude(self.mod)
self.inputs = []
self.f = relay.GlobalVar(name)
self.mod[self.f] = relay.Function(self.inputs, self.build_impl(**kwargs), None)
def input(self, i):
self.inputs.append(i)
return i
def build_impl(self, vocab_size, embedding_size):
# Input declarations
embeddings_type = TensorType(shape = (vocab_size, embedding_size), dtype = "float32")
embeddings = self.input(Var("embeddings", embeddings_type))
word_index_type = TensorType(shape = (), dtype = "int32")
word_index = self.input(Var("word_index", word_index_type))
index_expr = tvm.relay.gather_nd(embeddings, word_index)
return index_expr
mod = GatherNDTest(name = "gather_test", vocab_size = 200, embedding_size = 400).mod
print(mod)
‘’’
The backtrace from GDB is here (https://pastebin.com/cvbtn1tm). I am using commit e3eff20d8a91c2b986548f780493a8e30b0c2b7a on master.
In any case, I cannot figure out how to pass arguments to gather_nd, so I can index the embeddings tensor to get a single word embedding. Thanks!