jax

1066aa7b - Update XLA dependency to use revision http://github.com/openxla/xla/commit/80529bdfc1c3bc206d0809a474031e9c9008c9f0