Bug report with Loop

[JamesZFS]

A case where I found a possible bug with drjit Loop feature:

import drjit as dr
from drjit.llvm.ad import UInt, Float, Loop

def f(x):  # returns x + 0 + 1 + ... + 9
    i = UInt(0)
    loop = Loop("Loop", lambda: (x, i))
    while loop(i < 10):
        x = x + i
        i += 1
    return x

a = dr.zeros(Float, 10)
dr.make_opaque(a)
a_old_index = a.index

print(f'{a.index=}, {a}')
b = f(a)
print(f'{b.index=}, {b}')
assert dr.allclose(b, [45] * 10)
assert a.index != a_old_index  # drjit.loop seems to change the index of the captured variable to that of a placeholder
print(f'{a.index=}, {a}')  # runtime error: a is the placeholder, which cannot be evaluated

Adding a copy in the first line of def f(x): solves this issue, but it just appears strange to me that after instantiating Loop, the original jit variable becomes a placeholder without any notice.

[Speierers]

This is a well understood issue of the current implementation of Dr.Jit symbolic loops. So far we have always assumed that the loop state variables wouldn't be referenced elsewhere (e.g. otherwise would be copied).

In the future we might change the API of dr.Loop to make this assumption for explicit.

In this specific example you expect a == b, in which case it would make more sense to do a = f(a). Otherwise adding x = Float(x) before the loop initialization will also ensure that you don't end up with a placeholder variable in a.

Let's keep it open as it is an important issue and should be fixed in the future somehow.

[JamesZFS]

I understand the issue and consideration, thanks!