jax

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