jax

60d71983 - Update XLA release.

Update XLA release. Includes https://github.com/tensorflow/tensorflow/commit/a66de7da15a4a5a4a1b3b8ad528f4d5960bc9195 which fixes #616 but requires a Jaxlib rebuild.