zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-14 20:41:13 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix format (#1539)

Browse Source
This commit is contained in:
Awni Hannun 2024-10-28 18:29:14 -07:00 committed by GitHub
parent 015c247393
commit d2ff04a4f2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
python/src/array.cpp
View File

@ -848,12 +848,19 @@ void init_array(nb::module_& m) {
.def( .def(
"__format__", "__format__",
[](array& a, nb::object format_spec) { [](array& a, nb::object format_spec) {
if (a.ndim() > 0) { if (nb::len(nb::str(format_spec)) > 0 && a.ndim() > 0) {
throw nb::type_error( throw nb::type_error(
"unsupported format string passed to mx.array.__format__"); "unsupported format string passed to mx.array.__format__");
} else if (a.ndim() == 0) {
auto obj = to_scalar(a);
return nb::cast<std::string>(
nb::handle(PyObject_Format(obj.ptr(), format_spec.ptr())));
} else {
nb::gil_scoped_release nogil;
std::ostringstream os;
os << a;
return os.str();
} }
auto obj = to_scalar(a);
return nb::str(PyObject_Format(obj.ptr(), format_spec.ptr()));
}) })
.def( .def(
"flatten", "flatten",

3
python/tests/test_array.py
View File

@ -1885,6 +1885,9 @@ class TestArray(mlx_tests.MLXTestCase):
with self.assertRaises(TypeError): with self.assertRaises(TypeError):
s = f"{a:.2f}" s = f"{a:.2f}"
a = mx.array([1, 2, 3])
self.assertEqual(f"{a}", "array([1, 2, 3], dtype=int32)")
def test_deep_graphs(self): def test_deep_graphs(self):
# The following tests should simply run cleanly without a segfault or # The following tests should simply run cleanly without a segfault or
# crash due to exceeding recursion depth limits. # crash due to exceeding recursion depth limits.