How to index ArrayXf using index of type UInt32

[XiaoXinyyx]

Is there any way to implement the code below?
(In symbolic loop, the idx is a state variable, so I keep it as UInt32)

from drjit.cuda import ArrayXf, UInt32, Float

A = ArrayXf([1,2],[2,3],[3,4],[4,5])
idx= UInt32(1, 3)
value = Float(99, 100)
A[idx] = value # Exception occured: Complex slicing operations are currently only supported on tensors.

[rtabbara]

Hi @XiaoXinyyx ,

I think the most concise way is to use TensorXf instead

import drjit as dr
from drjit.cuda import TensorXf, UInt32, Float, ArrayXf

A = TensorXf([[1,2],[2,3],[3,4],[4,5]])
idx= UInt32(1, 3)
value = Float([99, 100] * dr.width(idx))
A[idx, :] = value

print(A)
'''
[[1, 2],
 [99, 100],
 [3, 4],
 [99, 100]]
'''

But I would recommend reading the Dr.Jit section on tensor limitations here and why you may encounter issues if you're inside a symbolic loop.

If you still wanted to stick to an ArrayXf, then I would suggest instead maybe changing the layout of your data. Something like

B = ArrayXf([1,2,3,4],[2,3,4,5])
idx= UInt32(1, 3)
value = ArrayXf([99], [100])
dr.scatter(B, value, idx)
print(B)
'''
[[1, 2],
 [99, 100],
 [3, 4],
 [99, 100]]
'''

[XiaoXinyyx]

Thank you for your answer. I've tried using TensorXf or dr.scatter, but I want to get a result like
[1, 99, 3, 4], [2, 3, 4, 100] or the transpose of that. But I think there might not be a simple way to get this result in one step.

[rtabbara]

Oh apologies. Yes, you would probably have to split this up into multiple scatters. Using example B from before

B = ArrayXf([1,2,3,4],[2,3,4,5])
dr.scatter(B[0], Float(99), UInt32(1))
dr.scatter(B[1], Float(100), UInt32(3))

To be honest, when you perform a scatter on an ArrayXf it internally will perform

dr.scatter(B[0], val[0], idx)
dr.scatter(B[1], val[1], idx)
...

anyway, it's just that for your particular case, you want to write to different indices depending on the components for your ArrayXf (i.e. B[0], B[1], etc.)

[XiaoXinyyx]

I reviewed the document of dr.jit again and I found drjit.alloc_local maybe a proper solution? Because an array in local memory can be accessed using computed indices

[rtabbara]

drjit.alloc_local maybe a proper solution?

I'm not sure. The documentation on local memory is quite detailed and includes an example usage, so that should give a good idea whether or not it's suitable for your needs.

I'm sorry, I'm not sure I can provide any further advice beyond my previous suggestions. As a general comment, if you want to concisely express more complex tensor slicing operations, as the documentation mentions, Dr.Jit may be ill-suited for your needs and you might be better off considering tensor libraries such as PyTorch.