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

Fix cumulative operations when axis=None (#2653)

Browse Source
This commit is contained in:
AN Long
2025-10-09 07:25:38 +09:00
committed by GitHub
parent f5d4397e5c
commit 85a8824a8c
3 changed files with 81 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
python/src/array.cpp
View File

@@ -1275,10 +1275,7 @@ void init_array(nb::module_& m) {
if (axis) {
return mx::logcumsumexp(a, *axis, reverse, inclusive, s);
} else {
// TODO: Implement that in the C++ API as well. See concatenate
// above.
return mx::logcumsumexp(
mx::reshape(a, {-1}, s), 0, reverse, inclusive, s);
return mx::logcumsumexp(a, reverse, inclusive, s);
}
},
"axis"_a = nb::none(),
@@ -1408,9 +1405,7 @@ void init_array(nb::module_& m) {
if (axis) {
return mx::cumsum(a, *axis, reverse, inclusive, s);
} else {
// TODO: Implement that in the C++ API as well. See concatenate
// above.
return mx::cumsum(reshape(a, {-1}, s), 0, reverse, inclusive, s);
return mx::cumsum(a, reverse, inclusive, s);
}
},
"axis"_a = nb::none(),
@@ -1429,10 +1424,7 @@ void init_array(nb::module_& m) {
if (axis) {
return mx::cumprod(a, *axis, reverse, inclusive, s);
} else {
// TODO: Implement that in the C++ API as well. See concatenate
// above.
return mx::cumprod(
mx::reshape(a, {-1}, s), 0, reverse, inclusive, s);
return mx::cumprod(a, reverse, inclusive, s);
}
},
"axis"_a = nb::none(),
@@ -1451,10 +1443,7 @@ void init_array(nb::module_& m) {
if (axis) {
return mx::cummax(a, *axis, reverse, inclusive, s);
} else {
// TODO: Implement that in the C++ API as well. See concatenate
// above.
return mx::cummax(
mx::reshape(a, {-1}, s), 0, reverse, inclusive, s);
return mx::cummax(a, reverse, inclusive, s);
}
},
"axis"_a = nb::none(),
@@ -1473,10 +1462,7 @@ void init_array(nb::module_& m) {
if (axis) {
return mx::cummin(a, *axis, reverse, inclusive, s);
} else {
// TODO: Implement that in the C++ API as well. See concatenate
// above.
return mx::cummin(
mx::reshape(a, {-1}, s), 0, reverse, inclusive, s);
return mx::cummin(a, reverse, inclusive, s);
}
},
"axis"_a = nb::none(),