xla

a9b4fadf - [Pallas] Make gmm functional (#7117)

[Pallas] Make gmm functional (#7117) Summary: This should be the last PR needed to make gmm functional. Now, everything should be executed in the xla devices and produces the correct output. Test Plan: python test/test_gmm.py