mlx_rs/ops/
cumulative.rs

1use crate::error::Result;
2use crate::utils::guard::Guarded;
3use crate::{Array, Stream};
4use mlx_internal_macros::generate_macro;
5
6/// Axis and scan direction for [`Array::logcumsumexp`].
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub struct LogCumsumExpOptions {
9    /// `None` flattens the input before scanning.
10    pub axis: Option<i32>,
11
12    /// Scan in reverse order.
13    pub reverse: bool,
14
15    /// Include the current value in each output position.
16    pub inclusive: bool,
17}
18
19impl Default for LogCumsumExpOptions {
20    fn default() -> Self {
21        Self {
22            axis: None,
23            reverse: false,
24            inclusive: true,
25        }
26    }
27}
28fn optional_dtype_none() -> mlx_sys::mlx_optional_dtype {
29    mlx_sys::mlx_optional_dtype {
30        value: mlx_sys::mlx_dtype__MLX_FLOAT32,
31        has_value: false,
32    }
33}
34
35impl Array {
36    /// Compute a stable cumulative `LogAddExp` scan.
37    ///
38    /// This does not form `log(cumsum(exp(x)))`. Exclusive scans use negative infinity as the
39    /// shifted seed.
40    ///
41    /// ```rust
42    /// use mlx_rs::{array, ops::LogCumsumExpOptions};
43    ///
44    /// let output = array!([0.0, 1.0, 2.0])
45    ///     .logcumsumexp(LogCumsumExpOptions::default())
46    ///     .unwrap();
47    /// assert_eq!(output.shape(), &[3]);
48    /// ```
49    pub fn logcumsumexp(&self, options: LogCumsumExpOptions) -> Result<Array> {
50        let stream = Stream::thread_local_or_default();
51        match options.axis {
52            Some(axis) => Array::try_from_op(|res| unsafe {
53                mlx_sys::mlx_logcumsumexp_axis(
54                    res,
55                    self.as_ptr(),
56                    axis,
57                    options.reverse,
58                    options.inclusive,
59                    stream.as_ref().as_ptr(),
60                )
61            }),
62            None => Array::try_from_op(|res| unsafe {
63                mlx_sys::mlx_logcumsumexp(
64                    res,
65                    self.as_ptr(),
66                    options.reverse,
67                    options.inclusive,
68                    stream.as_ref().as_ptr(),
69                )
70            }),
71        }
72    }
73
74    /// Return the cumulative maximum of the elements along the given axis returning an error if the inputs are invalid.
75    ///
76    /// # Params
77    ///
78    /// - axis: Optional axis to compute the cumulative maximum over. If unspecified the cumulative maximum of the flattened array is returned.
79    /// - reverse: If true, the cumulative maximum is computed in reverse - defaults to false if unspecified.
80    /// - inclusive: If true, the i-th element of the output includes the i-th element of the input - defaults to true if unspecified.
81    ///
82    /// # Example
83    ///
84    /// ```rust
85    /// use mlx_rs::Array;
86    /// let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
87    ///
88    /// // result is [[5, 8], [5, 9]] -- cumulative max along the columns
89    /// let result = array.cummax(0, None, None).unwrap();
90    /// ```
91    pub fn cummax(
92        &self,
93        axis: impl Into<Option<i32>>,
94        reverse: impl Into<Option<bool>>,
95        inclusive: impl Into<Option<bool>>,
96    ) -> Result<Array> {
97        let stream = Stream::thread_local_or_default();
98        let stream = stream.as_ref();
99
100        match axis.into() {
101            Some(axis) => Array::try_from_op(|res| unsafe {
102                mlx_sys::mlx_cummax_axis(
103                    res,
104                    self.as_ptr(),
105                    axis,
106                    reverse.into().unwrap_or(false),
107                    inclusive.into().unwrap_or(true),
108                    stream.as_ptr(),
109                )
110            }),
111            None => Array::try_from_op(|res| unsafe {
112                mlx_sys::mlx_cummax(
113                    res,
114                    self.as_ptr(),
115                    reverse.into().unwrap_or(false),
116                    inclusive.into().unwrap_or(true),
117                    stream.as_ptr(),
118                )
119            }),
120        }
121    }
122
123    /// Compatibility shim for [`cummax`].
124    #[deprecated(
125        since = "0.26.0",
126        note = "use `with_stream` or `with_device` around `cummax`"
127    )]
128    pub fn cummax_device(
129        &self,
130        axis: impl Into<Option<i32>>,
131        reverse: impl Into<Option<bool>>,
132        inclusive: impl Into<Option<bool>>,
133        stream: impl AsRef<Stream>,
134    ) -> Result<Array> {
135        crate::with_stream(stream.as_ref(), || self.cummax(axis, reverse, inclusive))
136    }
137
138    /// Return the cumulative minimum of the elements along the given axis returning an error if the inputs are invalid.
139    ///
140    /// # Params
141    ///
142    /// - axis: Optional axis to compute the cumulative minimum over. If unspecified the cumulative maximum of the flattened array is returned.
143    /// - reverse: If true, the cumulative minimum is computed in reverse - defaults to false if unspecified.
144    /// - inclusive: If true, the i-th element of the output includes the i-th element of the input - defaults to true if unspecified.
145    ///
146    /// # Example
147    ///
148    /// ```rust
149    /// use mlx_rs::Array;
150    /// let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
151    ///
152    /// // result is [[5, 8], [4, 8]] -- cumulative min along the columns
153    /// let result = array.cummin(0, None, None).unwrap();
154    /// ```
155    pub fn cummin(
156        &self,
157        axis: impl Into<Option<i32>>,
158        reverse: impl Into<Option<bool>>,
159        inclusive: impl Into<Option<bool>>,
160    ) -> Result<Array> {
161        let stream = Stream::thread_local_or_default();
162        let stream = stream.as_ref();
163
164        match axis.into() {
165            Some(axis) => Array::try_from_op(|res| unsafe {
166                mlx_sys::mlx_cummin_axis(
167                    res,
168                    self.as_ptr(),
169                    axis,
170                    reverse.into().unwrap_or(false),
171                    inclusive.into().unwrap_or(true),
172                    stream.as_ptr(),
173                )
174            }),
175            None => Array::try_from_op(|res| unsafe {
176                mlx_sys::mlx_cummin(
177                    res,
178                    self.as_ptr(),
179                    reverse.into().unwrap_or(false),
180                    inclusive.into().unwrap_or(true),
181                    stream.as_ptr(),
182                )
183            }),
184        }
185    }
186
187    /// Compatibility shim for [`cummin`].
188    #[deprecated(
189        since = "0.26.0",
190        note = "use `with_stream` or `with_device` around `cummin`"
191    )]
192    pub fn cummin_device(
193        &self,
194        axis: impl Into<Option<i32>>,
195        reverse: impl Into<Option<bool>>,
196        inclusive: impl Into<Option<bool>>,
197        stream: impl AsRef<Stream>,
198    ) -> Result<Array> {
199        crate::with_stream(stream.as_ref(), || self.cummin(axis, reverse, inclusive))
200    }
201
202    /// Return the cumulative product of the elements along the given axis returning an error if the inputs are invalid.
203    ///
204    /// # Params
205    ///
206    /// - axis: Optional axis to compute the cumulative product over. If unspecified the cumulative maximum of the flattened array is returned.
207    /// - reverse: If true, the cumulative product is computed in reverse - defaults to false if unspecified.
208    /// - inclusive: If true, the i-th element of the output includes the i-th element of the input - defaults to true if unspecified.
209    ///
210    /// # Example
211    ///
212    /// ```rust
213    /// use mlx_rs::Array;
214    /// let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
215    ///
216    /// // result is [[5, 8], [20, 72]] -- cumulative min along the columns
217    /// let result = array.cumprod(0, None, None).unwrap();
218    /// ```
219    pub fn cumprod(
220        &self,
221        axis: impl Into<Option<i32>>,
222        reverse: impl Into<Option<bool>>,
223        inclusive: impl Into<Option<bool>>,
224    ) -> Result<Array> {
225        let stream = Stream::thread_local_or_default();
226        let stream = stream.as_ref();
227
228        match axis.into() {
229            Some(axis) => Array::try_from_op(|res| unsafe {
230                mlx_sys::mlx_cumprod_axis(
231                    res,
232                    self.as_ptr(),
233                    axis,
234                    reverse.into().unwrap_or(false),
235                    inclusive.into().unwrap_or(true),
236                    optional_dtype_none(),
237                    stream.as_ptr(),
238                )
239            }),
240            None => Array::try_from_op(|res| unsafe {
241                mlx_sys::mlx_cumprod(
242                    res,
243                    self.as_ptr(),
244                    reverse.into().unwrap_or(false),
245                    inclusive.into().unwrap_or(true),
246                    optional_dtype_none(),
247                    stream.as_ptr(),
248                )
249            }),
250        }
251    }
252
253    /// Compatibility shim for [`cumprod`].
254    #[deprecated(
255        since = "0.26.0",
256        note = "use `with_stream` or `with_device` around `cumprod`"
257    )]
258    pub fn cumprod_device(
259        &self,
260        axis: impl Into<Option<i32>>,
261        reverse: impl Into<Option<bool>>,
262        inclusive: impl Into<Option<bool>>,
263        stream: impl AsRef<Stream>,
264    ) -> Result<Array> {
265        crate::with_stream(stream.as_ref(), || self.cumprod(axis, reverse, inclusive))
266    }
267
268    /// Return the cumulative sum of the elements along the given axis returning an error if the inputs are invalid.
269    ///
270    /// # Params
271    ///
272    /// - axis: Optional axis to compute the cumulative sum over. If unspecified the cumulative maximum of the flattened array is returned.
273    /// - reverse: If true, the cumulative sum is computed in reverse - defaults to false if unspecified.
274    /// - inclusive: If true, the i-th element of the output includes the i-th element of the input - defaults to true if unspecified.
275    ///
276    /// # Example
277    ///
278    /// ```rust
279    /// use mlx_rs::Array;
280    /// let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
281    ///
282    /// // result is [[5, 8], [9, 17]] -- cumulative min along the columns
283    /// let result = array.cumsum(0, None, None).unwrap();
284    /// ```
285    pub fn cumsum(
286        &self,
287        axis: impl Into<Option<i32>>,
288        reverse: impl Into<Option<bool>>,
289        inclusive: impl Into<Option<bool>>,
290    ) -> Result<Array> {
291        let stream = Stream::thread_local_or_default();
292        let stream = stream.as_ref();
293
294        match axis.into() {
295            Some(axis) => Array::try_from_op(|res| unsafe {
296                mlx_sys::mlx_cumsum_axis(
297                    res,
298                    self.as_ptr(),
299                    axis,
300                    reverse.into().unwrap_or(false),
301                    inclusive.into().unwrap_or(true),
302                    optional_dtype_none(),
303                    stream.as_ptr(),
304                )
305            }),
306            None => Array::try_from_op(|res| unsafe {
307                mlx_sys::mlx_cumsum(
308                    res,
309                    self.as_ptr(),
310                    reverse.into().unwrap_or(false),
311                    inclusive.into().unwrap_or(true),
312                    optional_dtype_none(),
313                    stream.as_ptr(),
314                )
315            }),
316        }
317    }
318
319    /// Compatibility shim for [`cumsum`].
320    #[deprecated(
321        since = "0.26.0",
322        note = "use `with_stream` or `with_device` around `cumsum`"
323    )]
324    pub fn cumsum_device(
325        &self,
326        axis: impl Into<Option<i32>>,
327        reverse: impl Into<Option<bool>>,
328        inclusive: impl Into<Option<bool>>,
329        stream: impl AsRef<Stream>,
330    ) -> Result<Array> {
331        crate::with_stream(stream.as_ref(), || self.cumsum(axis, reverse, inclusive))
332    }
333}
334
335/// See [`Array::cummax`]
336pub fn cummax(
337    a: impl AsRef<Array>,
338    axis: impl Into<Option<i32>>,
339    reverse: impl Into<Option<bool>>,
340    inclusive: impl Into<Option<bool>>,
341) -> Result<Array> {
342    a.as_ref().cummax(axis, reverse, inclusive)
343}
344
345/// Compatibility shim for [`cummax`].
346#[generate_macro(customize(forwarding_shim = true))]
347#[deprecated(
348    since = "0.26.0",
349    note = "use `with_stream` or `with_device` around `cummax`"
350)]
351pub fn cummax_device(
352    a: impl AsRef<Array>,
353    #[optional] axis: impl Into<Option<i32>>,
354    #[optional] reverse: impl Into<Option<bool>>,
355    #[optional] inclusive: impl Into<Option<bool>>,
356    #[optional] stream: impl AsRef<Stream>,
357) -> Result<Array> {
358    crate::with_stream(stream.as_ref(), || cummax(a, axis, reverse, inclusive))
359}
360
361/// See [`Array::cummin`]
362pub fn cummin(
363    a: impl AsRef<Array>,
364    axis: impl Into<Option<i32>>,
365    reverse: impl Into<Option<bool>>,
366    inclusive: impl Into<Option<bool>>,
367) -> Result<Array> {
368    a.as_ref().cummin(axis, reverse, inclusive)
369}
370
371/// Compatibility shim for [`cummin`].
372#[generate_macro(customize(forwarding_shim = true))]
373#[deprecated(
374    since = "0.26.0",
375    note = "use `with_stream` or `with_device` around `cummin`"
376)]
377pub fn cummin_device(
378    a: impl AsRef<Array>,
379    #[optional] axis: impl Into<Option<i32>>,
380    #[optional] reverse: impl Into<Option<bool>>,
381    #[optional] inclusive: impl Into<Option<bool>>,
382    #[optional] stream: impl AsRef<Stream>,
383) -> Result<Array> {
384    crate::with_stream(stream.as_ref(), || cummin(a, axis, reverse, inclusive))
385}
386
387/// See [`Array::cumprod`]
388pub fn cumprod(
389    a: impl AsRef<Array>,
390    axis: impl Into<Option<i32>>,
391    reverse: impl Into<Option<bool>>,
392    inclusive: impl Into<Option<bool>>,
393) -> Result<Array> {
394    a.as_ref().cumprod(axis, reverse, inclusive)
395}
396
397/// Compatibility shim for [`cumprod`].
398#[generate_macro(customize(forwarding_shim = true))]
399#[deprecated(
400    since = "0.26.0",
401    note = "use `with_stream` or `with_device` around `cumprod`"
402)]
403pub fn cumprod_device(
404    a: impl AsRef<Array>,
405    #[optional] axis: impl Into<Option<i32>>,
406    #[optional] reverse: impl Into<Option<bool>>,
407    #[optional] inclusive: impl Into<Option<bool>>,
408    #[optional] stream: impl AsRef<Stream>,
409) -> Result<Array> {
410    crate::with_stream(stream.as_ref(), || cumprod(a, axis, reverse, inclusive))
411}
412
413/// See [`Array::cumsum`]
414pub fn cumsum(
415    a: impl AsRef<Array>,
416    axis: impl Into<Option<i32>>,
417    reverse: impl Into<Option<bool>>,
418    inclusive: impl Into<Option<bool>>,
419) -> Result<Array> {
420    a.as_ref().cumsum(axis, reverse, inclusive)
421}
422
423/// Compatibility shim for [`cumsum`].
424#[generate_macro(customize(forwarding_shim = true))]
425#[deprecated(
426    since = "0.26.0",
427    note = "use `with_stream` or `with_device` around `cumsum`"
428)]
429pub fn cumsum_device(
430    a: impl AsRef<Array>,
431    #[optional] axis: impl Into<Option<i32>>,
432    #[optional] reverse: impl Into<Option<bool>>,
433    #[optional] inclusive: impl Into<Option<bool>>,
434    #[optional] stream: impl AsRef<Stream>,
435) -> Result<Array> {
436    crate::with_stream(stream.as_ref(), || cumsum(a, axis, reverse, inclusive))
437}
438
439#[cfg(test)]
440mod tests {
441    use super::*;
442    use pretty_assertions::assert_eq;
443
444    #[test]
445    fn test_cummax() {
446        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
447
448        let result = array.cummax(0, None, None).unwrap();
449        assert_eq!(result.shape(), &[2, 2]);
450        assert_eq!(result.as_slice::<i32>(), &[5, 8, 5, 9]);
451
452        let result = array.cummax(1, None, None).unwrap();
453        assert_eq!(result.shape(), &[2, 2]);
454        assert_eq!(result.as_slice::<i32>(), &[5, 8, 4, 9]);
455
456        let result = array.cummax(None, None, None).unwrap();
457        assert_eq!(result.shape(), &[4]);
458        assert_eq!(result.as_slice::<i32>(), &[5, 8, 8, 9]);
459
460        let result = array.cummax(0, Some(true), None).unwrap();
461        assert_eq!(result.shape(), &[2, 2]);
462        assert_eq!(result.as_slice::<i32>(), &[5, 9, 4, 9]);
463
464        let result = array.cummax(0, None, Some(true)).unwrap();
465        assert_eq!(result.shape(), &[2, 2]);
466        assert_eq!(result.as_slice::<i32>(), &[5, 8, 5, 9]);
467    }
468
469    #[test]
470    fn test_cummax_out_of_bounds() {
471        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
472        let result = array.cummax(2, None, None);
473        assert!(result.is_err());
474    }
475
476    #[test]
477    fn test_cummin() {
478        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
479
480        let result = array.cummin(0, None, None).unwrap();
481        assert_eq!(result.shape(), &[2, 2]);
482        assert_eq!(result.as_slice::<i32>(), &[5, 8, 4, 8]);
483
484        let result = array.cummin(1, None, None).unwrap();
485        assert_eq!(result.shape(), &[2, 2]);
486        assert_eq!(result.as_slice::<i32>(), &[5, 5, 4, 4]);
487
488        let result = array.cummin(None, None, None).unwrap();
489        assert_eq!(result.shape(), &[4]);
490        assert_eq!(result.as_slice::<i32>(), &[5, 5, 4, 4]);
491
492        let result = array.cummin(0, Some(true), None).unwrap();
493        assert_eq!(result.shape(), &[2, 2]);
494        assert_eq!(result.as_slice::<i32>(), &[4, 8, 4, 9]);
495
496        let result = array.cummin(0, None, Some(true)).unwrap();
497        assert_eq!(result.shape(), &[2, 2]);
498        assert_eq!(result.as_slice::<i32>(), &[5, 8, 4, 8]);
499    }
500
501    #[test]
502    fn test_cummin_out_of_bounds() {
503        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
504        let result = array.cummin(2, None, None);
505        assert!(result.is_err());
506    }
507
508    #[test]
509    fn test_cumprod() {
510        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
511
512        let result = array.cumprod(0, None, None).unwrap();
513        assert_eq!(result.shape(), &[2, 2]);
514        assert_eq!(result.as_slice::<i32>(), &[5, 8, 20, 72]);
515
516        let result = array.cumprod(1, None, None).unwrap();
517        assert_eq!(result.shape(), &[2, 2]);
518        assert_eq!(result.as_slice::<i32>(), &[5, 40, 4, 36]);
519
520        let result = array.cumprod(None, None, None).unwrap();
521        assert_eq!(result.shape(), &[4]);
522        assert_eq!(result.as_slice::<i32>(), &[5, 40, 160, 1440]);
523
524        let result = array.cumprod(0, Some(true), None).unwrap();
525        assert_eq!(result.shape(), &[2, 2]);
526        assert_eq!(result.as_slice::<i32>(), &[20, 72, 4, 9]);
527
528        let result = array.cumprod(0, None, Some(true)).unwrap();
529        assert_eq!(result.shape(), &[2, 2]);
530        assert_eq!(result.as_slice::<i32>(), &[5, 8, 20, 72]);
531    }
532
533    #[test]
534    fn test_cumprod_out_of_bounds() {
535        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
536        let result = array.cumprod(2, None, None);
537        assert!(result.is_err());
538    }
539
540    #[test]
541    fn test_cumsum() {
542        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
543
544        let result = array.cumsum(0, None, None).unwrap();
545        assert_eq!(result.shape(), &[2, 2]);
546        assert_eq!(result.as_slice::<i32>(), &[5, 8, 9, 17]);
547
548        let result = array.cumsum(1, None, None).unwrap();
549        assert_eq!(result.shape(), &[2, 2]);
550        assert_eq!(result.as_slice::<i32>(), &[5, 13, 4, 13]);
551
552        let result = array.cumsum(None, None, None).unwrap();
553        assert_eq!(result.shape(), &[4]);
554        assert_eq!(result.as_slice::<i32>(), &[5, 13, 17, 26]);
555
556        let result = array.cumsum(0, Some(true), None).unwrap();
557        assert_eq!(result.shape(), &[2, 2]);
558        assert_eq!(result.as_slice::<i32>(), &[9, 17, 4, 9]);
559
560        let result = array.cumsum(0, None, Some(true)).unwrap();
561        assert_eq!(result.shape(), &[2, 2]);
562        assert_eq!(result.as_slice::<i32>(), &[5, 8, 9, 17]);
563    }
564
565    #[test]
566    fn test_cumsum_out_of_bounds() {
567        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
568        let result = array.cumsum(2, None, None);
569        assert!(result.is_err());
570    }
571}
mlx_rs/ops/cumulative.rs

mlx_rs/ops/
cumulative.rs
