Test math mode via __FAST_MATH__ instead of exp(-inf)

exp(-inf) agrees across math modes, so the old test passed regardless
of the selected mode. Branch on the Metal compiler's __FAST_MATH__ macro
instead, which is only defined under fast math, so the test fails if the
math mode is not forwarded to the compiler. Reuse one kernel name across
modes to exercise the library-cache rebuild path.

Author: shubhxho

python/tests/test_fast.py

@@ -1047,27 +1047,47 @@ def test_custom_metal_kernel_invalid_compile_options(self):
             )
 
     @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
-    def test_custom_metal_kernel_safe_math_mode(self):
-        kernel = mx.fast.metal_kernel(
-            name="safe_math_mode",
-            input_names=["inp"],
-            output_names=["out"],
-            source="""
-                uint elem = thread_position_in_grid.x;
-                out[elem] = metal::exp(inp[elem]);
-            """,
-            compile_options={"math_mode": "safe"},
-        )
-        a = mx.array([-float("inf"), 0.0], dtype=mx.float32)
-        out = kernel(
-            inputs=[a],
-            grid=(a.size, 1, 1),
-            threadgroup=(a.size, 1, 1),
-            output_shapes=[a.shape],
-            output_dtypes=[a.dtype],
-            stream=mx.gpu,
-        )[0]
-        self.assertTrue(mx.array_equal(out, mx.array([0.0, 1.0])))
+    def test_custom_metal_kernel_math_mode(self):
+        # Numerical special cases such as exp(-inf) can agree between math
+        # modes, so they don't reliably detect whether the mode was applied.
+        # Branch on the compiler's __FAST_MATH__ macro instead: it is defined
+        # only when fast math is enabled, so the test fails if the selected
+        # math mode is not forwarded to the Metal compiler.
+        source = """
+            uint elem = thread_position_in_grid.x;
+            #if defined(__FAST_MATH__) && __FAST_MATH__
+            out[elem] = 1.0f;
+            #else
+            out[elem] = 0.0f;
+            #endif
+        """
+
+        a = mx.zeros((4,), dtype=mx.float32)
+        expected = {
+            "safe": mx.zeros_like(a),
+            "fast": mx.ones_like(a),
+        }
+
+        # Reuse the same kernel name across modes so the library cache is forced
+        # to rebuild when the math mode changes, guarding against a stale build
+        # being returned for a different mode.
+        for mode, expected_out in expected.items():
+            kernel = mx.fast.metal_kernel(
+                name="math_mode",
+                input_names=["inp"],
+                output_names=["out"],
+                source=source,
+                compile_options={"math_mode": mode},
+            )
+            out = kernel(
+                inputs=[a],
+                grid=(a.size, 1, 1),
+                threadgroup=(a.size, 1, 1),
+                output_shapes=[a.shape],
+                output_dtypes=[a.dtype],
+                stream=mx.gpu,
+            )[0]
+            self.assertTrue(mx.array_equal(out, expected_out))
 
     @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
     def test_custom_kernel_mixed_dtypes(self):