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

fix array from list for > 32 bit types (#501)

Browse Source
This commit is contained in:
Awni Hannun 2024-01-19 15:49:25 -08:00 committed by GitHub
parent ddf50113c5
commit 6bf779e72b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 30 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
python/src/array.cpp
View File

@ -229,9 +229,28 @@ array array_from_list(
return array(vals.begin(), shape, specified_type.value_or(bool_));
}
case pyint: {
std::vector<int> vals;
fill_vector(pl, vals);
return array(vals.begin(), shape, specified_type.value_or(int32));
auto dtype = specified_type.value_or(int32);
if (dtype == int64) {
std::vector<int64_t> vals;
fill_vector(pl, vals);
return array(vals.begin(), shape, dtype);
} else if (dtype == uint64) {
std::vector<uint64_t> vals;
fill_vector(pl, vals);
return array(vals.begin(), shape, dtype);
} else if (dtype == uint32) {
std::vector<uint32_t> vals;
fill_vector(pl, vals);
return array(vals.begin(), shape, dtype);
} else if (is_floating_point(dtype)) {
std::vector<float> vals;
fill_vector(pl, vals);
return array(vals.begin(), shape, dtype);
} else {
std::vector<int> vals;
fill_vector(pl, vals);
return array(vals.begin(), shape, dtype);
}
}
case pyfloat: {
std::vector<float> vals;

8
python/tests/test_array.py
View File

@ -226,6 +226,14 @@ class TestArray(mlx_tests.MLXTestCase):
x = mx.array([1 + 0j, 2j, True, 0], mx.complex64)
self.assertEqual(x.tolist(), [1 + 0j, 2j, 1 + 0j, 0j])
xnp = np.array([0, 4294967295], dtype=np.uint32)
x = mx.array([0, 4294967295], dtype=mx.uint32)
self.assertTrue(np.array_equal(x, xnp))
xnp = np.array([0, 4294967295], dtype=np.float32)
x = mx.array([0, 4294967295], dtype=mx.float32)
self.assertTrue(np.array_equal(x, xnp))
def test_construction_from_lists_of_mlx_arrays(self):
dtypes = [
mx.bool_,