Why jp.choose needs mode = 'wrap' or 'clip' and does not work with 'raise'?

[MiladInk]

I was trying out jp.choose in my code and when I wanted to jit it jax complains that I should use 'wrap' or clip and not 'raise'. I do not understand the reasoning behind that. Can anyone elaborate? I appreciate it a lot.

[mattjj]

Thanks for the question!

I think 'raise' could work in some situations but not others. In particular it wouldn't work when code isn't being staged out by jit, scan, or similar.

By "staging" I mean "rather than evaluating numerical values in arrays as Python code is executed (e.g. in jax.numpy function applications), instead delaying such evaluation so that it can be carried out separately from Python execution, e.g. after optimized compilation". (We should add a good definition of "staging" to the JAX Glossary.) jax.jit performs staging, as does jax.lax.scan and other higher-order primitives (i.e. ones which accept functions/callables as arguments). That means when using jit/scan/etc numerical values aren't available while executing Python, and in particular for jnp.choose the values of the indices aren't available for us to check and potentially raise an exception as jnp.choose is applied.

Technically distinct from staging is how vmap works. For vmap values are computed as Python executes, but they're abstracted and hidden, so to make functions like jnp.choose check their values we'd need to break some abstraction barriers. In other words, in addition to staging, using jax.vmap with jnp.choose would also effectively prevent us from handling the 'raise' mode.

So in other words, I would expectjnp.choose with mode='raise' to work fine when not using jax.jit/jax.lax.scan or jax.vmap (e.g. perhaps only using AD transformations like jax.grad), but to raise errors when those transformations are involved. Indeed the code looks consistent with that.

To make things like mode='raise' work with staging (or vmap), we'd need to be able to stage out some kind of runtime errors. That is, we need to delay the error check so that it happens in the compiled code (rather than while Python is executing). That means our compiled code needs to be able to express runtime errors, which to date it has not.

But we have some features related to staged-out runtime errors just coming online now, like jax.experimental.checkify and robust Python callbacks. It might be worth considering using those features in jax.numpy to handle situations like mode='raise'. cc @LenaMartens @sharadmv

Does that make sense?