Does TVM support dynamic range sum?

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.

I ran into this yesterday. Can someone more knowledgeable about TVM tell us if this is a limitation of the compute and/or reduce_axis API or of the actual code generator (or both)?

A possible workaround I was looking into is using the lower level ir_builder directly as is done in the SPMM kernel which seems to have dynamic loop bounds here.

From some offline conversations, it sounds like these use cases are encouraged to use the (hybrid script capability](https://docs.tvm.ai/langref/hybrid_script.html). I’ll report back once I try it.

Hi, have you solved it ? I would like to use some dynamic computed loop counts inside te.compute and don’t know how to do it, I guess it is the same issue as yours. Looking forward to your feedback, thanks.

Hi @huangteng, unfortunately dynamic indices are not supported in the te.compute API today, but can be supported with both the IRBuilder interface and the hybrid script and tvm script (which is designed to replace hybrid script) interfaces.

Hi jknight, thanks for the reply, not sure if I understood correctly, but I think dynamic indices works since the tvm.build pass. However I have a feeling that this affects the autoTVM task and will throw some strange backtrace, I have mentioned it and written the simplest reproduce code in this post, could you please help to take a look ? Appreciated (I stucked into that problem for nearly a month frowning: ).