Enable the C++ jax.jit fast code-path by default.
Enable the C++ jax.jit fast code-path by default.

Here is the list of things I think we should do/revisit at some point in the future:

1. Most importantly, think about Lazy expressions. I think they are powerful, and they can likely enable some optimization impossible otherwise (e.g. stack a ShardedDeviceArray and split it). We likely want to add support for them in the C++ jax.jit.

2.  Remove trivial computation support (currently fallback to Python), at least in the omnistaging code-path, as it will no longer necessary with omnistaging, and that the complexity do not justify the feature.

3. Not sure what to do about jax.jit(pmap). It currently fallback to Python when the Executable has more than one device. I understand it executes this on all cores (as a pmap), but will return the first result (from the first core). I am tempted to think this is a non-feature, that no one is looking for and that it can be achieved by doing pmap(f)(x)[0] instead. Having a single simple way of doing one thing is usually beneficial.

4. Not sure what to do about `DeviceConstant` support (currently fallback to Python). I do not understand yet why they exist and problem they solve.

5. Revisit the stikiness for jax.jit(f, device=device)(x). I think it should fail if x is stiky to another device, to let the user aware of the copies. My humble opinion is that, with an efficient device_put, we should write jitted_f(device_put(x, device)).
Improving performance is key for JAX users, and the way JAX let you control what lives where is wonderful. Helping the user be aware of their copies is helping in this direction/Doing copies for the users hurts performance.

(The support for `jax_debug_nans` could be improved.)

With the C++ jit:
└──╼ benchy  //third_party/py/jax/benchmarks:api_benchmark
name                                   time/op
jit_trivial_dispatch                   35.2µs ± 2%
jit_trivial                            36.8µs ± 2%
jit_simple_dispatch                    16.1µs ± 5%
jit_simple                             29.3µs ± 8%
jit_simple_many_args_dispatch           149µs ± 8%
jit_simple_many_args                    153µs ± 9%
jit_dispatch_without_transfer           144µs ± 7%
jit_dispatch_with_transfer              150µs ± 5%

Without:
└──╼ benchy  //third_party/py/jax/benchmarks:api_benchmark
name                                   time/op
jit_trivial_dispatch                   29.9µs ± 2%
jit_trivial                            31.8µs ± 4%
jit_simple_dispatch                    59.7µs ± 1%
jit_simple                             66.9µs ± 5%
jit_simple_many_args_dispatch           368µs ± 3%
jit_simple_many_args                    367µs ± 3%
jit_dispatch_without_transfer           346µs ± 9%
jit_dispatch_with_transfer              392µs ± 9%

See also https://github.com/google/jax/pull/4169 for context.

PiperOrigin-RevId: 336716707

google/jax,google/jax,google/jax,tensorflow/probability,tensorflow/probability,google/jax