zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 09:21:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix array initialization from list (#942)

Browse Source 
* Fix array initialization from list

* Change the error message in the test
This commit is contained in:
Angelos Katharopoulos 2024-04-01 06:27:52 -07:00 committed by GitHub
parent 110d9b149d
commit 02fedbf1da
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 23 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
python/src/array.cpp
View File

@ -162,14 +162,6 @@ PyScalarT validate_shape(
shape,
idx + 1,
all_python_primitive_elements);
} else if (nb::isinstance<nb::bool_>(l)) {
t = pybool;
} else if (nb::isinstance<nb::int_>(l)) {
t = pyint;
} else if (nb::isinstance<nb::float_>(l)) {
t = pyfloat;
} else if (PyComplex_Check(l.ptr())) {
t = pycomplex;
} else if (nb::isinstance<array>(l)) {
all_python_primitive_elements = false;
auto arr = nb::cast<array>(l);
@ -184,10 +176,25 @@ PyScalarT validate_shape(
"Initialization encountered non-uniform length.");
}
} else {
std::ostringstream msg;
msg << "Invalid type " << nb::type_name(l.type()).c_str()
<< " received in array initialization.";
throw std::invalid_argument(msg.str());
if (nb::isinstance<nb::bool_>(l)) {
t = pybool;
} else if (nb::isinstance<nb::int_>(l)) {
t = pyint;
} else if (nb::isinstance<nb::float_>(l)) {
t = pyfloat;
} else if (PyComplex_Check(l.ptr())) {
t = pycomplex;
} else {
std::ostringstream msg;
msg << "Invalid type " << nb::type_name(l.type()).c_str()
<< " received in array initialization.";
throw std::invalid_argument(msg.str());
}
if (idx + 1 != shape.size()) {
throw std::invalid_argument(
"Initialization encountered non-uniform length.");
}
}
type = std::max(type, t);
}
@ -1440,4 +1447,4 @@ void init_array(nb::module_& m) {
R"pbdoc(
Extract a diagonal or construct a diagonal matrix.
)pbdoc");
}
}

4
python/tests/test_array.py
View File

@ -391,7 +391,9 @@ class TestArray(mlx_tests.MLXTestCase):
# shape check from `stack()`
with self.assertRaises(ValueError) as e:
mx.array([x, 1.0])
self.assertEqual(str(e.exception), "All arrays must have the same shape")
self.assertEqual(
str(e.exception), "Initialization encountered non-uniform length."
)
# shape check from `validate_shape`
with self.assertRaises(ValueError) as e: