xla

8e6a5e5d - composibility of assume_pure and call_jax (#8989)