[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
pytorch · May 28, 2024 · 65b5ace · 65b5ace
1 parent 1a8c2fe
commit 65b5ace
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/test/test_gmm.py b/test/test_gmm.py
@@ -122,7 +122,7 @@ def test_gmm(self):
 
       atol, rtol = self._tolerances(lhs_dtype, rhs_dtype, out_dtype)
       np.testing.assert_allclose(
-          ref_out, np.array(out[0].cpu()), rtol=rtol, atol=atol)
+          ref_out, np.array(out.cpu()), rtol=rtol, atol=atol)
 
     # Make sure gmm doesn't fallback.
     self.assertNotIn("aten::", met.short_metrics_report())

diff --git a/torch_xla/experimental/custom_kernel.py b/torch_xla/experimental/custom_kernel.py
@@ -746,7 +746,7 @@ def gmm(lhs: torch.Tensor, rhs: torch.Tensor,
       group_offsets.to("xla"),
       group_ids.to("xla"),
       m_tile_ids.to("xla"), group_offset_torch, lhs, rhs
-  ], payload, [torch.Size([m, n])], [lhs.dtype])
+  ], payload, [torch.Size([m, n])], [lhs.dtype])[0]
 
 
 def non_xla_attetion(q, k, v, attention_type):