scatter_nd
might work if we assume that input_put_
usually sets the values to empty tensor.
ctx = tvm.cpu(0)
target = 'llvm'
dtype = 'float32'
shape = (3,3)
tp = relay.TensorType(shape)
data = np.random.rand(*shape).astype(dtype)
hs = [0, 1, 2, 2]
ws = [0, 1, 1, 2]
vs = [2.0, 4.0, 7.0, 9.0]
indices = tvm.nd.array([hs, ws], ctx)
values = tvm.nd.array(np.array(vs).astype(dtype), ctx)
indices_re = relay.Constant(indices)
values_re = relay.Constant(values)
y = relay.scatter_nd(values_re, indices_re, shape)
x = relay.var("x", tp, dtype=dtype)
func = relay.Function([x], y)
intrp = relay.create_executor("graph", ctx=ctx, target=target)
intrp.evaluate(func)(data)
<tvm.nd.NDArray shape=(3, 3), cpu(0)>
array([[2., 0., 0.],
[0., 4., 0.],
[0., 7., 9.]], dtype=float32)