xla

f39434ab - [Kernel] Use call_jax to simplify the gmm pallas kernel wrapper (#9180)