jax

4ea536c2 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/ac24cc91d12634f06c21a30f170c6a117ff615f9