jax

0f31a1b0 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/9b979f4aa0aff2bd68756d1c5d23d250213b462d