[MPS] Avoid outputing zeros from exponential_ for MPS (#159386)

kurtamohler · pytorchmergebot · commit 70d2e9ba455c · 2025-07-30T00:20:31.000Z
Fixes #159103 Pull Request resolved: #159386 Approved by: https://github.com/malfet
diff --git a/aten/src/ATen/native/mps/operations/Distributions.mm b/aten/src/ATen/native/mps/operations/Distributions.mm
@@ -418,8 +418,9 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, std::optional<Generator
     MPSGraphTensor* logTensor = [mpsGraph logarithmWithTensor:subtractTensor name:nil];
     return [mpsGraph divisionWithPrimaryTensor:logTensor secondaryTensor:minusLambdaTensor name:nil];
   };
+  auto eps = std::numeric_limits<float>::epsilon();
   return mps::random_mps_impl<double>(self,
-                                      0.0,
+                                      eps,
                                       1.0,
                                       std::nullopt,
                                       std::nullopt,
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -7681,6 +7681,12 @@ def test_exponential_1(self):
         self.assertEqual(Exponential(0.2).sample((1,)).size(), (1,))
         self.assertEqual(Exponential(50.0).sample((1,)).size(), (1,))
 
+    @parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+    def test_exponential_nonzero(self, dtype):
+        for _ in range(100):
+            a = torch.empty(32_000, device="mps", dtype=dtype).exponential_()
+            self.assertTrue((a != 0).all())
+
     # Test add
     def test_add_sub(self):
         def helper(shape, alpha, op_name, inplace):
