jax

c8c228c6 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/53b74f6d91b5be992f853c3e6e8c1780357f42c5