Support overriding implementations of JAX functions within a scope

[shoyer]

The idea here is to support overriding the implementation of higher order JAX
functions within a limited scope, for libraries such as Haiku and Flax that
implement their own versions of this functions that support mutation.

Usage example:

with jax.override_context({'lax.scan': my_scan}):
  # all calls to lax.scan() are replaced by my_scan()
  ...

Ultimately, it would be great to replace this with a unified interface for
mutable state in JAX, but this could be a convenient temporary measure.

[juliuskunze]

@shoyer This feature could potentially allow a large speedup of jnp functions running on the NumPy backend: numpy_eval()(jnp.<some_fun>) could be mapped directly to np.<some_fun>, instead of tracing jnp.<some_fun> and calling the NumPy implementations of the encountered lax primitives. It might make #3893 fast out-of-the-box, removing the need to keep optimized NumPy implementations of shape rules.

[shoyer]

@juliuskunze I would be a little reluctant to support directly overriding jax.numpy functions. JAX works so well in part because transformations are written at the LAX primitive level, which is a much more uniform and rationalized interface than jax.numpy. That's the main reason why I suggested this interface as a temporary measure -- it would be much better (if possible) to figure out a uniform interface for threading variables holding mutable state through transformations.

[NeilGirdhar]

Whoa, this is awesome!! This would resolve my problem in Haiku (linked above) and it seems very versatile too. What's the future of this pull request?

[jekbradbury]

I think I support this. We haven't made any progress on a unified state mechanism since the last time this was proposed, and it's better than external libraries directly monkeypatching JAX. Note that most libraries that provide state-aware versions of functions like jit, scan, etc. add additional knobs/functionality to those functions (and not necessarily in a mutually compatible way), so this feature doesn't solve all aspects of the problem.

[shoyer]

Since drafting this, we actually have use-case for this in our own application code where haiku.scan wasn't quite right and we wanted our own version.

My one thought is that we should probably make it clearer that it's experimental. If we don't think it's worth making a new experimental module, perhaps we should just include "experimental" in the name, e.g., jax.experimental_override_context? Or perhaps it can just live in the top-level jax.experimental namespace along with the enable_x64/disable_x64 context managers.

[shoyer]

This is ready for review now.

[NeilGirdhar]

Wow, I appreciate how fast the Jax team responds to issues!

I'm just curious, would it be worth adding the eval_shape override to the list? Haiku seems to provide an override.

[froystig]

Placing override_context in experimental sounds good to me.

[shoyer]

I'm just curious, would it be worth adding the eval_shape override to the list? Haiku seems to provide an override.

Yes, absolutely!

I think I got them all, but let me know if you notice any others that are missing. (It's also easy to extend this in the future, of course.)

[shoyer]

Tests are passing -- is anyone from the JAX team up for a review?

[shoyer]

After discussion with the whole JAX team, we decided not to go further with this PR at this time. It solves some real problems (writing library code that uses higher order functions that wants to be compatible with stateful transformations), but:

1. Clearly isn't the ultimate solution (which would be some universal system for managing state inside JAX)
2. It is still something that is going to require application specific choices. For example, if you scan over a neural net layer, should the weights be scanned over or not?

For now, the recommendation is to write your own, library specific implementations of higher order functions like scan and your own scope for overrides, if necessary.