xla

65b5ace8 - [Pallas] Make gmm output a tensor (#7120)

[Pallas] Make gmm output a tensor (#7120) Summary: Somehow gmm outputs a list intead of a tensor. Let's output a tensor. Test Plan: python test/test_gmm.py