[MPS] tril op not handling infs correctly (#150479)

pytorchbot · Isalia20 · web-flow · commit a3cd7b0cc4d7 · 2025-04-01T16:30:38.000-07:00
[MPS] tril op not handling infs correctly (#149866) Fixes #149813 Pull Request resolved: #149866 Approved by: https://github.com/malfet (cherry picked from commit ba46643) Co-authored-by: Isalia20 <irakli.salia854@gmail.com>
diff --git a/aten/src/ATen/native/mps/operations/TriangularOps.mm b/aten/src/ATen/native/mps/operations/TriangularOps.mm
@@ -107,7 +107,12 @@
                                                          numLowerTensor:negDiagMinusOneTensor
                                                          numUpperTensor:minusOneTensor
                                                                    name:nil];
-        outputTensor = [mpsGraph subtractionWithPrimaryTensor:inputTensor secondaryTensor:complementTensor name:nil];
+        MPSGraphTensor* zeroTensor = [mpsGraph constantWithScalar:0.0 dataType:getMPSDataType(self)];
+        MPSGraphTensor* mask = [mpsGraph equalWithPrimaryTensor:complementTensor secondaryTensor:zeroTensor name:nil];
+        outputTensor = [mpsGraph selectWithPredicateTensor:mask
+                                       truePredicateTensor:inputTensor
+                                      falsePredicateTensor:zeroTensor
+                                                      name:nil];
       }
 
       newCachedGraph->inputTensor_ = inputTensor;
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -7849,13 +7849,21 @@ def helper(shape, diag=0):
             self.assertEqual(tril_result, tril_result_cpu)
             self.assertEqual(x.grad, cpu_x.grad)
 
-        helper((2, 8, 4, 5))
-        helper((2, 8, 4, 5), diag=1)
-        helper((2, 8, 4, 5), diag=2)
-        helper((2, 8, 4, 5), diag=3)
-        helper((2, 8, 4, 5), diag=-1)
-        helper((2, 8, 4, 5), diag=-2)
-        helper((2, 8, 4, 5), diag=-3)
+        for diag in [0, 1, 2, 3, -1, -2, -3]:
+            helper((2, 8, 4, 5), diag=diag)
+
+        def helper_nans_infs(value, diag_vals=(0, 1, -2)):
+            """For nans and infs"""
+            mps_tensor = torch.full((2, 2, 5, 5), value, device="mps")
+            cpu_tensor = torch.full((2, 2, 5, 5), value, device="cpu")
+            for diag in diag_vals:
+                mps_result = torch.tril(mps_tensor, diagonal=diag)
+                cpu_result = torch.tril(cpu_tensor, diagonal=diag)
+                self.assertEqual(mps_result, cpu_result, f"Mismatch for diag={diag}")
+
+        helper_nans_infs(float("inf"))
+        helper_nans_infs(float("-inf"))
+        helper_nans_infs(float("nan"))
 
     # test eye
     def test_eye(self):
