diff --git a/mlx/backend/common/compiled.h b/mlx/backend/common/compiled.h index 8c6466da03..94ba38b8f6 100644 --- a/mlx/backend/common/compiled.h +++ b/mlx/backend/common/compiled.h @@ -37,9 +37,9 @@ void print_float_constant(std::ostream& os, const array& x) { auto old_precision = os.precision(); if constexpr (std::is_same_v) { - os << std::setprecision(std::numeric_limits::digits10 + 1); + os << std::setprecision(std::numeric_limits::max_digits10); } else { - os << std::setprecision(std::numeric_limits::digits10 + 1); + os << std::setprecision(std::numeric_limits::max_digits10); } os << value << std::setprecision(old_precision); } diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 03d0027d7e..a68ff2d5ec 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -87,6 +87,12 @@ def test_compile_nonfinite_constants(self): self.assertEqual(out[0].item(), 1.0) self.assertEqual(out[1].item(), float("-inf")) + def test_compile_float_constant_precision(self): + x = mx.ones((4,), dtype=mx.float32) + for constant in (1 / 3, 128**-0.5, 0.7071067811865476): + fun = lambda x, constant=constant: (x * x) * constant + self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x))) + def test_compile_tuple_output_in_thread(self): @mx.compile def fun(x): @@ -1309,14 +1315,13 @@ def f(x): def test_double_constant(self): with mx.stream(mx.cpu): - x = mx.array(1.0, dtype=mx.float64) + x = mx.array([1.0], dtype=mx.float64) + constant = math.nextafter(1.0, 2.0) def fun(x): - return (x + math.pi) * 2.0 + return (x * x) * constant - y = fun(x).item() - y_compiled = mx.compile(fun)(x).item() - self.assertEqual(y, y_compiled) + self.assertTrue(mx.array_equal(fun(x), mx.compile(fun)(x))) def test_shared_broadcast(self): def fun(x, y, z):