jax

c3b53a26 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/3cc8846c10052cc1c32c4db87866eac4e4cdbccd