What I wanted to do is to multiply with a varying length vector, something like this:
def _compute(i, j):
k = tvm.reduce_axis((0, indptr[i+1]-indptr[i]), name='k')
return tvm.sum(data[indptr[i]+k]*weight[j,indices[indptr[i]+k]], axis=k)
matmul = tvm.compute(oshape, _compute, tag="dense", name='out')
This would obviously produce an TVMError: Not all Vars are passed in api_args: 'i' does not appeared in api_args
.
I think the problem is with the definition of k
within _compute
function, but how can I perform sum in a dynamic range?
Thanks in advance.