[Pallas] Make gmm functional

[alanwaketan]

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

[JackCaoG]

AssertionError: "min must be less than max." does not match "min must be less than or equal to max." lol seems like you need to update the test.

[alanwaketan]

Thanks, Jack.