jax

b1272723 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/59a18ff2183fd80a0b0af123eff42cd85ba72bb8