Is it possible to index a python list with a jax array?

[pbsmoothie4life]

I'm wondering if there's a way to index a python list with a jax array? I know that one can index a list with a python int, by passing that argument as a static_argnum to jit, but what about the case when the index comes from a jit'ed function? A minimally working example showing the issue is the following....

from jax import jit
from functools import partial

@jit
def ind(n):
    return n

@partial(jit, static_argnums=(1,))
def fun(x,n):
    return x[n]

n = ind(3)

x = list(range(10))

fun(x, n)

I understand that in this specific example I can turn x into a jnp.array and the example will work, but I'm interested in the case where x is a general Pytree, and so I need x to remain a list or some other nestable object. I also understand that if I don't jit the ind function, this example will work, but in my case the ind function is computationally expensive so I'd like to jit it. I've noticed that @jakevdp answers a lot of questions of these types. I feel like I might be able to use jax.lax.switch to do this, but I also feel like there may be a more direct solution that I'm missing. Thanks in advance for any help!

[mattjj]

Thanks for the question!

How about just calling int on the array?

from jax import jit
from functools import partial

@jit
def ind(n):
    return n

@partial(jit, static_argnums=(1,))
def fun(x,n):
    return x[n]

n = ind(3)

x = list(range(10))

fun(x, int(n))  # CHANGE

[pbsmoothie4life]

Thank you for the quick reply, I really appreciate it! This solution works in the toy example I brought up, but when I try to use it in my actual setting it breaks with a concretization error on the int function (see below). I'm going to continue along these lines though. Thanks again!

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[])>with<BatchTrace(level=3/1)> with
val = Traced<ShapedArray(int32[64])>with<DynamicJaxprTrace(level=2/1)>
batch_dim = 0
The problem arose with the int function. If trying to convert the data type of a value, try using x.astype(int) or jnp.array(x, int) instead.

[mattjj]

We should improve that error message. But I think it's saying that we're out-of-luck because vmap batching is involved. In general that means every element in the batch might select a different list element, but there's no way for us to represent that in a vectorized way.

[pbsmoothie4life]

Ahh that's what I was worried about, the 64 is indeed the batch size. I saw a post on here where someone brought up an issue with vmapping over lists in a similar way, and was thinking that could be the issue here.

Okay, I might need to go the jax.lax.switch route after all, and just allocate the function arguments to be the maximum sized tree to cover the space of all possible trees.

Thanks again for your very thoughtful help, I really appreciate it!

[jakevdp]

One way you could index into a list with a traced index is to use lax.switch; for example:

from jax import jit, lax

@jit
def fun(x,n):
  # return x[n]  # fails with TracerIntegerConversionError
  return lax.switch(n, [lambda xi=xi: xi for xi in x])

x = list(range(10))
fun(x, 3)

The mental model of what's going on here is that JAX creates ten different functions with constant outputs, pushes them all to the compiler, and then at runtime the compiler chooses which function to run based on the index value. This works, but will be pretty inefficient as the size of the list grows. In general, it's much more efficient to convert x to an array and index it directly:

@jit
def fun2(x, n):
  return x[n]
fun2(jnp.array(x), 3)

[mattjj]

Hah I hadn't thought of that! Neat trick!

It does require that the list elements are all the same shape, whereas the static_argnums version re-specializes and hence can handle distinct shapes/dtypes of list elements, at the cost of re-tracing and re-compiling when the index changes (whereas this lax.switch version only needs to trace/compile once).

[pbsmoothie4life]

Thanks so much for the quick response! I tried creating that exact list of functions as well, but it has an odd behavior where all of the functions map to 9.

for f in [lambda: xi for xi in list(range(10))]:
    print(f())

yields 9, 10 times; idk if it will work differently when embedded in lax.switch though. @mattjj brings up a good point I hadn't thought of, which is going about this jax.lax.switch route will require all of the function branches to take arguments of the same tree-structures, which is a deal breaker in my setting so I'm going to go the static_argnums route.

[jakevdp]

Oh, that's my mistake, sorry. It's a Python variable scope issue. I edited the answer to fix it.

[mattjj]

I think this is the classic Python-loops-dont-create-new-binders-they-reuse-them bug. This should work:

for f in [(lambda xi: lambda: xi)(xi) for xi in list(range(10))]:
    print(f())

You'd need that with a lax.switch also.