I would like ask for some advice. I want to generate TIR functions via the Python API (TIR ir_builder) (called from Relax), which take a list of tensors as input, (which might have different shapes), and does some computation on them. For example concatenation, where the input tensors’ shape might differ, or stack operation, or summation of a list of tensors (topi.elemwise_sum).
I want to express operations like this, but currently unsure how to do it. Currently I am trying to generate a Python list with the Tensors, what I can use in a Python loop to unroll the operation along the list items. But I was wondering if it would be possible to do this without unrollling. Is there a data structure I could use instead of the Python list, which would support the ir_builder’s LoopVars, as index? Because currently, if I use a python list to unroll the tensors, it can only be indexed with Python variables.