jax

db378d24 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/796e731794b58a5eb7f7b0f6bc382a9e28578a3e