zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-17 17:28:10 +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
2 changed files with 9 additions and 2 deletions
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++;