zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-02 23:31:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Compile float64 functions on CPU (#2311)

Browse Source
This commit is contained in:
Awni Hannun 2025-06-24 10:18:52 -07:00 committed by GitHub
parent 5adf185f86
commit 81bb9a2a9e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 24 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/common/compiled.cpp
View File

@ -14,6 +14,8 @@ void print_constant(std::ostream& os, const array& x) {
return print_float_constant<float16_t>(os, x);
case bfloat16:
return print_float_constant<bfloat16_t>(os, x);
case float64:
return print_float_constant<double>(os, x);
case complex64:
return print_complex_constant<complex64_t>(os, x);
case int8:
@ -50,6 +52,8 @@ std::string get_type_string(Dtype d) {
return "float16_t";
case bfloat16:
return "bfloat16_t";
case float64:
return "double";
case complex64:
return "complex64_t";
case bool_:

8
mlx/backend/common/compiled.h
View File

@ -18,8 +18,12 @@ std::string get_type_string(Dtype d);
template <typename T>
void print_float_constant(std::ostream& os, const array& x) {
auto old_precision = os.precision();
os << std::setprecision(std::numeric_limits<float>::digits10 + 1)
<< x.item<T>() << std::setprecision(old_precision);
if constexpr (std::is_same_v<T, double>) {
os << std::setprecision(std::numeric_limits<double>::digits10 + 1);
} else {
os << std::setprecision(std::numeric_limits<float>::digits10 + 1);
}
os << x.item<T>() << std::setprecision(old_precision);
}
template <typename T>

2
python/src/convert.cpp
View File

@ -205,6 +205,8 @@ nb::object to_scalar(mx::array& a) {
return nb::cast(static_cast<float>(a.item<mx::bfloat16_t>()));
case mx::complex64:
return nb::cast(a.item<std::complex<float>>());
case mx::float64:
return nb::cast(a.item<double>());
default:
throw nb::type_error("type cannot be converted to Python scalar.");
}

12
python/tests/test_compile.py
View File

@ -2,6 +2,7 @@
import gc
import io
import math
import unittest
from functools import partial
@ -979,6 +980,17 @@ class TestCompile(mlx_tests.MLXTestCase):
self.assertEqual(mem_pre, mem_post)
def test_double_constant(self):
with mx.stream(mx.cpu):
x = mx.array(1.0, dtype=mx.float64)
def fun(x):
return (x + math.pi) * 2.0
y = fun(x).item()
y_compiled = mx.compile(fun)(x).item()
self.assertEqual(y, y_compiled)
if __name__ == "__main__":
mlx_tests.MLXTestRunner()