jax

f2dc38d2 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/ce8d98f849c2c774122b520d51976671c166e622