[Python] Array slicing inside `jax.lax.while_loop yields error` "array boolean indices must be...

Please consider the following toy example, which mimics what I'm trying to achieve in my real-world application:

import flax
import jax

def op(x):
return x - 1

@flax.struct.dataclass
class adaptive_state:
i: jax.numpy.ndarray
key: jax.random.PRNGKey

def loop(carry):
state, active_mask = carry

key, step_key = jax.random.split(state.key)
cond = jax.random.bernoulli(step_key, .5, state.i[active_mask].shape)

i_next = op(state.i[active_mask])

return state.replace(
i = state.i.at[active_mask].set(jax.numpy.where(cond, i_next, state.i[active_mask])),
key = key
), active_mask.at[active_mask].set(jax.numpy.where(cond, i_next > 0, active_mask[active_mask]))

init_state = adaptive_state(jax.numpy.array([47, 11, 815, 3], dtype = int), jax.random.PRNGKey(0))
state = jax.lax.while_loop(lambda carry: jax.numpy.any(carry[1]), loop, (init_state, jax.numpy.ones(init_state.i.shape[0], dtype = bool)))[0]


In my real-world application op is a computationally expensive task which is why I don't want to perform it on already "completely processed" (i.e. "inactive") elements. In they toy example the task is simply to reduce the elements by 1 until they are 0. Also: I do NOT want to use something like jax.vmap to perform op on individual active elements separately, since my real-world op is already optimized for batch processing.

Is there any way to achieve what I'm looking for? The current code yields the error IndexError("Array boolean indices must be concrete."). It works as expected when I replace jax.lax.while_loop by a dummy while_loop like

def while_loop(cond_fun, body_fun, init_val):
val = init_val
while cond_fun(val):
val = body_fun(val)
return val

Continue reading...