zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 17:31: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

23
python/src/array.cpp
View File

@ -162,14 +162,6 @@ PyScalarT validate_shape(
shape, shape,
idx + 1, idx + 1,
all_python_primitive_elements); 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)) { } else if (nb::isinstance<array>(l)) {
all_python_primitive_elements = false; all_python_primitive_elements = false;
auto arr = nb::cast<array>(l); auto arr = nb::cast<array>(l);
@ -183,12 +175,27 @@ PyScalarT validate_shape(
throw std::invalid_argument( throw std::invalid_argument(
"Initialization encountered non-uniform length."); "Initialization encountered non-uniform length.");
} }
} 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 { } else {
std::ostringstream msg; std::ostringstream msg;
msg << "Invalid type " << nb::type_name(l.type()).c_str() msg << "Invalid type " << nb::type_name(l.type()).c_str()
<< " received in array initialization."; << " received in array initialization.";
throw std::invalid_argument(msg.str()); 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); type = std::max(type, t);
} }
return type; return type;

4
python/tests/test_array.py
View File

@ -391,7 +391,9 @@ class TestArray(mlx_tests.MLXTestCase):
# shape check from `stack()` # shape check from `stack()`
with self.assertRaises(ValueError) as e: with self.assertRaises(ValueError) as e:
mx.array([x, 1.0]) 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` # shape check from `validate_shape`
with self.assertRaises(ValueError) as e: with self.assertRaises(ValueError) as e: