jax

174025a6 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/39fbe6fa509bd4c93fc35e7a9690b6e283fe634b