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

Fix unsqueeze with None (#899)

Browse Source 
* Fix unsqueeze with None

* Clean unnecessary files
This commit is contained in:
Abdussamet Türker 2024-03-26 23:59:44 +03:00 committed by GitHub
parent 570f2bf29e
commit 5611e1a95e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
python/src/indexing.cpp
View File

@ -367,6 +367,7 @@ array mlx_get_item_nd(array src, const nb::tuple& entries) {
}
bool squeeze_needed = false;
bool unsqueeze_needed = false;
// Slice handling
{
@ -395,17 +396,19 @@ array mlx_get_item_nd(array src, const nb::tuple& entries) {
}
axis++;
} else {
unsqueeze_needed = true;
}
}
src = slice(src, starts, ends, strides);
}
// Unsqueeze handling
if (remaining_indices.size() > src.ndim() || squeeze_needed) {
if (unsqueeze_needed || squeeze_needed) {
std::vector<int> out_shape;
int axis = 0;
for (auto& idx : remaining_indices) {
if (idx.is_none()) {
if (unsqueeze_needed && idx.is_none()) {
out_shape.push_back(1);
} else if (squeeze_needed && nb::isinstance<nb::int_>(idx)) {
axis++;

4
python/tests/test_array.py
View File

@ -763,6 +763,10 @@ class TestArray(mlx_tests.MLXTestCase):
a_sliced_npy = np.asarray(a_sliced_mlx)
self.assertTrue(np.array_equal(a_sliced_npy, a_npy[None]))
a_sliced_mlx = a_mlx[:, None]
a_sliced_npy = np.asarray(a_sliced_mlx)
self.assertTrue(np.array_equal(a_sliced_npy, a_npy[:, None]))
# Multi dim indexing, all ints
self.assertEqual(a_mlx[0, 0].item(), 0)
self.assertEqual(a_mlx[0, 0].ndim, 0)