xla

aeed61a9 - [Pallas] Introduce GMM(torch.autograd.Function) (#7152)

[Pallas] Introduce GMM(torch.autograd.Function) (#7152) Summary: This pull request make GMM as a torch.autograd.Function such that we can use torch.autograd.backward instead of manual backpropagation. Test Plan: python test/test_gmm.py