jax

8552ae99 - Update XLA dependency to use revision http://github.com/openxla/xla/commit/1924eec184be486bce582db95964636696ea3f03