Address review: drop include_initial from cumsum/cumprod

Per maintainer request (#3731), remove the include_initial arg. It needed an
extra concatenate to be correct, which is the inefficient pattern angeloskath
flagged. cumsum/cumprod keep the efficient dtype arg only; cumulative_sum and
cumulative_prod remain pure aliases of them.

Author: katlun-lgtm

mlx/ops.cpp

@@ -3942,7 +3942,6 @@ array cumsum(
     bool reverse /* = false*/,
     bool inclusive /* = true*/,
     std::optional<Dtype> dtype /* = std::nullopt*/,
-    bool include_initial /* = false*/,
     StreamOrDevice s /* = {}*/) {
   int ndim = a.ndim();
   if (axis >= ndim || axis < -ndim) {
@@ -3954,31 +3953,22 @@ array cumsum(
   axis = (axis + a.ndim()) % a.ndim();
   auto x = dtype ? astype(a, *dtype, s) : a;
   auto out_type = x.dtype() == bool_ ? int32 : x.dtype();
-  auto out = array(
+  return array(
       x.shape(),
       out_type,
       std::make_shared<Scan>(
           to_stream(s), Scan::ReduceType::Sum, axis, reverse, inclusive),
       {x});
-  if (include_initial) {
-    Shape init_shape = out.shape();
-    init_shape[axis] = 1;
-    auto init = zeros(init_shape, out.dtype(), s);
-    out = reverse ? concatenate({out, init}, axis, s)
-                  : concatenate({init, out}, axis, s);
-  }
-  return out;
 }
 
 array cumsum(
     const array& a,
     bool reverse /* = false*/,
     bool inclusive /* = true*/,
     std::optional<Dtype> dtype /* = std::nullopt*/,
-    bool include_initial /* = false*/,
     StreamOrDevice s /* = {}*/) {
   return cumsum(
-      flatten(a, to_stream(s)), 0, reverse, inclusive, dtype, include_initial, to_stream(s));
+      flatten(a, to_stream(s)), 0, reverse, inclusive, dtype, to_stream(s));
 }
 
 array cumprod(
@@ -3987,7 +3977,6 @@ array cumprod(
     bool reverse /* = false*/,
     bool inclusive /* = true*/,
     std::optional<Dtype> dtype /* = std::nullopt*/,
-    bool include_initial /* = false*/,
     StreamOrDevice s /* = {}*/) {
   int ndim = a.ndim();
   if (axis >= ndim || axis < -ndim) {
@@ -3998,30 +3987,21 @@ array cumprod(
   }
   axis = (axis + a.ndim()) % a.ndim();
   auto x = dtype ? astype(a, *dtype, s) : a;
-  auto out = array(
+  return array(
       x.shape(),
       x.dtype(),
       std::make_shared<Scan>(
           to_stream(s), Scan::ReduceType::Prod, axis, reverse, inclusive),
       {x});
-  if (include_initial) {
-    Shape init_shape = out.shape();
-    init_shape[axis] = 1;
-    auto init = ones(init_shape, out.dtype(), s);
-    out = reverse ? concatenate({out, init}, axis, s)
-                  : concatenate({init, out}, axis, s);
-  }
-  return out;
 }
 
 array cumprod(
     const array& a,
     bool reverse /* = false*/,
     bool inclusive /* = true*/,
     std::optional<Dtype> dtype /* = std::nullopt*/,
-    bool include_initial /* = false*/,
     StreamOrDevice s /* = {}*/) {
-  return cumprod(flatten(a, s), 0, reverse, inclusive, dtype, include_initial, s);
+  return cumprod(flatten(a, s), 0, reverse, inclusive, dtype, s);
 }
 
 array cummax(

mlx/ops.h

@@ -1334,7 +1334,6 @@ MLX_API array cumsum(
     bool reverse = false,
     bool inclusive = true,
     std::optional<Dtype> dtype = std::nullopt,
-    bool include_initial = false,
     StreamOrDevice s = {});
 
 /** Cumulative sum of an array along the given axis. */
@@ -1344,7 +1343,6 @@ MLX_API array cumsum(
     bool reverse = false,
     bool inclusive = true,
     std::optional<Dtype> dtype = std::nullopt,
-    bool include_initial = false,
     StreamOrDevice s = {});
 
 /** Cumulative product of an array. */
@@ -1353,7 +1351,6 @@ MLX_API array cumprod(
     bool reverse = false,
     bool inclusive = true,
     std::optional<Dtype> dtype = std::nullopt,
-    bool include_initial = false,
     StreamOrDevice s = {});
 
 /** Cumulative product of an array along the given axis. */
@@ -1363,7 +1360,6 @@ MLX_API array cumprod(
     bool reverse = false,
     bool inclusive = true,
     std::optional<Dtype> dtype = std::nullopt,
-    bool include_initial = false,
     StreamOrDevice s = {});
 
 /** Cumulative max of an array. */

python/src/ops.cpp

@@ -3448,23 +3448,21 @@ void init_ops(nb::module_& m) {
          bool reverse,
          bool inclusive,
          std::optional<mx::Dtype> dtype,
-         bool include_initial,
          mx::StreamOrDevice s) {
         if (axis) {
-          return mx::cumsum(a, *axis, reverse, inclusive, dtype, include_initial, s);
+          return mx::cumsum(a, *axis, reverse, inclusive, dtype, s);
         }
-        return mx::cumsum(a, reverse, inclusive, dtype, include_initial, s);
+        return mx::cumsum(a, reverse, inclusive, dtype, s);
       },
       nb::arg(),
       "axis"_a = nb::none(),
       nb::kw_only(),
       "reverse"_a = false,
       "inclusive"_a = true,
       "dtype"_a = nb::none(),
-      "include_initial"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def cumsum(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, dtype: Optional[Dtype] = None, include_initial: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cumsum(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, dtype: Optional[Dtype] = None, stream: Union[None, Stream, Device] = None) -> array"),
       R"pbdoc(
         Return the cumulative sum of the elements along the given axis.
 
@@ -3477,8 +3475,6 @@ void init_ops(nb::module_& m) {
           inclusive (bool): The i-th element of the output includes the i-th
             element of the input.
           dtype (Dtype, optional): Cast the input to this type before summing.
-          include_initial (bool): Prepend the identity element (0) so the
-            output has one extra element along the given axis.
 
         Returns:
           array: The output array.
@@ -3490,23 +3486,21 @@ void init_ops(nb::module_& m) {
          bool reverse,
          bool inclusive,
          std::optional<mx::Dtype> dtype,
-         bool include_initial,
          mx::StreamOrDevice s) {
         if (axis) {
-          return mx::cumprod(a, *axis, reverse, inclusive, dtype, include_initial, s);
+          return mx::cumprod(a, *axis, reverse, inclusive, dtype, s);
         }
-        return mx::cumprod(a, reverse, inclusive, dtype, include_initial, s);
+        return mx::cumprod(a, reverse, inclusive, dtype, s);
       },
       nb::arg(),
       "axis"_a = nb::none(),
       nb::kw_only(),
       "reverse"_a = false,
       "inclusive"_a = true,
       "dtype"_a = nb::none(),
-      "include_initial"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def cumprod(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, dtype: Optional[Dtype] = None, include_initial: bool = False, stream: Union[None, Stream, Device] = None) -> array"),
+          "def cumprod(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, dtype: Optional[Dtype] = None, stream: Union[None, Stream, Device] = None) -> array"),
       R"pbdoc(
         Return the cumulative product of the elements along the given axis.
 
@@ -3519,8 +3513,6 @@ void init_ops(nb::module_& m) {
           inclusive (bool): The i-th element of the output includes the i-th
             element of the input.
           dtype (Dtype, optional): Cast the input to this type before multiplying.
-          include_initial (bool): Prepend the identity element (1) so the
-            output has one extra element along the given axis.
 
         Returns:
           array: The output array.

python/tests/test_ops.py

@@ -3449,21 +3449,11 @@ def test_to_from_fp8(self):
     def test_cumulative_sum_prod(self):
         a = mx.array([1, 2, 3, 4])
         self.assertEqual(mx.cumulative_sum(a).tolist(), [1, 3, 6, 10])
-        self.assertEqual(
-            mx.cumulative_sum(a, include_initial=True).tolist(), [0, 1, 3, 6, 10]
-        )
         self.assertEqual(mx.cumulative_prod(a).tolist(), [1, 2, 6, 24])
-        self.assertEqual(
-            mx.cumulative_prod(a, include_initial=True).tolist(), [1, 1, 2, 6, 24]
-        )
 
         m = mx.array([[1, 2], [3, 4]])
         self.assertEqual(mx.cumulative_sum(m, axis=0).tolist(), [[1, 2], [4, 6]])
         self.assertEqual(mx.cumulative_sum(m, axis=1).tolist(), [[1, 3], [3, 7]])
-        self.assertEqual(
-            mx.cumulative_sum(m, axis=1, include_initial=True).tolist(),
-            [[0, 1, 3], [0, 3, 7]],
-        )
         # axis=None flattens.
         self.assertEqual(mx.cumulative_sum(m).tolist(), [1, 3, 6, 10])
         self.assertEqual(mx.cumulative_sum(a, dtype=mx.float32).dtype, mx.float32)