zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-25 01:41:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Pickle: workaround for bfloat16

Browse Source
This commit is contained in:
Luca Arnaboldi 2024-03-04 23:45:17 +01:00
parent 6b6b4f0a5f
commit 25afe6bc09
2 changed files with 19 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
python/src/array.cpp
View File

@ -776,13 +776,25 @@ void init_array(py::module_& m) {
.def("__iter__", [](const array& a) { return ArrayPythonIterator(a); })
.def(py::pickle(
[](array& a) { // __getstate__
return py::array(buffer_info(a));
},
[](py::array npa) { // __setstate__
if (not py::isinstance<py::array>(npa)) {
throw std::runtime_error("Invalid state!");
auto dtype = a.dtype();
if (dtype == bfloat16) {
array b = astype(a, float32);
return py::make_tuple(
dtype_to_array_protocol(dtype), py::array(buffer_info(b)));
} else {
return py::make_tuple(
dtype_to_array_protocol(dtype), py::array(buffer_info(a)));
}
return np_array_to_mlx(npa, std::nullopt);
},
[](py::tuple t) { // __setstate__
if (t.size() != 2 or not py::isinstance<py::array>(t[1]) or
not py::isinstance<py::str>(t[0])) {
throw std::runtime_error("Invalid state passed to __setstate__!");
}
return np_array_to_mlx(
t[1].cast<py::array>(),
dtype_from_array_protocol(t[0].cast<std::string>()));
}))
.def("__copy__", [](const array& self) { return array(self); })
.def(

2
python/tests/test_array.py
View File

@ -671,7 +671,7 @@ class TestArray(mlx_tests.MLXTestCase):
mx.uint64,
mx.float16,
mx.float32,
# mx.bfloat16,
mx.bfloat16,
mx.complex64,
]