jax

83a74efd - Update XLA dependency to use revision http://github.com/openxla/xla/commit/d27a833d899fb8561cb823594f0a4d920c2fd67f