Since when is this form of dynamic indexing possible under jit?

[simon-bachhuber]

Out of curiosity: Has this always been possible? I thought this is where jax.lax.dynamic_* was required.

@jax.jit
def f(idxs, data):
    res = 0
    for i in idxs:
        res += data[i]
    return res

f(jnp.array([0, 1, 0]), jnp.array([1, 3]))

[soraros]

I think yes, as the shape of res is static. But in fact if call jax.make_jaxpr on the function, I won't be surprised if the __getitem__ is indeed lowered into dynamic_slice.

[jakevdp]

I believe this kind of dynamic indexing (x[i] where x and i are traced, and i is integer-typed) has always been supported in JIT, but dynamic iteration (for i in idxs where idxs is traced) has only been possible since #8043.

[simon-bachhuber]

I feel like you can create some surprising behaviour imo.
E.g.

@jax.jit
def f(parent_array):
    # now it works..
    ys = jnp.array([1, 2, 3])
    # now it doesn't
    ys = [1, 2, 3]

    res = 0
    for i in range(len(parent_array)):
        parent = parent_array[i]
        res += ys[parent]
    return res

f(jnp.array([0, 1, 2]),)

I somehow expected this to either not work or work both times.

[soraros]

Python lists need the index to be a int, and we can't just convert traced value into int.
Maybe you could say more about why you find this behaviour surprising?