linen.vmap gives error "TypeError: list indices must be integers or slices, not BatchTracer"

[marcvanzee]

The input is a tensor of indexes.

[marcvanzee]

Answer by @levskaya:

You're trying to index into a python list with a tracer. You should make sure your vmapped input is in a jnp.array(...):

def foo1(x):
  return 1*x

def foo2(x):
  return 2*x

@jax.vmap
def bar(x, idx):
  y = jnp.array([foo1(x), foo2(x)])
  return y[idx]

x = jnp.array([0.,1.,2.])
idx = jnp.array([0,1,0])
bar(x, idx)