jax

9818cea0 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/b3a27413b814c30b825690b5a67fa1f0fc1cb0c5