LogCumSumExp (#2069)

2025-04-13 11:27:29 +03:00
parent 7275ac7523
commit e9e268336b
15 changed files with 209 additions and 3 deletions
@@ -1202,6 +1202,28 @@ void init_array(nb::module_& m) {
          nb::kw_only(),
          "stream"_a = nb::none(),
          "See :func:`max`.")
+      .def(
+          "logcumsumexp",
+          [](const mx::array& a,
+             std::optional<int> axis,
+             bool reverse,
+             bool inclusive,
+             mx::StreamOrDevice s) {
+            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);
+            }
+          },
+          "axis"_a = nb::none(),
+          nb::kw_only(),
+          "reverse"_a = false,
+          "inclusive"_a = true,
+          "stream"_a = nb::none(),
+          "See :func:`logcumsumexp`.")
      .def(
          "logsumexp",
          [](const mx::array& a,