jax

ff99faf2 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/7d9b40ba458f389d978603fb78d7fe4bff02100c