jax

e6f26ff2 - Deprecate `jax.xla_computation`. Use JAX AOT APIs to get the equivalent of `jax.xla_computation` functionality.