jax

0d37f593 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/70ff8e0a1d970987ab63e1d8a9af7c70646eb522