jax

c76c9e73 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/20fc52c010400f5466804ce8230872e80a31fe8b