Guidance on conditionally applying function across batch

[bycn]

Minimal example:

We have a pytree attribute state.done with shape (batch,) of type boolean as well as other attributes. The desired behavior is, if example i is true, then apply fn1 on the ith examples of state's attributes, otherwise fn2. Conceptually, it is just jax.lax.cond applied across a batch (like jnp.where).

Actual context/goal:

I'm writing a parallelizable reinforcement learning environment. A common pattern is to "reset" the state after a number of (possibly variable) steps (each step is a function call carrying some state.) Ideally, this reset creates a new rng that changes the state for the next episode (set of steps until next reset). In the example described above, I'm trying to create this behavior by taking the result of a vmapped step function, checking if any of their outputted state is "done", and calling a reset function on the ones that do. In brax, this is sidestepped as they do not modify the rng after a reset, instead simply reverting the state back to the state at the initial step. This is easy to implement and is working efficiently for my environment as well, but it is undesirable from a learning perspective.

Would appreciate any pointers and happy to clarify anything :)

[jakevdp]

This sounds like the precise use-case of lax.switch.

[bycn]

Sorry, I'm not sure how it would apply? lax.switch selects from many functions and applies one to one input. I'm trying to do lax.cond mapped across a batch, but not sure if there is a better way

[soraros]

I think @jakevdp meant something like

def apply_fn(b, x):
  return lax.switch(b, [fn2, fn1], x)

vmap(apply_fn)(state.done, state.example)

[bycn]

Ah, this would be similar to directly mapping cond, I guess it is the recommended approach

[soraros]

Indeed, they are both compiled to select, and generate largely identical code.