mlx_rs/ops/
reduction.rs

1use crate::array::Array;
2use crate::error::Result;
3use crate::utils::axes_or_default_to_all;
4use crate::utils::guard::Guarded;
5use crate::{Axes, Stream};
6use mlx_internal_macros::generate_macro;
7
8/// Axis selection and dimension retention for [`Array::count_nonzero`].
9#[derive(Debug, Clone, Default, PartialEq, Eq)]
10pub struct CountNonzeroOptions {
11    /// Axes to reduce.
12    pub axes: Axes,
13
14    /// Keep reduced axes as singleton dimensions.
15    pub keep_dims: bool,
16}
17
18static EMPTY_AXES_DUMMY: [i32; 1] = [0];
19
20impl Array {
21    /// Count values that compare unequal to zero.
22    ///
23    /// The result is the sum of an `i32` nonzero mask. An explicitly empty axis list returns that
24    /// elementwise mask instead of reducing it.
25    ///
26    /// ```rust
27    /// use mlx_rs::{array, ops::CountNonzeroOptions, Axes, Dtype};
28    ///
29    /// let input = array!([[0, 2], [3, 0]]);
30    /// let output = input
31    ///     .count_nonzero(CountNonzeroOptions {
32    ///         axes: Axes::Axis(0),
33    ///         keep_dims: false,
34    ///     })
35    ///     .unwrap();
36    /// assert_eq!(output.dtype(), Dtype::Int32);
37    /// ```
38    pub fn count_nonzero(&self, options: CountNonzeroOptions) -> Result<Array> {
39        let stream = Stream::thread_local_or_default();
40        match options.axes {
41            Axes::All => Array::try_from_op(|res| unsafe {
42                mlx_sys::mlx_count_nonzero(
43                    res,
44                    self.as_ptr(),
45                    options.keep_dims,
46                    stream.as_ref().as_ptr(),
47                )
48            }),
49            Axes::Axis(axis) => Array::try_from_op(|res| unsafe {
50                mlx_sys::mlx_count_nonzero_axis(
51                    res,
52                    self.as_ptr(),
53                    axis,
54                    options.keep_dims,
55                    stream.as_ref().as_ptr(),
56                )
57            }),
58            Axes::Axes(axes) => {
59                let axes_ptr = if axes.is_empty() {
60                    EMPTY_AXES_DUMMY.as_ptr()
61                } else {
62                    axes.as_ptr()
63                };
64                Array::try_from_op(|res| unsafe {
65                    mlx_sys::mlx_count_nonzero_axes(
66                        res,
67                        self.as_ptr(),
68                        axes_ptr,
69                        axes.len(),
70                        options.keep_dims,
71                        stream.as_ref().as_ptr(),
72                    )
73                })
74            }
75        }
76    }
77
78    /// An `and` reduction over the given axes returning an error if the axes are invalid.
79    ///
80    /// # Params
81    ///
82    /// - axes: The axes to reduce over -- defaults to all axes if not provided
83    /// - keep_dims: Whether to keep the reduced dimensions -- defaults to false if not provided
84    ///
85    /// # Example
86    ///
87    /// ```rust
88    /// use mlx_rs::Array;
89    /// let a = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], &[3, 4]);
90    /// let mut b = a.all_axes(&[0], None).unwrap();
91    ///
92    /// let results: &[bool] = b.as_slice();
93    /// // results == [false, true, true, true]
94    /// ```
95    pub fn all_axes(&self, axes: &[i32], keep_dims: impl Into<Option<bool>>) -> Result<Array> {
96        let stream = Stream::thread_local_or_default();
97        Array::try_from_op(|res| unsafe {
98            mlx_sys::mlx_all_axes(
99                res,
100                self.as_ptr(),
101                axes.as_ptr(),
102                axes.len(),
103                keep_dims.into().unwrap_or(false),
104                stream.as_ref().as_ptr(),
105            )
106        })
107    }
108
109    /// Compatibility shim for [`all_axes`].
110    #[deprecated(
111        since = "0.26.0",
112        note = "use `with_stream` or `with_device` around `all_axes`"
113    )]
114    pub fn all_axes_device(
115        &self,
116        axes: &[i32],
117        keep_dims: impl Into<Option<bool>>,
118        stream: impl AsRef<Stream>,
119    ) -> Result<Array> {
120        crate::with_stream(stream.as_ref(), || self.all_axes(axes, keep_dims))
121    }
122
123    /// Similar to [`Array::all_axes`] but only reduces over a single axis.
124    pub fn all_axis(&self, axis: i32, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
125        let stream = Stream::thread_local_or_default();
126        Array::try_from_op(|res| unsafe {
127            mlx_sys::mlx_all_axis(
128                res,
129                self.as_ptr(),
130                axis,
131                keep_dims.into().unwrap_or(false),
132                stream.as_ref().as_ptr(),
133            )
134        })
135    }
136
137    /// Compatibility shim for [`all_axis`].
138    #[deprecated(
139        since = "0.26.0",
140        note = "use `with_stream` or `with_device` around `all_axis`"
141    )]
142    pub fn all_axis_device(
143        &self,
144        axis: i32,
145        keep_dims: impl Into<Option<bool>>,
146        stream: impl AsRef<Stream>,
147    ) -> Result<Array> {
148        crate::with_stream(stream.as_ref(), || self.all_axis(axis, keep_dims))
149    }
150
151    /// Similar to [`Array::all_axes`] but reduces over all axes.
152    pub fn all(&self, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
153        let stream = Stream::thread_local_or_default();
154        Array::try_from_op(|res| unsafe {
155            mlx_sys::mlx_all(
156                res,
157                self.as_ptr(),
158                keep_dims.into().unwrap_or(false),
159                stream.as_ref().as_ptr(),
160            )
161        })
162    }
163
164    /// Compatibility shim for [`all`].
165    #[deprecated(
166        since = "0.26.0",
167        note = "use `with_stream` or `with_device` around `all`"
168    )]
169    pub fn all_device(
170        &self,
171        keep_dims: impl Into<Option<bool>>,
172        stream: impl AsRef<Stream>,
173    ) -> Result<Array> {
174        crate::with_stream(stream.as_ref(), || self.all(keep_dims))
175    }
176
177    /// A `product` reduction over the given axes returning an error if the axes are invalid.
178    ///
179    /// # Params
180    ///
181    /// - axes: axes to reduce over
182    /// - keep_dims: Whether to keep the reduced dimensions -- defaults to false if not provided
183    ///
184    /// # Example
185    ///
186    /// ```rust
187    /// use mlx_rs::Array;
188    /// let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
189    ///
190    /// // result is [20, 72]
191    /// let result = array.prod_axes(&[0], None).unwrap();
192    /// ```
193    pub fn prod_axes(&self, axes: &[i32], keep_dims: impl Into<Option<bool>>) -> Result<Array> {
194        let stream = Stream::thread_local_or_default();
195        Array::try_from_op(|res| unsafe {
196            mlx_sys::mlx_prod_axes(
197                res,
198                self.as_ptr(),
199                axes.as_ptr(),
200                axes.len(),
201                keep_dims.into().unwrap_or(false),
202                stream.as_ref().as_ptr(),
203            )
204        })
205    }
206
207    /// Compatibility shim for [`prod_axes`].
208    #[deprecated(
209        since = "0.26.0",
210        note = "use `with_stream` or `with_device` around `prod_axes`"
211    )]
212    pub fn prod_axes_device(
213        &self,
214        axes: &[i32],
215        keep_dims: impl Into<Option<bool>>,
216        stream: impl AsRef<Stream>,
217    ) -> Result<Array> {
218        crate::with_stream(stream.as_ref(), || self.prod_axes(axes, keep_dims))
219    }
220
221    /// Similar to [`Array::prod_axes`] but only reduces over a single axis.
222    pub fn prod_axis(&self, axis: i32, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
223        let stream = Stream::thread_local_or_default();
224        Array::try_from_op(|res| unsafe {
225            mlx_sys::mlx_prod_axis(
226                res,
227                self.as_ptr(),
228                axis,
229                keep_dims.into().unwrap_or(false),
230                stream.as_ref().as_ptr(),
231            )
232        })
233    }
234
235    /// Compatibility shim for [`prod_axis`].
236    #[deprecated(
237        since = "0.26.0",
238        note = "use `with_stream` or `with_device` around `prod_axis`"
239    )]
240    pub fn prod_axis_device(
241        &self,
242        axis: i32,
243        keep_dims: impl Into<Option<bool>>,
244        stream: impl AsRef<Stream>,
245    ) -> Result<Array> {
246        crate::with_stream(stream.as_ref(), || self.prod_axis(axis, keep_dims))
247    }
248
249    /// Similar to [`Array::prod_axes`] but reduces over all axes.
250    pub fn prod(&self, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
251        let stream = Stream::thread_local_or_default();
252        Array::try_from_op(|res| unsafe {
253            mlx_sys::mlx_prod(
254                res,
255                self.as_ptr(),
256                keep_dims.into().unwrap_or(false),
257                stream.as_ref().as_ptr(),
258            )
259        })
260    }
261
262    /// Compatibility shim for [`prod`].
263    #[deprecated(
264        since = "0.26.0",
265        note = "use `with_stream` or `with_device` around `prod`"
266    )]
267    pub fn prod_device(
268        &self,
269        keep_dims: impl Into<Option<bool>>,
270        stream: impl AsRef<Stream>,
271    ) -> Result<Array> {
272        crate::with_stream(stream.as_ref(), || self.prod(keep_dims))
273    }
274
275    /// A `max` reduction over the given axes returning an error if the axes are invalid.
276    ///
277    /// # Params
278    ///
279    /// - axes: axes to reduce over
280    /// - keep_dims: Whether to keep the reduced dimensions -- defaults to false if not provided
281    ///
282    /// # Example
283    ///
284    /// ```rust
285    /// use mlx_rs::Array;
286    /// let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
287    ///
288    /// // result is [5, 9]
289    /// let result = array.max_axes(&[0], None).unwrap();
290    /// ```
291    pub fn max_axes(&self, axes: &[i32], keep_dims: impl Into<Option<bool>>) -> Result<Array> {
292        let stream = Stream::thread_local_or_default();
293        Array::try_from_op(|res| unsafe {
294            mlx_sys::mlx_max_axes(
295                res,
296                self.as_ptr(),
297                axes.as_ptr(),
298                axes.len(),
299                keep_dims.into().unwrap_or(false),
300                stream.as_ref().as_ptr(),
301            )
302        })
303    }
304
305    /// Compatibility shim for [`max_axes`].
306    #[deprecated(
307        since = "0.26.0",
308        note = "use `with_stream` or `with_device` around `max_axes`"
309    )]
310    pub fn max_axes_device(
311        &self,
312        axes: &[i32],
313        keep_dims: impl Into<Option<bool>>,
314        stream: impl AsRef<Stream>,
315    ) -> Result<Array> {
316        crate::with_stream(stream.as_ref(), || self.max_axes(axes, keep_dims))
317    }
318
319    /// Similar to [`Array::max_axes`] but only reduces over a single axis.
320    pub fn max_axis(&self, axis: i32, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
321        let stream = Stream::thread_local_or_default();
322        Array::try_from_op(|res| unsafe {
323            mlx_sys::mlx_max_axis(
324                res,
325                self.as_ptr(),
326                axis,
327                keep_dims.into().unwrap_or(false),
328                stream.as_ref().as_ptr(),
329            )
330        })
331    }
332
333    /// Compatibility shim for [`max_axis`].
334    #[deprecated(
335        since = "0.26.0",
336        note = "use `with_stream` or `with_device` around `max_axis`"
337    )]
338    pub fn max_axis_device(
339        &self,
340        axis: i32,
341        keep_dims: impl Into<Option<bool>>,
342        stream: impl AsRef<Stream>,
343    ) -> Result<Array> {
344        crate::with_stream(stream.as_ref(), || self.max_axis(axis, keep_dims))
345    }
346
347    /// Similar to [`Array::max_axes`] but reduces over all axes.
348    pub fn max(&self, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
349        let stream = Stream::thread_local_or_default();
350        Array::try_from_op(|res| unsafe {
351            mlx_sys::mlx_max(
352                res,
353                self.as_ptr(),
354                keep_dims.into().unwrap_or(false),
355                stream.as_ref().as_ptr(),
356            )
357        })
358    }
359
360    /// Compatibility shim for [`max`].
361    #[deprecated(
362        since = "0.26.0",
363        note = "use `with_stream` or `with_device` around `max`"
364    )]
365    pub fn max_device(
366        &self,
367        keep_dims: impl Into<Option<bool>>,
368        stream: impl AsRef<Stream>,
369    ) -> Result<Array> {
370        crate::with_stream(stream.as_ref(), || self.max(keep_dims))
371    }
372
373    /// Sum reduce the array over the given axes returning an error if the axes are invalid.
374    ///
375    /// # Params
376    ///
377    /// - axes: axes to reduce over
378    /// - keep_dims: if `true`, keep the reduces axes as singleton dimensions
379    ///
380    /// # Example
381    ///
382    /// ```rust
383    /// use mlx_rs::Array;
384    /// let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
385    ///
386    /// // result is [9, 17]
387    /// let result = array.sum_axes(&[0], None).unwrap();
388    /// ```
389    pub fn sum_axes(&self, axes: &[i32], keep_dims: impl Into<Option<bool>>) -> Result<Array> {
390        let stream = Stream::thread_local_or_default();
391        Array::try_from_op(|res| unsafe {
392            mlx_sys::mlx_sum_axes(
393                res,
394                self.as_ptr(),
395                axes.as_ptr(),
396                axes.len(),
397                keep_dims.into().unwrap_or(false),
398                stream.as_ref().as_ptr(),
399            )
400        })
401    }
402
403    /// Compatibility shim for [`sum_axes`].
404    #[deprecated(
405        since = "0.26.0",
406        note = "use `with_stream` or `with_device` around `sum_axes`"
407    )]
408    pub fn sum_axes_device(
409        &self,
410        axes: &[i32],
411        keep_dims: impl Into<Option<bool>>,
412        stream: impl AsRef<Stream>,
413    ) -> Result<Array> {
414        crate::with_stream(stream.as_ref(), || self.sum_axes(axes, keep_dims))
415    }
416
417    /// Similar to [`Array::sum_axes`] but only reduces over a single axis.
418    pub fn sum_axis(&self, axis: i32, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
419        let stream = Stream::thread_local_or_default();
420        Array::try_from_op(|res| unsafe {
421            mlx_sys::mlx_sum_axis(
422                res,
423                self.as_ptr(),
424                axis,
425                keep_dims.into().unwrap_or(false),
426                stream.as_ref().as_ptr(),
427            )
428        })
429    }
430
431    /// Compatibility shim for [`sum_axis`].
432    #[deprecated(
433        since = "0.26.0",
434        note = "use `with_stream` or `with_device` around `sum_axis`"
435    )]
436    pub fn sum_axis_device(
437        &self,
438        axis: i32,
439        keep_dims: impl Into<Option<bool>>,
440        stream: impl AsRef<Stream>,
441    ) -> Result<Array> {
442        crate::with_stream(stream.as_ref(), || self.sum_axis(axis, keep_dims))
443    }
444
445    /// Similar to [`Array::sum_axes`] but reduces over all axes.
446    pub fn sum(&self, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
447        let stream = Stream::thread_local_or_default();
448        Array::try_from_op(|res| unsafe {
449            mlx_sys::mlx_sum(
450                res,
451                self.as_ptr(),
452                keep_dims.into().unwrap_or(false),
453                stream.as_ref().as_ptr(),
454            )
455        })
456    }
457
458    /// Compatibility shim for [`sum`].
459    #[deprecated(
460        since = "0.26.0",
461        note = "use `with_stream` or `with_device` around `sum`"
462    )]
463    pub fn sum_device(
464        &self,
465        keep_dims: impl Into<Option<bool>>,
466        stream: impl AsRef<Stream>,
467    ) -> Result<Array> {
468        crate::with_stream(stream.as_ref(), || self.sum(keep_dims))
469    }
470
471    /// A `mean` reduction over the given axes returning an error if the axes are invalid.
472    ///
473    /// # Params
474    ///
475    /// - axes: axes to reduce over
476    /// - keep_dims: Whether to keep the reduced dimensions -- defaults to false if not provided
477    ///
478    /// # Example
479    ///
480    /// ```rust
481    /// use mlx_rs::Array;
482    /// let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
483    ///
484    /// // result is [4.5, 8.5]
485    /// let result = array.mean_axes(&[0], None).unwrap();
486    /// ```
487    pub fn mean_axes(&self, axes: &[i32], keep_dims: impl Into<Option<bool>>) -> Result<Array> {
488        let stream = Stream::thread_local_or_default();
489        let axes = axes_or_default_to_all(axes, self.ndim() as i32);
490        Array::try_from_op(|res| unsafe {
491            mlx_sys::mlx_mean_axes(
492                res,
493                self.as_ptr(),
494                axes.as_ptr(),
495                axes.len(),
496                keep_dims.into().unwrap_or(false),
497                stream.as_ref().as_ptr(),
498            )
499        })
500    }
501
502    /// Compatibility shim for [`mean_axes`].
503    #[deprecated(
504        since = "0.26.0",
505        note = "use `with_stream` or `with_device` around `mean_axes`"
506    )]
507    pub fn mean_axes_device(
508        &self,
509        axes: &[i32],
510        keep_dims: impl Into<Option<bool>>,
511        stream: impl AsRef<Stream>,
512    ) -> Result<Array> {
513        crate::with_stream(stream.as_ref(), || self.mean_axes(axes, keep_dims))
514    }
515
516    /// Similar to [`Array::mean_axes`] but only reduces over a single axis.
517    pub fn mean_axis(&self, axis: i32, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
518        let stream = Stream::thread_local_or_default();
519        Array::try_from_op(|res| unsafe {
520            mlx_sys::mlx_mean_axis(
521                res,
522                self.as_ptr(),
523                axis,
524                keep_dims.into().unwrap_or(false),
525                stream.as_ref().as_ptr(),
526            )
527        })
528    }
529
530    /// Compatibility shim for [`mean_axis`].
531    #[deprecated(
532        since = "0.26.0",
533        note = "use `with_stream` or `with_device` around `mean_axis`"
534    )]
535    pub fn mean_axis_device(
536        &self,
537        axis: i32,
538        keep_dims: impl Into<Option<bool>>,
539        stream: impl AsRef<Stream>,
540    ) -> Result<Array> {
541        crate::with_stream(stream.as_ref(), || self.mean_axis(axis, keep_dims))
542    }
543
544    /// Similar to [`Array::mean_axes`] but reduces over all axes.
545    pub fn mean(&self, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
546        let stream = Stream::thread_local_or_default();
547        Array::try_from_op(|res| unsafe {
548            mlx_sys::mlx_mean(
549                res,
550                self.as_ptr(),
551                keep_dims.into().unwrap_or(false),
552                stream.as_ref().as_ptr(),
553            )
554        })
555    }
556
557    /// Compatibility shim for [`mean`].
558    #[deprecated(
559        since = "0.26.0",
560        note = "use `with_stream` or `with_device` around `mean`"
561    )]
562    pub fn mean_device(
563        &self,
564        keep_dims: impl Into<Option<bool>>,
565        stream: impl AsRef<Stream>,
566    ) -> Result<Array> {
567        crate::with_stream(stream.as_ref(), || self.mean(keep_dims))
568    }
569
570    /// A `min` reduction over the given axes returning an error if the axes are invalid.
571    ///
572    /// # Params
573    ///
574    /// - axes: axes to reduce over
575    /// - keep_dims: Whether to keep the reduced dimensions -- defaults to false if not provided
576    ///
577    /// # Example
578    ///
579    /// ```rust
580    /// use mlx_rs::Array;
581    /// let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
582    ///
583    /// // result is [4, 8]
584    /// let result = array.min_axes(&[0], None).unwrap();
585    /// ```
586    pub fn min_axes(&self, axes: &[i32], keep_dims: impl Into<Option<bool>>) -> Result<Array> {
587        let stream = Stream::thread_local_or_default();
588        Array::try_from_op(|res| unsafe {
589            mlx_sys::mlx_min_axes(
590                res,
591                self.as_ptr(),
592                axes.as_ptr(),
593                axes.len(),
594                keep_dims.into().unwrap_or(false),
595                stream.as_ref().as_ptr(),
596            )
597        })
598    }
599
600    /// Compatibility shim for [`min_axes`].
601    #[deprecated(
602        since = "0.26.0",
603        note = "use `with_stream` or `with_device` around `min_axes`"
604    )]
605    pub fn min_axes_device(
606        &self,
607        axes: &[i32],
608        keep_dims: impl Into<Option<bool>>,
609        stream: impl AsRef<Stream>,
610    ) -> Result<Array> {
611        crate::with_stream(stream.as_ref(), || self.min_axes(axes, keep_dims))
612    }
613
614    /// Similar to [`Array::min_axes`] but only reduces over a single axis.
615    pub fn min_axis(&self, axis: i32, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
616        let stream = Stream::thread_local_or_default();
617        Array::try_from_op(|res| unsafe {
618            mlx_sys::mlx_min_axis(
619                res,
620                self.as_ptr(),
621                axis,
622                keep_dims.into().unwrap_or(false),
623                stream.as_ref().as_ptr(),
624            )
625        })
626    }
627
628    /// Compatibility shim for [`min_axis`].
629    #[deprecated(
630        since = "0.26.0",
631        note = "use `with_stream` or `with_device` around `min_axis`"
632    )]
633    pub fn min_axis_device(
634        &self,
635        axis: i32,
636        keep_dims: impl Into<Option<bool>>,
637        stream: impl AsRef<Stream>,
638    ) -> Result<Array> {
639        crate::with_stream(stream.as_ref(), || self.min_axis(axis, keep_dims))
640    }
641
642    /// Similar to [`Array::min_axes`] but reduces over all axes.
643    pub fn min(&self, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
644        let stream = Stream::thread_local_or_default();
645        Array::try_from_op(|res| unsafe {
646            mlx_sys::mlx_min(
647                res,
648                self.as_ptr(),
649                keep_dims.into().unwrap_or(false),
650                stream.as_ref().as_ptr(),
651            )
652        })
653    }
654
655    /// Compatibility shim for [`min`].
656    #[deprecated(
657        since = "0.26.0",
658        note = "use `with_stream` or `with_device` around `min`"
659    )]
660    pub fn min_device(
661        &self,
662        keep_dims: impl Into<Option<bool>>,
663        stream: impl AsRef<Stream>,
664    ) -> Result<Array> {
665        crate::with_stream(stream.as_ref(), || self.min(keep_dims))
666    }
667
668    /// Compute the variance(s) over the given axes returning an error if the axes are invalid.
669    ///
670    /// # Params
671    ///
672    /// - axes: axes to reduce over
673    /// - keep_dims: if `true`, keep the reduces axes as singleton dimensions
674    /// - ddof: the divisor to compute the variance is `N - ddof`
675    pub fn var_axes(
676        &self,
677        axes: &[i32],
678        keep_dims: impl Into<Option<bool>>,
679        ddof: impl Into<Option<i32>>,
680    ) -> Result<Array> {
681        let stream = Stream::thread_local_or_default();
682        Array::try_from_op(|res| unsafe {
683            mlx_sys::mlx_var_axes(
684                res,
685                self.as_ptr(),
686                axes.as_ptr(),
687                axes.len(),
688                keep_dims.into().unwrap_or(false),
689                ddof.into().unwrap_or(0),
690                stream.as_ref().as_ptr(),
691            )
692        })
693    }
694
695    /// Compatibility shim for [`var_axes`].
696    #[deprecated(
697        since = "0.26.0",
698        note = "use `with_stream` or `with_device` around `var_axes`"
699    )]
700    pub fn var_axes_device(
701        &self,
702        axes: &[i32],
703        keep_dims: impl Into<Option<bool>>,
704        ddof: impl Into<Option<i32>>,
705        stream: impl AsRef<Stream>,
706    ) -> Result<Array> {
707        crate::with_stream(stream.as_ref(), || self.var_axes(axes, keep_dims, ddof))
708    }
709
710    /// Similar to [`Array::var_axes`] but only reduces over a single axis.
711    pub fn var_axis(
712        &self,
713        axis: i32,
714        keep_dims: impl Into<Option<bool>>,
715        ddof: impl Into<Option<i32>>,
716    ) -> Result<Array> {
717        let stream = Stream::thread_local_or_default();
718        Array::try_from_op(|res| unsafe {
719            mlx_sys::mlx_var_axis(
720                res,
721                self.as_ptr(),
722                axis,
723                keep_dims.into().unwrap_or(false),
724                ddof.into().unwrap_or(0),
725                stream.as_ref().as_ptr(),
726            )
727        })
728    }
729
730    /// Compatibility shim for [`var_axis`].
731    #[deprecated(
732        since = "0.26.0",
733        note = "use `with_stream` or `with_device` around `var_axis`"
734    )]
735    pub fn var_axis_device(
736        &self,
737        axis: i32,
738        keep_dims: impl Into<Option<bool>>,
739        ddof: impl Into<Option<i32>>,
740        stream: impl AsRef<Stream>,
741    ) -> Result<Array> {
742        crate::with_stream(stream.as_ref(), || self.var_axis(axis, keep_dims, ddof))
743    }
744
745    /// Similar to [`Array::var_axes`] but reduces over all axes.
746    pub fn var(
747        &self,
748        keep_dims: impl Into<Option<bool>>,
749        ddof: impl Into<Option<i32>>,
750    ) -> Result<Array> {
751        let stream = Stream::thread_local_or_default();
752        Array::try_from_op(|res| unsafe {
753            mlx_sys::mlx_var(
754                res,
755                self.as_ptr(),
756                keep_dims.into().unwrap_or(false),
757                ddof.into().unwrap_or(0),
758                stream.as_ref().as_ptr(),
759            )
760        })
761    }
762
763    /// Compatibility shim for [`var`].
764    #[deprecated(
765        since = "0.26.0",
766        note = "use `with_stream` or `with_device` around `var`"
767    )]
768    pub fn var_device(
769        &self,
770        keep_dims: impl Into<Option<bool>>,
771        ddof: impl Into<Option<i32>>,
772        stream: impl AsRef<Stream>,
773    ) -> Result<Array> {
774        crate::with_stream(stream.as_ref(), || self.var(keep_dims, ddof))
775    }
776
777    /// Compute the median over the given axes.
778    ///
779    /// # Params
780    ///
781    /// - axes: axes to reduce over
782    /// - keep_dims: Whether to keep the reduced dimensions -- defaults to false if not provided
783    pub fn median_axes(&self, axes: &[i32], keep_dims: impl Into<Option<bool>>) -> Result<Array> {
784        let stream = Stream::thread_local_or_default();
785        Array::try_from_op(|res| unsafe {
786            mlx_sys::mlx_median_axes(
787                res,
788                self.as_ptr(),
789                axes.as_ptr(),
790                axes.len(),
791                keep_dims.into().unwrap_or(false),
792                stream.as_ref().as_ptr(),
793            )
794        })
795    }
796
797    /// Compatibility shim for [`median_axes`].
798    #[deprecated(
799        since = "0.26.0",
800        note = "use `with_stream` or `with_device` around `median_axes`"
801    )]
802    pub fn median_axes_device(
803        &self,
804        axes: &[i32],
805        keep_dims: impl Into<Option<bool>>,
806        stream: impl AsRef<Stream>,
807    ) -> Result<Array> {
808        crate::with_stream(stream.as_ref(), || self.median_axes(axes, keep_dims))
809    }
810
811    /// Similar to [`Array::median_axes`] but only reduces over a single axis.
812    pub fn median_axis(&self, axis: i32, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
813        let stream = Stream::thread_local_or_default();
814        Array::try_from_op(|res| unsafe {
815            mlx_sys::mlx_median_axis(
816                res,
817                self.as_ptr(),
818                axis,
819                keep_dims.into().unwrap_or(false),
820                stream.as_ref().as_ptr(),
821            )
822        })
823    }
824
825    /// Compatibility shim for [`median_axis`].
826    #[deprecated(
827        since = "0.26.0",
828        note = "use `with_stream` or `with_device` around `median_axis`"
829    )]
830    pub fn median_axis_device(
831        &self,
832        axis: i32,
833        keep_dims: impl Into<Option<bool>>,
834        stream: impl AsRef<Stream>,
835    ) -> Result<Array> {
836        crate::with_stream(stream.as_ref(), || self.median_axis(axis, keep_dims))
837    }
838
839    /// Compute the median over all axes.
840    pub fn median(&self, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
841        let stream = Stream::thread_local_or_default();
842        Array::try_from_op(|res| unsafe {
843            mlx_sys::mlx_median(
844                res,
845                self.as_ptr(),
846                keep_dims.into().unwrap_or(false),
847                stream.as_ref().as_ptr(),
848            )
849        })
850    }
851
852    /// Compatibility shim for [`median`].
853    #[deprecated(
854        since = "0.26.0",
855        note = "use `with_stream` or `with_device` around `median`"
856    )]
857    pub fn median_device(
858        &self,
859        keep_dims: impl Into<Option<bool>>,
860        stream: impl AsRef<Stream>,
861    ) -> Result<Array> {
862        crate::with_stream(stream.as_ref(), || self.median(keep_dims))
863    }
864
865    /// A `log-sum-exp` reduction over the given axes returning an error if the axes are invalid.
866    ///
867    /// The log-sum-exp reduction is a numerically stable version of using the individual operations.
868    ///
869    /// # Params
870    ///
871    /// - axes: axes to reduce over
872    /// - keep_dims: Whether to keep the reduced dimensions -- defaults to false if not provided
873    pub fn logsumexp_axes(
874        &self,
875        axes: &[i32],
876        keep_dims: impl Into<Option<bool>>,
877    ) -> Result<Array> {
878        let stream = Stream::thread_local_or_default();
879        Array::try_from_op(|res| unsafe {
880            mlx_sys::mlx_logsumexp_axes(
881                res,
882                self.as_ptr(),
883                axes.as_ptr(),
884                axes.len(),
885                keep_dims.into().unwrap_or(false),
886                stream.as_ref().as_ptr(),
887            )
888        })
889    }
890
891    /// Compatibility shim for [`logsumexp_axes`].
892    #[deprecated(
893        since = "0.26.0",
894        note = "use `with_stream` or `with_device` around `logsumexp_axes`"
895    )]
896    pub fn logsumexp_axes_device(
897        &self,
898        axes: &[i32],
899        keep_dims: impl Into<Option<bool>>,
900        stream: impl AsRef<Stream>,
901    ) -> Result<Array> {
902        crate::with_stream(stream.as_ref(), || self.logsumexp_axes(axes, keep_dims))
903    }
904
905    /// Similar to [`Array::logsumexp_axes`] but only reduces over a single axis.
906    pub fn logsumexp_axis(&self, axis: i32, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
907        let stream = Stream::thread_local_or_default();
908        Array::try_from_op(|res| unsafe {
909            mlx_sys::mlx_logsumexp_axis(
910                res,
911                self.as_ptr(),
912                axis,
913                keep_dims.into().unwrap_or(false),
914                stream.as_ref().as_ptr(),
915            )
916        })
917    }
918
919    /// Compatibility shim for [`logsumexp_axis`].
920    #[deprecated(
921        since = "0.26.0",
922        note = "use `with_stream` or `with_device` around `logsumexp_axis`"
923    )]
924    pub fn logsumexp_axis_device(
925        &self,
926        axis: i32,
927        keep_dims: impl Into<Option<bool>>,
928        stream: impl AsRef<Stream>,
929    ) -> Result<Array> {
930        crate::with_stream(stream.as_ref(), || self.logsumexp_axis(axis, keep_dims))
931    }
932
933    /// Similar to [`Array::logsumexp_axes`] but reduces over all axes.
934    pub fn logsumexp(&self, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
935        let stream = Stream::thread_local_or_default();
936        Array::try_from_op(|res| unsafe {
937            mlx_sys::mlx_logsumexp(
938                res,
939                self.as_ptr(),
940                keep_dims.into().unwrap_or(false),
941                stream.as_ref().as_ptr(),
942            )
943        })
944    }
945
946    /// Compatibility shim for [`logsumexp`].
947    #[deprecated(
948        since = "0.26.0",
949        note = "use `with_stream` or `with_device` around `logsumexp`"
950    )]
951    pub fn logsumexp_device(
952        &self,
953        keep_dims: impl Into<Option<bool>>,
954        stream: impl AsRef<Stream>,
955    ) -> Result<Array> {
956        crate::with_stream(stream.as_ref(), || self.logsumexp(keep_dims))
957    }
958}
959
960/// See [`Array::all_axes`]
961pub fn all_axes(
962    array: impl AsRef<Array>,
963    axes: &[i32],
964    keep_dims: impl Into<Option<bool>>,
965) -> Result<Array> {
966    array.as_ref().all_axes(axes, keep_dims)
967}
968
969/// Compatibility shim for [`all_axes`].
970#[generate_macro(customize(forwarding_shim = true))]
971#[deprecated(
972    since = "0.26.0",
973    note = "use `with_stream` or `with_device` around `all_axes`"
974)]
975pub fn all_axes_device(
976    array: impl AsRef<Array>,
977    axes: &[i32],
978    #[optional] keep_dims: impl Into<Option<bool>>,
979    #[optional] stream: impl AsRef<Stream>,
980) -> Result<Array> {
981    crate::with_stream(stream.as_ref(), || all_axes(array, axes, keep_dims))
982}
983
984/// See [`Array::all_axis`]
985pub fn all_axis(
986    array: impl AsRef<Array>,
987    axis: i32,
988    keep_dims: impl Into<Option<bool>>,
989) -> Result<Array> {
990    array.as_ref().all_axis(axis, keep_dims)
991}
992
993/// Compatibility shim for [`all_axis`].
994#[generate_macro(customize(forwarding_shim = true))]
995#[deprecated(
996    since = "0.26.0",
997    note = "use `with_stream` or `with_device` around `all_axis`"
998)]
999pub fn all_axis_device(
1000    array: impl AsRef<Array>,
1001    axis: i32,
1002    #[optional] keep_dims: impl Into<Option<bool>>,
1003    #[optional] stream: impl AsRef<Stream>,
1004) -> Result<Array> {
1005    crate::with_stream(stream.as_ref(), || all_axis(array, axis, keep_dims))
1006}
1007
1008/// See [`Array::all`]
1009pub fn all(array: impl AsRef<Array>, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
1010    array.as_ref().all(keep_dims)
1011}
1012
1013/// Compatibility shim for [`all`].
1014#[generate_macro(customize(forwarding_shim = true))]
1015#[deprecated(
1016    since = "0.26.0",
1017    note = "use `with_stream` or `with_device` around `all`"
1018)]
1019pub fn all_device(
1020    array: impl AsRef<Array>,
1021    #[optional] keep_dims: impl Into<Option<bool>>,
1022    #[optional] stream: impl AsRef<Stream>,
1023) -> Result<Array> {
1024    crate::with_stream(stream.as_ref(), || all(array, keep_dims))
1025}
1026
1027/// See [`Array::prod_axes`]
1028pub fn prod_axes(
1029    array: impl AsRef<Array>,
1030    axes: &[i32],
1031    keep_dims: impl Into<Option<bool>>,
1032) -> Result<Array> {
1033    array.as_ref().prod_axes(axes, keep_dims)
1034}
1035
1036/// Compatibility shim for [`prod_axes`].
1037#[generate_macro(customize(forwarding_shim = true))]
1038#[deprecated(
1039    since = "0.26.0",
1040    note = "use `with_stream` or `with_device` around `prod_axes`"
1041)]
1042pub fn prod_axes_device(
1043    array: impl AsRef<Array>,
1044    axes: &[i32],
1045    #[optional] keep_dims: impl Into<Option<bool>>,
1046    #[optional] stream: impl AsRef<Stream>,
1047) -> Result<Array> {
1048    crate::with_stream(stream.as_ref(), || prod_axes(array, axes, keep_dims))
1049}
1050
1051/// See [`Array::prod_axis`]
1052pub fn prod_axis(
1053    array: impl AsRef<Array>,
1054    axis: i32,
1055    keep_dims: impl Into<Option<bool>>,
1056) -> Result<Array> {
1057    array.as_ref().prod_axis(axis, keep_dims)
1058}
1059
1060/// Compatibility shim for [`prod_axis`].
1061#[generate_macro(customize(forwarding_shim = true))]
1062#[deprecated(
1063    since = "0.26.0",
1064    note = "use `with_stream` or `with_device` around `prod_axis`"
1065)]
1066pub fn prod_axis_device(
1067    array: impl AsRef<Array>,
1068    axis: i32,
1069    #[optional] keep_dims: impl Into<Option<bool>>,
1070    #[optional] stream: impl AsRef<Stream>,
1071) -> Result<Array> {
1072    crate::with_stream(stream.as_ref(), || prod_axis(array, axis, keep_dims))
1073}
1074
1075/// See [`Array::prod`]
1076pub fn prod(array: impl AsRef<Array>, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
1077    array.as_ref().prod(keep_dims)
1078}
1079
1080/// Compatibility shim for [`prod`].
1081#[generate_macro(customize(forwarding_shim = true))]
1082#[deprecated(
1083    since = "0.26.0",
1084    note = "use `with_stream` or `with_device` around `prod`"
1085)]
1086pub fn prod_device(
1087    array: impl AsRef<Array>,
1088    #[optional] keep_dims: impl Into<Option<bool>>,
1089    #[optional] stream: impl AsRef<Stream>,
1090) -> Result<Array> {
1091    crate::with_stream(stream.as_ref(), || prod(array, keep_dims))
1092}
1093
1094/// See [`Array::max_axes`]
1095pub fn max_axes(
1096    array: impl AsRef<Array>,
1097    axes: &[i32],
1098    keep_dims: impl Into<Option<bool>>,
1099) -> Result<Array> {
1100    array.as_ref().max_axes(axes, keep_dims)
1101}
1102
1103/// Compatibility shim for [`max_axes`].
1104#[generate_macro(customize(forwarding_shim = true))]
1105#[deprecated(
1106    since = "0.26.0",
1107    note = "use `with_stream` or `with_device` around `max_axes`"
1108)]
1109pub fn max_axes_device(
1110    array: impl AsRef<Array>,
1111    axes: &[i32],
1112    #[optional] keep_dims: impl Into<Option<bool>>,
1113    #[optional] stream: impl AsRef<Stream>,
1114) -> Result<Array> {
1115    crate::with_stream(stream.as_ref(), || max_axes(array, axes, keep_dims))
1116}
1117
1118/// See [`Array::max_axis`]
1119pub fn max_axis(
1120    array: impl AsRef<Array>,
1121    axis: i32,
1122    keep_dims: impl Into<Option<bool>>,
1123) -> Result<Array> {
1124    array.as_ref().max_axis(axis, keep_dims)
1125}
1126
1127/// Compatibility shim for [`max_axis`].
1128#[generate_macro(customize(forwarding_shim = true))]
1129#[deprecated(
1130    since = "0.26.0",
1131    note = "use `with_stream` or `with_device` around `max_axis`"
1132)]
1133pub fn max_axis_device(
1134    array: impl AsRef<Array>,
1135    axis: i32,
1136    #[optional] keep_dims: impl Into<Option<bool>>,
1137    #[optional] stream: impl AsRef<Stream>,
1138) -> Result<Array> {
1139    crate::with_stream(stream.as_ref(), || max_axis(array, axis, keep_dims))
1140}
1141
1142/// See [`Array::max`]
1143pub fn max(array: impl AsRef<Array>, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
1144    array.as_ref().max(keep_dims)
1145}
1146
1147/// Compatibility shim for [`max`].
1148#[generate_macro(customize(forwarding_shim = true))]
1149#[deprecated(
1150    since = "0.26.0",
1151    note = "use `with_stream` or `with_device` around `max`"
1152)]
1153pub fn max_device(
1154    array: impl AsRef<Array>,
1155    #[optional] keep_dims: impl Into<Option<bool>>,
1156    #[optional] stream: impl AsRef<Stream>,
1157) -> Result<Array> {
1158    crate::with_stream(stream.as_ref(), || max(array, keep_dims))
1159}
1160
1161/// Compute the standard deviation(s) over the given axes.
1162///
1163/// # Params
1164///
1165/// - `a`: Input array
1166/// - `axes`: Optional axis or axes to reduce over. If unspecified this defaults to reducing over
1167///   the entire array.
1168/// - `keep_dims`: Keep reduced axes as singleton dimensions, defaults to False.
1169/// - `ddof`: The divisor to compute the variance is `N - ddof`, defaults to `0`.
1170pub fn std_axes(
1171    a: impl AsRef<Array>,
1172    axes: &[i32],
1173    keep_dims: impl Into<Option<bool>>,
1174    ddof: impl Into<Option<i32>>,
1175) -> Result<Array> {
1176    let stream = Stream::thread_local_or_default();
1177    let a = a.as_ref();
1178    let keep_dims = keep_dims.into().unwrap_or(false);
1179    let ddof = ddof.into().unwrap_or(0);
1180    Array::try_from_op(|res| unsafe {
1181        mlx_sys::mlx_std_axes(
1182            res,
1183            a.as_ptr(),
1184            axes.as_ptr(),
1185            axes.len(),
1186            keep_dims,
1187            ddof,
1188            stream.as_ref().as_ptr(),
1189        )
1190    })
1191}
1192
1193/// Compatibility shim for [`std_axes`].
1194#[generate_macro(customize(forwarding_shim = true))]
1195#[deprecated(
1196    since = "0.26.0",
1197    note = "use `with_stream` or `with_device` around `std_axes`"
1198)]
1199pub fn std_axes_device(
1200    a: impl AsRef<Array>,
1201    axes: &[i32],
1202    #[optional] keep_dims: impl Into<Option<bool>>,
1203    #[optional] ddof: impl Into<Option<i32>>,
1204    #[optional] stream: impl AsRef<Stream>,
1205) -> Result<Array> {
1206    crate::with_stream(stream.as_ref(), || std_axes(a, axes, keep_dims, ddof))
1207}
1208
1209/// Similar to [`std_axes`] but only reduces over a single axis.
1210pub fn std_axis(
1211    a: impl AsRef<Array>,
1212    axis: i32,
1213    keep_dims: impl Into<Option<bool>>,
1214    ddof: impl Into<Option<i32>>,
1215) -> Result<Array> {
1216    let stream = Stream::thread_local_or_default();
1217    let a = a.as_ref();
1218    let keep_dims = keep_dims.into().unwrap_or(false);
1219    let ddof = ddof.into().unwrap_or(0);
1220    Array::try_from_op(|res| unsafe {
1221        mlx_sys::mlx_std_axis(
1222            res,
1223            a.as_ptr(),
1224            axis,
1225            keep_dims,
1226            ddof,
1227            stream.as_ref().as_ptr(),
1228        )
1229    })
1230}
1231
1232/// Compatibility shim for [`std_axis`].
1233#[generate_macro(customize(forwarding_shim = true))]
1234#[deprecated(
1235    since = "0.26.0",
1236    note = "use `with_stream` or `with_device` around `std_axis`"
1237)]
1238pub fn std_axis_device(
1239    a: impl AsRef<Array>,
1240    axis: i32,
1241    #[optional] keep_dims: impl Into<Option<bool>>,
1242    #[optional] ddof: impl Into<Option<i32>>,
1243    #[optional] stream: impl AsRef<Stream>,
1244) -> Result<Array> {
1245    crate::with_stream(stream.as_ref(), || std_axis(a, axis, keep_dims, ddof))
1246}
1247
1248/// Similar to [`std_axes`] but reduces over all axes.
1249pub fn std(
1250    a: impl AsRef<Array>,
1251    keep_dims: impl Into<Option<bool>>,
1252    ddof: impl Into<Option<i32>>,
1253) -> Result<Array> {
1254    let stream = Stream::thread_local_or_default();
1255    let a = a.as_ref();
1256    let keep_dims = keep_dims.into().unwrap_or(false);
1257    let ddof = ddof.into().unwrap_or(0);
1258    Array::try_from_op(|res| unsafe {
1259        mlx_sys::mlx_std(res, a.as_ptr(), keep_dims, ddof, stream.as_ref().as_ptr())
1260    })
1261}
1262
1263/// Compatibility shim for [`std`].
1264#[generate_macro(customize(forwarding_shim = true))]
1265#[deprecated(
1266    since = "0.26.0",
1267    note = "use `with_stream` or `with_device` around `std`"
1268)]
1269pub fn std_device(
1270    a: impl AsRef<Array>,
1271    #[optional] keep_dims: impl Into<Option<bool>>,
1272    #[optional] ddof: impl Into<Option<i32>>,
1273    #[optional] stream: impl AsRef<Stream>,
1274) -> Result<Array> {
1275    crate::with_stream(stream.as_ref(), || std(a, keep_dims, ddof))
1276}
1277
1278/// See [`Array::sum_axes`]
1279pub fn sum_axes(
1280    array: impl AsRef<Array>,
1281    axes: &[i32],
1282    keep_dims: impl Into<Option<bool>>,
1283) -> Result<Array> {
1284    array.as_ref().sum_axes(axes, keep_dims)
1285}
1286
1287/// Compatibility shim for [`sum_axes`].
1288#[generate_macro(customize(forwarding_shim = true))]
1289#[deprecated(
1290    since = "0.26.0",
1291    note = "use `with_stream` or `with_device` around `sum_axes`"
1292)]
1293pub fn sum_axes_device(
1294    array: impl AsRef<Array>,
1295    axes: &[i32],
1296    #[optional] keep_dims: impl Into<Option<bool>>,
1297    #[optional] stream: impl AsRef<Stream>,
1298) -> Result<Array> {
1299    crate::with_stream(stream.as_ref(), || sum_axes(array, axes, keep_dims))
1300}
1301
1302/// See [`Array::sum_axis`]
1303pub fn sum_axis(
1304    array: impl AsRef<Array>,
1305    axis: i32,
1306    keep_dims: impl Into<Option<bool>>,
1307) -> Result<Array> {
1308    array.as_ref().sum_axis(axis, keep_dims)
1309}
1310
1311/// Compatibility shim for [`sum_axis`].
1312#[generate_macro(customize(forwarding_shim = true))]
1313#[deprecated(
1314    since = "0.26.0",
1315    note = "use `with_stream` or `with_device` around `sum_axis`"
1316)]
1317pub fn sum_axis_device(
1318    array: impl AsRef<Array>,
1319    axis: i32,
1320    #[optional] keep_dims: impl Into<Option<bool>>,
1321    #[optional] stream: impl AsRef<Stream>,
1322) -> Result<Array> {
1323    crate::with_stream(stream.as_ref(), || sum_axis(array, axis, keep_dims))
1324}
1325
1326/// See [`Array::sum`]
1327pub fn sum(array: impl AsRef<Array>, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
1328    array.as_ref().sum(keep_dims)
1329}
1330
1331/// Compatibility shim for [`sum`].
1332#[generate_macro(customize(forwarding_shim = true))]
1333#[deprecated(
1334    since = "0.26.0",
1335    note = "use `with_stream` or `with_device` around `sum`"
1336)]
1337pub fn sum_device(
1338    array: impl AsRef<Array>,
1339    #[optional] keep_dims: impl Into<Option<bool>>,
1340    #[optional] stream: impl AsRef<Stream>,
1341) -> Result<Array> {
1342    crate::with_stream(stream.as_ref(), || sum(array, keep_dims))
1343}
1344
1345/// See [`Array::mean_axes`]
1346pub fn mean_axes(
1347    array: impl AsRef<Array>,
1348    axes: &[i32],
1349    keep_dims: impl Into<Option<bool>>,
1350) -> Result<Array> {
1351    array.as_ref().mean_axes(axes, keep_dims)
1352}
1353
1354/// Compatibility shim for [`mean_axes`].
1355#[generate_macro(customize(forwarding_shim = true))]
1356#[deprecated(
1357    since = "0.26.0",
1358    note = "use `with_stream` or `with_device` around `mean_axes`"
1359)]
1360pub fn mean_axes_device(
1361    array: impl AsRef<Array>,
1362    axes: &[i32],
1363    #[optional] keep_dims: impl Into<Option<bool>>,
1364    #[optional] stream: impl AsRef<Stream>,
1365) -> Result<Array> {
1366    crate::with_stream(stream.as_ref(), || mean_axes(array, axes, keep_dims))
1367}
1368
1369/// See [`Array::mean_axis`]
1370pub fn mean_axis(
1371    array: impl AsRef<Array>,
1372    axis: i32,
1373    keep_dims: impl Into<Option<bool>>,
1374) -> Result<Array> {
1375    array.as_ref().mean_axis(axis, keep_dims)
1376}
1377
1378/// Compatibility shim for [`mean_axis`].
1379#[generate_macro(customize(forwarding_shim = true))]
1380#[deprecated(
1381    since = "0.26.0",
1382    note = "use `with_stream` or `with_device` around `mean_axis`"
1383)]
1384pub fn mean_axis_device(
1385    array: impl AsRef<Array>,
1386    axis: i32,
1387    #[optional] keep_dims: impl Into<Option<bool>>,
1388    #[optional] stream: impl AsRef<Stream>,
1389) -> Result<Array> {
1390    crate::with_stream(stream.as_ref(), || mean_axis(array, axis, keep_dims))
1391}
1392
1393/// See [`Array::mean`]
1394pub fn mean(array: impl AsRef<Array>, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
1395    array.as_ref().mean(keep_dims)
1396}
1397
1398/// Compatibility shim for [`mean`].
1399#[generate_macro(customize(forwarding_shim = true))]
1400#[deprecated(
1401    since = "0.26.0",
1402    note = "use `with_stream` or `with_device` around `mean`"
1403)]
1404pub fn mean_device(
1405    array: impl AsRef<Array>,
1406    #[optional] keep_dims: impl Into<Option<bool>>,
1407    #[optional] stream: impl AsRef<Stream>,
1408) -> Result<Array> {
1409    crate::with_stream(stream.as_ref(), || mean(array, keep_dims))
1410}
1411
1412/// See [`Array::min`]
1413pub fn min_axes(
1414    array: impl AsRef<Array>,
1415    axes: &[i32],
1416    keep_dims: impl Into<Option<bool>>,
1417) -> Result<Array> {
1418    array.as_ref().min_axes(axes, keep_dims)
1419}
1420
1421/// Compatibility shim for [`min_axes`].
1422#[generate_macro(customize(forwarding_shim = true))]
1423#[deprecated(
1424    since = "0.26.0",
1425    note = "use `with_stream` or `with_device` around `min_axes`"
1426)]
1427pub fn min_axes_device(
1428    array: impl AsRef<Array>,
1429    axes: &[i32],
1430    #[optional] keep_dims: impl Into<Option<bool>>,
1431    #[optional] stream: impl AsRef<Stream>,
1432) -> Result<Array> {
1433    crate::with_stream(stream.as_ref(), || min_axes(array, axes, keep_dims))
1434}
1435
1436/// See [`Array::min_axis`]
1437pub fn min_axis(
1438    array: impl AsRef<Array>,
1439    axis: i32,
1440    keep_dims: impl Into<Option<bool>>,
1441) -> Result<Array> {
1442    array.as_ref().min_axis(axis, keep_dims)
1443}
1444
1445/// Compatibility shim for [`min_axis`].
1446#[generate_macro(customize(forwarding_shim = true))]
1447#[deprecated(
1448    since = "0.26.0",
1449    note = "use `with_stream` or `with_device` around `min_axis`"
1450)]
1451pub fn min_axis_device(
1452    array: impl AsRef<Array>,
1453    axis: i32,
1454    #[optional] keep_dims: impl Into<Option<bool>>,
1455    #[optional] stream: impl AsRef<Stream>,
1456) -> Result<Array> {
1457    crate::with_stream(stream.as_ref(), || min_axis(array, axis, keep_dims))
1458}
1459
1460/// See [`Array::min`]
1461pub fn min(array: impl AsRef<Array>, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
1462    array.as_ref().min(keep_dims)
1463}
1464
1465/// Compatibility shim for [`min`].
1466#[generate_macro(customize(forwarding_shim = true))]
1467#[deprecated(
1468    since = "0.26.0",
1469    note = "use `with_stream` or `with_device` around `min`"
1470)]
1471pub fn min_device(
1472    array: impl AsRef<Array>,
1473    #[optional] keep_dims: impl Into<Option<bool>>,
1474    #[optional] stream: impl AsRef<Stream>,
1475) -> Result<Array> {
1476    crate::with_stream(stream.as_ref(), || min(array, keep_dims))
1477}
1478
1479/// See [`Array::var_axes`]
1480pub fn var_axes(
1481    array: impl AsRef<Array>,
1482    axes: &[i32],
1483    keep_dims: impl Into<Option<bool>>,
1484    ddof: impl Into<Option<i32>>,
1485) -> Result<Array> {
1486    array.as_ref().var_axes(axes, keep_dims, ddof)
1487}
1488
1489/// Compatibility shim for [`var_axes`].
1490#[generate_macro(customize(forwarding_shim = true))]
1491#[deprecated(
1492    since = "0.26.0",
1493    note = "use `with_stream` or `with_device` around `var_axes`"
1494)]
1495pub fn var_axes_device(
1496    array: impl AsRef<Array>,
1497    axes: &[i32],
1498    #[optional] keep_dims: impl Into<Option<bool>>,
1499    #[optional] ddof: impl Into<Option<i32>>,
1500    #[optional] stream: impl AsRef<Stream>,
1501) -> Result<Array> {
1502    crate::with_stream(stream.as_ref(), || var_axes(array, axes, keep_dims, ddof))
1503}
1504
1505/// See [`Array::var_axis`]
1506pub fn var_axis(
1507    array: impl AsRef<Array>,
1508    axis: i32,
1509    keep_dims: impl Into<Option<bool>>,
1510    ddof: impl Into<Option<i32>>,
1511) -> Result<Array> {
1512    array.as_ref().var_axis(axis, keep_dims, ddof)
1513}
1514
1515/// Compatibility shim for [`var_axis`].
1516#[generate_macro(customize(forwarding_shim = true))]
1517#[deprecated(
1518    since = "0.26.0",
1519    note = "use `with_stream` or `with_device` around `var_axis`"
1520)]
1521pub fn var_axis_device(
1522    array: impl AsRef<Array>,
1523    axis: i32,
1524    #[optional] keep_dims: impl Into<Option<bool>>,
1525    #[optional] ddof: impl Into<Option<i32>>,
1526    #[optional] stream: impl AsRef<Stream>,
1527) -> Result<Array> {
1528    crate::with_stream(stream.as_ref(), || var_axis(array, axis, keep_dims, ddof))
1529}
1530
1531/// See [`Array::var`]
1532pub fn var(
1533    array: impl AsRef<Array>,
1534    keep_dims: impl Into<Option<bool>>,
1535    ddof: impl Into<Option<i32>>,
1536) -> Result<Array> {
1537    array.as_ref().var(keep_dims, ddof)
1538}
1539
1540/// Compatibility shim for [`var`].
1541#[generate_macro(customize(forwarding_shim = true))]
1542#[deprecated(
1543    since = "0.26.0",
1544    note = "use `with_stream` or `with_device` around `var`"
1545)]
1546pub fn var_device(
1547    array: impl AsRef<Array>,
1548    #[optional] keep_dims: impl Into<Option<bool>>,
1549    #[optional] ddof: impl Into<Option<i32>>,
1550    #[optional] stream: impl AsRef<Stream>,
1551) -> Result<Array> {
1552    crate::with_stream(stream.as_ref(), || var(array, keep_dims, ddof))
1553}
1554
1555/// See [`Array::median_axes`]
1556pub fn median_axes(
1557    array: impl AsRef<Array>,
1558    axes: &[i32],
1559    keep_dims: impl Into<Option<bool>>,
1560) -> Result<Array> {
1561    array.as_ref().median_axes(axes, keep_dims)
1562}
1563
1564/// Compatibility shim for [`median_axes`].
1565#[generate_macro(customize(forwarding_shim = true))]
1566#[deprecated(
1567    since = "0.26.0",
1568    note = "use `with_stream` or `with_device` around `median_axes`"
1569)]
1570pub fn median_axes_device(
1571    array: impl AsRef<Array>,
1572    axes: &[i32],
1573    #[optional] keep_dims: impl Into<Option<bool>>,
1574    #[optional] stream: impl AsRef<Stream>,
1575) -> Result<Array> {
1576    crate::with_stream(stream.as_ref(), || median_axes(array, axes, keep_dims))
1577}
1578
1579/// See [`Array::median_axis`]
1580pub fn median_axis(
1581    array: impl AsRef<Array>,
1582    axis: i32,
1583    keep_dims: impl Into<Option<bool>>,
1584) -> Result<Array> {
1585    array.as_ref().median_axis(axis, keep_dims)
1586}
1587
1588/// Compatibility shim for [`median_axis`].
1589#[generate_macro(customize(forwarding_shim = true))]
1590#[deprecated(
1591    since = "0.26.0",
1592    note = "use `with_stream` or `with_device` around `median_axis`"
1593)]
1594pub fn median_axis_device(
1595    array: impl AsRef<Array>,
1596    axis: i32,
1597    #[optional] keep_dims: impl Into<Option<bool>>,
1598    #[optional] stream: impl AsRef<Stream>,
1599) -> Result<Array> {
1600    crate::with_stream(stream.as_ref(), || median_axis(array, axis, keep_dims))
1601}
1602
1603/// See [`Array::median`]
1604pub fn median(array: impl AsRef<Array>, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
1605    array.as_ref().median(keep_dims)
1606}
1607
1608/// Compatibility shim for [`median`].
1609#[generate_macro(customize(forwarding_shim = true))]
1610#[deprecated(
1611    since = "0.26.0",
1612    note = "use `with_stream` or `with_device` around `median`"
1613)]
1614pub fn median_device(
1615    array: impl AsRef<Array>,
1616    #[optional] keep_dims: impl Into<Option<bool>>,
1617    #[optional] stream: impl AsRef<Stream>,
1618) -> Result<Array> {
1619    crate::with_stream(stream.as_ref(), || median(array, keep_dims))
1620}
1621
1622/// See [`Array::logsumexp_axes`]
1623pub fn logsumexp_axes(
1624    array: impl AsRef<Array>,
1625    axes: &[i32],
1626    keep_dims: impl Into<Option<bool>>,
1627) -> Result<Array> {
1628    array.as_ref().logsumexp_axes(axes, keep_dims)
1629}
1630
1631/// Compatibility shim for [`logsumexp_axes`].
1632#[generate_macro(customize(forwarding_shim = true))]
1633#[deprecated(
1634    since = "0.26.0",
1635    note = "use `with_stream` or `with_device` around `logsumexp_axes`"
1636)]
1637pub fn logsumexp_axes_device(
1638    array: impl AsRef<Array>,
1639    axes: &[i32],
1640    #[optional] keep_dims: impl Into<Option<bool>>,
1641    #[optional] stream: impl AsRef<Stream>,
1642) -> Result<Array> {
1643    crate::with_stream(stream.as_ref(), || logsumexp_axes(array, axes, keep_dims))
1644}
1645
1646/// See [`Array::logsumexp_axis`]
1647pub fn logsumexp_axis(
1648    array: impl AsRef<Array>,
1649    axis: i32,
1650    keep_dims: impl Into<Option<bool>>,
1651) -> Result<Array> {
1652    array.as_ref().logsumexp_axis(axis, keep_dims)
1653}
1654
1655/// Compatibility shim for [`logsumexp_axis`].
1656#[generate_macro(customize(forwarding_shim = true))]
1657#[deprecated(
1658    since = "0.26.0",
1659    note = "use `with_stream` or `with_device` around `logsumexp_axis`"
1660)]
1661pub fn logsumexp_axis_device(
1662    array: impl AsRef<Array>,
1663    axis: i32,
1664    #[optional] keep_dims: impl Into<Option<bool>>,
1665    #[optional] stream: impl AsRef<Stream>,
1666) -> Result<Array> {
1667    crate::with_stream(stream.as_ref(), || logsumexp_axis(array, axis, keep_dims))
1668}
1669
1670/// See [`Array::logsumexp`]
1671pub fn logsumexp(array: impl AsRef<Array>, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
1672    array.as_ref().logsumexp(keep_dims)
1673}
1674
1675/// Compatibility shim for [`logsumexp`].
1676#[generate_macro(customize(forwarding_shim = true))]
1677#[deprecated(
1678    since = "0.26.0",
1679    note = "use `with_stream` or `with_device` around `logsumexp`"
1680)]
1681pub fn logsumexp_device(
1682    array: impl AsRef<Array>,
1683    #[optional] keep_dims: impl Into<Option<bool>>,
1684    #[optional] stream: impl AsRef<Stream>,
1685) -> Result<Array> {
1686    crate::with_stream(stream.as_ref(), || logsumexp(array, keep_dims))
1687}
1688
1689#[cfg(test)]
1690mod tests {
1691    use super::*;
1692    use pretty_assertions::assert_eq;
1693
1694    #[test]
1695    fn test_all() {
1696        let array = Array::from_slice(&[true, false, true, false], &[2, 2]);
1697
1698        assert_eq!(array.all(None).unwrap().item_exact::<bool>(), false);
1699        assert_eq!(array.all(true).unwrap().shape(), &[1, 1]);
1700        assert_eq!(
1701            array.all_axes(&[0, 1], None).unwrap().item_exact::<bool>(),
1702            false
1703        );
1704
1705        let result = array.all_axis(0, None).unwrap();
1706        assert_eq!(result.as_slice::<bool>(), &[true, false]);
1707
1708        let result = array.all_axis(1, None).unwrap();
1709        assert_eq!(result.as_slice::<bool>(), &[false, false]);
1710    }
1711
1712    #[test]
1713    fn test_all_empty_axes() {
1714        let array = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], &[3, 4]);
1715        let all = array.all_axes(&[], None).unwrap();
1716
1717        let results: &[bool] = all.as_slice();
1718        assert_eq!(
1719            results,
1720            &[false, true, true, true, true, true, true, true, true, true, true, true]
1721        );
1722    }
1723
1724    #[test]
1725    fn test_prod() {
1726        let x = Array::from_slice(&[1, 2, 3, 3], &[2, 2]);
1727        assert_eq!(x.prod(None).unwrap().item_exact::<i32>(), 18);
1728
1729        let y = x.prod(true).unwrap();
1730        assert_eq!(y.item_exact::<i32>(), 18);
1731        assert_eq!(y.shape(), &[1, 1]);
1732
1733        let result = x.prod_axis(0, None).unwrap();
1734        assert_eq!(result.as_slice::<i32>(), &[3, 6]);
1735
1736        let result = x.prod_axis(1, None).unwrap();
1737        assert_eq!(result.as_slice::<i32>(), &[2, 9])
1738    }
1739
1740    #[test]
1741    fn test_prod_empty_axes() {
1742        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1743        let result = array.prod_axes(&[], None).unwrap();
1744
1745        let results: &[i32] = result.as_slice();
1746        assert_eq!(results, &[5, 8, 4, 9]);
1747    }
1748
1749    #[test]
1750    fn test_max() {
1751        let x = Array::from_slice(&[1, 2, 3, 4], &[2, 2]);
1752        assert_eq!(x.max(None).unwrap().item_exact::<i32>(), 4);
1753        let y = x.max(true).unwrap();
1754        assert_eq!(y.item_exact::<i32>(), 4);
1755        assert_eq!(y.shape(), &[1, 1]);
1756
1757        let result = x.max_axis(0, None).unwrap();
1758        assert_eq!(result.as_slice::<i32>(), &[3, 4]);
1759
1760        let result = x.max_axis(1, None).unwrap();
1761        assert_eq!(result.as_slice::<i32>(), &[2, 4]);
1762    }
1763
1764    #[test]
1765    fn test_max_empty_axes() {
1766        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1767        let result = array.max_axes(&[], None).unwrap();
1768
1769        let results: &[i32] = result.as_slice();
1770        assert_eq!(results, &[5, 8, 4, 9]);
1771    }
1772
1773    #[test]
1774    fn test_sum() {
1775        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1776        let result = array.sum_axis(0, None).unwrap();
1777
1778        let results: &[i32] = result.as_slice();
1779        assert_eq!(results, &[9, 17]);
1780    }
1781
1782    #[test]
1783    fn test_sum_empty_axes() {
1784        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1785        let result = array.sum_axes(&[], None).unwrap();
1786
1787        let results: &[i32] = result.as_slice();
1788        assert_eq!(results, &[5, 8, 4, 9]);
1789    }
1790
1791    #[test]
1792    fn test_mean() {
1793        let x = Array::from_slice(&[1, 2, 3, 4], &[2, 2]);
1794        assert_eq!(x.mean(None).unwrap().item_exact::<f32>(), 2.5);
1795        let y = x.mean(true).unwrap();
1796        assert_eq!(y.item_exact::<f32>(), 2.5);
1797        assert_eq!(y.shape(), &[1, 1]);
1798
1799        let result = x.mean_axis(0, None).unwrap();
1800        assert_eq!(result.as_slice::<f32>(), &[2.0, 3.0]);
1801
1802        let result = x.mean_axis(1, None).unwrap();
1803        assert_eq!(result.as_slice::<f32>(), &[1.5, 3.5]);
1804    }
1805
1806    #[test]
1807    fn test_mean_empty_axes() {
1808        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1809        let result = array.mean_axes(&[], None).unwrap();
1810
1811        let results: &[f32] = result.as_slice();
1812        assert_eq!(results, &[5.0, 8.0, 4.0, 9.0]);
1813    }
1814
1815    #[test]
1816    fn test_mean_out_of_bounds() {
1817        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1818        let result = array.mean_axis(2, None);
1819        assert!(result.is_err());
1820    }
1821
1822    #[test]
1823    fn test_min() {
1824        let x = Array::from_slice(&[1, 2, 3, 4], &[2, 2]);
1825        assert_eq!(x.min(None).unwrap().item_exact::<i32>(), 1);
1826        let y = x.min(true).unwrap();
1827        assert_eq!(y.item_exact::<i32>(), 1);
1828        assert_eq!(y.shape(), &[1, 1]);
1829
1830        let result = x.min_axis(0, None).unwrap();
1831        assert_eq!(result.as_slice::<i32>(), &[1, 2]);
1832
1833        let result = x.min_axis(1, None).unwrap();
1834        assert_eq!(result.as_slice::<i32>(), &[1, 3]);
1835    }
1836
1837    #[test]
1838    fn test_min_empty_axes() {
1839        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1840        let result = array.min_axes(&[], None).unwrap();
1841
1842        let results: &[i32] = result.as_slice();
1843        assert_eq!(results, &[5, 8, 4, 9]);
1844    }
1845
1846    #[test]
1847    fn test_var() {
1848        let x = Array::from_slice(&[1, 2, 3, 4], &[2, 2]);
1849        assert_eq!(x.var(None, None).unwrap().item_exact::<f32>(), 1.25);
1850        let y = x.var(true, None).unwrap();
1851        assert_eq!(y.item_exact::<f32>(), 1.25);
1852        assert_eq!(y.shape(), &[1, 1]);
1853
1854        let result = x.var_axis(0, None, None).unwrap();
1855        assert_eq!(result.as_slice::<f32>(), &[1.0, 1.0]);
1856
1857        let result = x.var_axis(1, None, None).unwrap();
1858        assert_eq!(result.as_slice::<f32>(), &[0.25, 0.25]);
1859
1860        let x = Array::from_slice(&[1.0, 2.0], &[2]);
1861        let out = x.var(None, Some(3)).unwrap();
1862        assert_eq!(out.item_exact::<f32>(), f32::INFINITY);
1863    }
1864
1865    #[test]
1866    fn test_var_empty_axes() {
1867        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1868        let result = array.var_axes(&[], None, 0).unwrap();
1869
1870        let results: &[f32] = result.as_slice();
1871        assert_eq!(results, &[0.0, 0.0, 0.0, 0.0]);
1872    }
1873
1874    #[test]
1875    fn test_log_sum_exp() {
1876        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1877        let result = array.logsumexp_axis(0, None).unwrap();
1878
1879        let results: &[f32] = result.as_slice();
1880        assert_eq!(results, &[5.3132615, 9.313262]);
1881    }
1882
1883    #[test]
1884    fn test_log_sum_exp_empty_axes() {
1885        let array = Array::from_slice(&[5, 8, 4, 9], &[2, 2]);
1886        let result = array.logsumexp_axes(&[], None).unwrap();
1887
1888        let results: &[f32] = result.as_slice();
1889        assert_eq!(results, &[5.0, 8.0, 4.0, 9.0]);
1890    }
1891
1892    // Tests adapted from Python test `test_ops.py/test_median`
1893    #[test]
1894    fn test_median() {
1895        // Test basic median over all elements (odd count)
1896        let x = Array::from_slice(&[0, 1, 2, 3, 4], &[5]);
1897        let out = x.median(None).unwrap();
1898        assert_eq!(out.shape(), &[] as &[i32]);
1899        assert_eq!(out.item_exact::<f32>(), 2.0);
1900
1901        // Test keepdims
1902        let out = x.median(true).unwrap();
1903        assert_eq!(out.shape(), &[1]);
1904
1905        // Test median with even count (should be average of two middle values)
1906        let x = Array::from_slice(&[0, 1, 2, 3, 4, 5], &[6]);
1907        let out = x.median(None).unwrap();
1908        assert!((out.item_exact::<f32>() - 2.5).abs() < 1e-5);
1909
1910        // Test median over specific axes
1911        use crate::random;
1912        random::seed(0).unwrap();
1913        let x = random::normal::<f32>(&[5, 5, 5, 5], None, None, None).unwrap();
1914
1915        let out = x.median_axes(&[0, 2], true).unwrap();
1916        assert_eq!(out.shape(), &[1, 5, 1, 5]);
1917
1918        let out = x.median_axes(&[1, 3], true).unwrap();
1919        assert_eq!(out.shape(), &[5, 1, 5, 1]);
1920
1921        // Test single axis
1922        let x = Array::from_slice(&[1, 5, 2, 4, 3, 6], &[2, 3]);
1923        let out = x.median_axis(0, None).unwrap();
1924        assert_eq!(out.shape(), &[3]);
1925
1926        let out = x.median_axis(1, None).unwrap();
1927        assert_eq!(out.shape(), &[2]);
1928    }
1929}
mlx_rs/ops/reduction.rs

mlx_rs/ops/
reduction.rs
