xla

[Kernel] Use call_jax to simplify the gmm pallas kernel wrapper

#9180

Merged