mlx_rs/ops/
logical.rs

1use crate::array::{Array, ArrayElement};
2use crate::error::Result;
3use crate::utils::guard::Guarded;
4use crate::Stream;
5use mlx_internal_macros::generate_macro;
6
7impl Array {
8    /// Element-wise equality returning an error if the arrays are not broadcastable.
9    ///
10    /// Equality comparison on two arrays with
11    /// [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting).
12    ///
13    /// # Params
14    ///
15    /// - other: array to compare
16    ///
17    /// # Example
18    ///
19    /// ```rust
20    /// use mlx_rs::Array;
21    /// let a = Array::from_slice(&[1, 2, 3], &[3]);
22    /// let b = Array::from_slice(&[1, 2, 3], &[3]);
23    /// let mut c = a.eq(&b).unwrap();
24    ///
25    /// let c_data: &[bool] = c.as_slice();
26    /// // c_data == [true, true, true]
27    /// ```
28    pub fn eq(&self, other: impl AsRef<Array>) -> Result<Array> {
29        let stream = Stream::thread_local_or_default();
30        Array::try_from_op(|res| unsafe {
31            mlx_sys::mlx_equal(
32                res,
33                self.as_ptr(),
34                other.as_ref().as_ptr(),
35                stream.as_ref().as_ptr(),
36            )
37        })
38    }
39
40    /// Compatibility shim for [`eq`].
41    #[deprecated(
42        since = "0.26.0",
43        note = "use `with_stream` or `with_device` around `eq`"
44    )]
45    pub fn eq_device(&self, other: impl AsRef<Array>, stream: impl AsRef<Stream>) -> Result<Array> {
46        crate::with_stream(stream.as_ref(), || self.eq(other))
47    }
48
49    /// Element-wise less than or equal returning an error if the arrays are not broadcastable.
50    ///
51    /// Less than or equal on two arrays with
52    /// [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting).
53    ///
54    /// # Params
55    ///
56    /// - other: array to compare
57    ///
58    /// # Example
59    ///
60    /// ```rust
61    /// use mlx_rs::Array;
62    /// let a = Array::from_slice(&[1, 2, 3], &[3]);
63    /// let b = Array::from_slice(&[1, 2, 3], &[3]);
64    /// let mut c = a.le(&b).unwrap();
65    ///
66    /// let c_data: &[bool] = c.as_slice();
67    /// // c_data == [true, true, true]
68    /// ```
69    pub fn le(&self, other: impl AsRef<Array>) -> Result<Array> {
70        let stream = Stream::thread_local_or_default();
71        Array::try_from_op(|res| unsafe {
72            mlx_sys::mlx_less_equal(
73                res,
74                self.as_ptr(),
75                other.as_ref().as_ptr(),
76                stream.as_ref().as_ptr(),
77            )
78        })
79    }
80
81    /// Compatibility shim for [`le`].
82    #[deprecated(
83        since = "0.26.0",
84        note = "use `with_stream` or `with_device` around `le`"
85    )]
86    pub fn le_device(&self, other: impl AsRef<Array>, stream: impl AsRef<Stream>) -> Result<Array> {
87        crate::with_stream(stream.as_ref(), || self.le(other))
88    }
89
90    /// Element-wise greater than or equal returning an error if the arrays are not broadcastable.
91    ///
92    /// Greater than or equal on two arrays with
93    /// [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting).
94    ///
95    /// # Params
96    ///
97    /// - other: array to compare
98    ///
99    /// # Example
100    ///
101    /// ```rust
102    /// use mlx_rs::Array;
103    /// let a = Array::from_slice(&[1, 2, 3], &[3]);
104    /// let b = Array::from_slice(&[1, 2, 3], &[3]);
105    /// let mut c = a.ge(&b).unwrap();
106    ///
107    /// let c_data: &[bool] = c.as_slice();
108    /// // c_data == [true, true, true]
109    /// ```
110    pub fn ge(&self, other: impl AsRef<Array>) -> Result<Array> {
111        let stream = Stream::thread_local_or_default();
112        Array::try_from_op(|res| unsafe {
113            mlx_sys::mlx_greater_equal(
114                res,
115                self.as_ptr(),
116                other.as_ref().as_ptr(),
117                stream.as_ref().as_ptr(),
118            )
119        })
120    }
121
122    /// Compatibility shim for [`ge`].
123    #[deprecated(
124        since = "0.26.0",
125        note = "use `with_stream` or `with_device` around `ge`"
126    )]
127    pub fn ge_device(&self, other: impl AsRef<Array>, stream: impl AsRef<Stream>) -> Result<Array> {
128        crate::with_stream(stream.as_ref(), || self.ge(other))
129    }
130
131    /// Element-wise not equal returning an error if the arrays are not broadcastable.
132    ///
133    /// Not equal on two arrays with
134    /// [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting).
135    ///
136    /// # Params
137    ///
138    /// - other: array to compare
139    ///
140    /// # Example
141    ///
142    /// ```rust
143    /// use mlx_rs::Array;
144    /// let a = Array::from_slice(&[1, 2, 3], &[3]);
145    /// let b = Array::from_slice(&[1, 2, 3], &[3]);
146    /// let mut c = a.ne(&b).unwrap();
147    ///
148    /// let c_data: &[bool] = c.as_slice();
149    /// // c_data == [false, false, false]
150    /// ```
151    pub fn ne(&self, other: impl AsRef<Array>) -> Result<Array> {
152        let stream = Stream::thread_local_or_default();
153        Array::try_from_op(|res| unsafe {
154            mlx_sys::mlx_not_equal(
155                res,
156                self.as_ptr(),
157                other.as_ref().as_ptr(),
158                stream.as_ref().as_ptr(),
159            )
160        })
161    }
162
163    /// Compatibility shim for [`ne`].
164    #[deprecated(
165        since = "0.26.0",
166        note = "use `with_stream` or `with_device` around `ne`"
167    )]
168    pub fn ne_device(&self, other: impl AsRef<Array>, stream: impl AsRef<Stream>) -> Result<Array> {
169        crate::with_stream(stream.as_ref(), || self.ne(other))
170    }
171
172    /// Element-wise less than returning an error if the arrays are not broadcastable.
173    ///
174    /// Less than on two arrays with [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting).
175    ///
176    /// # Params
177    ///
178    /// - other: array to compare
179    ///
180    /// # Example
181    ///
182    /// ```rust
183    /// use mlx_rs::Array;
184    /// let a = Array::from_slice(&[1, 2, 3], &[3]);
185    /// let b = Array::from_slice(&[1, 2, 3], &[3]);
186    /// let mut c = a.lt(&b).unwrap();
187    ///
188    /// let c_data: &[bool] = c.as_slice();
189    /// // c_data == [false, false, false]
190    /// ```
191    pub fn lt(&self, other: impl AsRef<Array>) -> Result<Array> {
192        let stream = Stream::thread_local_or_default();
193        Array::try_from_op(|res| unsafe {
194            mlx_sys::mlx_less(
195                res,
196                self.as_ptr(),
197                other.as_ref().as_ptr(),
198                stream.as_ref().as_ptr(),
199            )
200        })
201    }
202
203    /// Compatibility shim for [`lt`].
204    #[deprecated(
205        since = "0.26.0",
206        note = "use `with_stream` or `with_device` around `lt`"
207    )]
208    pub fn lt_device(&self, other: impl AsRef<Array>, stream: impl AsRef<Stream>) -> Result<Array> {
209        crate::with_stream(stream.as_ref(), || self.lt(other))
210    }
211
212    /// Element-wise greater than returning an error if the arrays are not broadcastable.
213    ///
214    /// Greater than on two arrays with [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting).
215    ///
216    /// # Params
217    ///
218    /// - other: array to compare
219    ///
220    /// # Example
221    ///
222    /// ```rust
223    /// use mlx_rs::Array;
224    /// let a = Array::from_slice(&[1, 2, 3], &[3]);
225    /// let b = Array::from_slice(&[1, 2, 3], &[3]);
226    /// let mut c = a.gt(&b).unwrap();
227    ///
228    /// let c_data: &[bool] = c.as_slice();
229    /// // c_data == [false, false, false]
230    /// ```
231    pub fn gt(&self, other: impl AsRef<Array>) -> Result<Array> {
232        let stream = Stream::thread_local_or_default();
233        Array::try_from_op(|res| unsafe {
234            mlx_sys::mlx_greater(
235                res,
236                self.as_ptr(),
237                other.as_ref().as_ptr(),
238                stream.as_ref().as_ptr(),
239            )
240        })
241    }
242
243    /// Compatibility shim for [`gt`].
244    #[deprecated(
245        since = "0.26.0",
246        note = "use `with_stream` or `with_device` around `gt`"
247    )]
248    pub fn gt_device(&self, other: impl AsRef<Array>, stream: impl AsRef<Stream>) -> Result<Array> {
249        crate::with_stream(stream.as_ref(), || self.gt(other))
250    }
251
252    /// Element-wise logical and returning an error if the arrays are not broadcastable.
253    ///
254    /// Logical and on two arrays with [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting).
255    ///
256    /// # Params
257    ///
258    /// - other: array to compare
259    ///
260    /// # Example
261    ///
262    /// ```rust
263    /// use mlx_rs::Array;
264    /// let a = Array::from_slice(&[true, false, true], &[3]);
265    /// let b = Array::from_slice(&[true, true, false], &[3]);
266    /// let mut c = a.logical_and(&b).unwrap();
267    ///
268    /// let c_data: &[bool] = c.as_slice();
269    /// // c_data == [true, false, false]
270    /// ```
271    pub fn logical_and(&self, other: impl AsRef<Array>) -> Result<Array> {
272        let stream = Stream::thread_local_or_default();
273        Array::try_from_op(|res| unsafe {
274            mlx_sys::mlx_logical_and(
275                res,
276                self.as_ptr(),
277                other.as_ref().as_ptr(),
278                stream.as_ref().as_ptr(),
279            )
280        })
281    }
282
283    /// Compatibility shim for [`logical_and`].
284    #[deprecated(
285        since = "0.26.0",
286        note = "use `with_stream` or `with_device` around `logical_and`"
287    )]
288    pub fn logical_and_device(
289        &self,
290        other: impl AsRef<Array>,
291        stream: impl AsRef<Stream>,
292    ) -> Result<Array> {
293        crate::with_stream(stream.as_ref(), || self.logical_and(other))
294    }
295
296    /// Element-wise logical or returning an error if the arrays are not broadcastable.
297    ///
298    /// Logical or on two arrays with [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting).
299    ///
300    /// # Params
301    ///
302    /// - other: array to compare
303    ///
304    /// # Example
305    ///
306    /// ```rust
307    /// use mlx_rs::Array;
308    /// let a = Array::from_slice(&[true, false, true], &[3]);
309    /// let b = Array::from_slice(&[true, true, false], &[3]);
310    /// let mut c = a.logical_or(&b).unwrap();
311    ///
312    /// let c_data: &[bool] = c.as_slice();
313    /// // c_data == [true, true, true]
314    /// ```
315    pub fn logical_or(&self, other: impl AsRef<Array>) -> Result<Array> {
316        let stream = Stream::thread_local_or_default();
317        Array::try_from_op(|res| unsafe {
318            mlx_sys::mlx_logical_or(
319                res,
320                self.as_ptr(),
321                other.as_ref().as_ptr(),
322                stream.as_ref().as_ptr(),
323            )
324        })
325    }
326
327    /// Compatibility shim for [`logical_or`].
328    #[deprecated(
329        since = "0.26.0",
330        note = "use `with_stream` or `with_device` around `logical_or`"
331    )]
332    pub fn logical_or_device(
333        &self,
334        other: impl AsRef<Array>,
335        stream: impl AsRef<Stream>,
336    ) -> Result<Array> {
337        crate::with_stream(stream.as_ref(), || self.logical_or(other))
338    }
339
340    /// Unary element-wise logical not.
341    ///
342    /// # Example
343    ///
344    /// ```rust
345    /// use mlx_rs::Array;
346    /// let a: Array = false.into();
347    /// let mut b = a.logical_not().unwrap();
348    ///
349    /// let b_data: &[bool] = b.as_slice();
350    /// // b_data == [true]
351    /// ```
352    pub fn logical_not(&self) -> Result<Array> {
353        let stream = Stream::thread_local_or_default();
354        Array::try_from_op(|res| unsafe {
355            mlx_sys::mlx_logical_not(res, self.as_ptr(), stream.as_ref().as_ptr())
356        })
357    }
358
359    /// Compatibility shim for [`logical_not`].
360    #[deprecated(
361        since = "0.26.0",
362        note = "use `with_stream` or `with_device` around `logical_not`"
363    )]
364    pub fn logical_not_device(&self, stream: impl AsRef<Stream>) -> Result<Array> {
365        crate::with_stream(stream.as_ref(), || self.logical_not())
366    }
367
368    /// Approximate comparison of two arrays returning an error if the inputs aren't valid.
369    ///
370    /// This evaluates the comparison result before returning a Rust `bool`.
371    ///
372    /// The arrays are considered equal if:
373    ///
374    /// ```text
375    /// all(abs(a - b) <= (atol + rtol * abs(b)))
376    /// ```
377    ///
378    /// # Params
379    ///
380    /// - other: array to compare
381    /// - rtol: relative tolerance = defaults to 1e-5 when None
382    /// - atol: absolute tolerance - defaults to 1e-8 when None
383    /// - equal_nan: whether to consider NaNs equal -- default is false when None
384    ///
385    /// # Example
386    ///
387    /// ```rust
388    /// use num_traits::Pow;
389    /// use mlx_rs::array;
390    /// let a = array!([0., 1., 2., 3.]).sqrt().unwrap();
391    /// let b = array!([0., 1., 2., 3.]).power(array!(0.5)).unwrap();
392    /// assert!(a.all_close(&b, None, None, None).unwrap());
393    /// ```
394    pub fn all_close(
395        &self,
396        other: impl AsRef<Array>,
397        rtol: impl Into<Option<f64>>,
398        atol: impl Into<Option<f64>>,
399        equal_nan: impl Into<Option<bool>>,
400    ) -> Result<bool> {
401        let stream = Stream::thread_local_or_default();
402        let result = Array::try_from_op(|res| unsafe {
403            mlx_sys::mlx_allclose(
404                res,
405                self.as_ptr(),
406                other.as_ref().as_ptr(),
407                rtol.into().unwrap_or(1e-5),
408                atol.into().unwrap_or(1e-8),
409                equal_nan.into().unwrap_or(false),
410                stream.as_ref().as_ptr(),
411            )
412        })?;
413        result.eval()?;
414        bool::array_item(&result)
415    }
416
417    /// Compatibility shim for [`all_close`].
418    #[deprecated(
419        since = "0.26.0",
420        note = "use `with_stream` or `with_device` around `all_close`"
421    )]
422    pub fn all_close_device(
423        &self,
424        other: impl AsRef<Array>,
425        rtol: impl Into<Option<f64>>,
426        atol: impl Into<Option<f64>>,
427        equal_nan: impl Into<Option<bool>>,
428        stream: impl AsRef<Stream>,
429    ) -> Result<bool> {
430        crate::with_stream(stream.as_ref(), || {
431            self.all_close(other, rtol, atol, equal_nan)
432        })
433    }
434
435    /// Returns a boolean array where two arrays are element-wise equal within a tolerance returning an error if the arrays are not broadcastable.
436    ///
437    /// Infinite values are considered equal if they have the same sign, NaN values are not equal unless
438    /// `equalNAN` is `true`.
439    ///
440    /// Two values are considered close if:
441    ///
442    /// ```text
443    /// abs(a - b) <= (atol + rtol * abs(b))
444    /// ```
445    ///
446    /// Unlike [self.array_eq] this function supports [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting).
447    pub fn is_close(
448        &self,
449        other: impl AsRef<Array>,
450        rtol: impl Into<Option<f64>>,
451        atol: impl Into<Option<f64>>,
452        equal_nan: impl Into<Option<bool>>,
453    ) -> Result<Array> {
454        let stream = Stream::thread_local_or_default();
455        Array::try_from_op(|res| unsafe {
456            mlx_sys::mlx_isclose(
457                res,
458                self.as_ptr(),
459                other.as_ref().as_ptr(),
460                rtol.into().unwrap_or(1e-5),
461                atol.into().unwrap_or(1e-8),
462                equal_nan.into().unwrap_or(false),
463                stream.as_ref().as_ptr(),
464            )
465        })
466    }
467
468    /// Compatibility shim for [`is_close`].
469    #[deprecated(
470        since = "0.26.0",
471        note = "use `with_stream` or `with_device` around `is_close`"
472    )]
473    pub fn is_close_device(
474        &self,
475        other: impl AsRef<Array>,
476        rtol: impl Into<Option<f64>>,
477        atol: impl Into<Option<f64>>,
478        equal_nan: impl Into<Option<bool>>,
479        stream: impl AsRef<Stream>,
480    ) -> Result<Array> {
481        crate::with_stream(stream.as_ref(), || {
482            self.is_close(other, rtol, atol, equal_nan)
483        })
484    }
485
486    /// Array equality check.
487    ///
488    /// Compare two arrays for equality. Returns `true` iff the arrays have
489    /// the same shape and their values are equal. The arrays need not have
490    /// the same type to be considered equal.
491    ///
492    /// # Params
493    ///
494    /// - other: array to compare
495    /// - equal_nan: whether to consider NaNs equal -- default is false when None
496    ///
497    /// # Example
498    ///
499    /// ```rust
500    /// use mlx_rs::Array;
501    /// let a = Array::from_slice(&[0, 1, 2, 3], &[4]);
502    /// let b = Array::from_slice(&[0., 1., 2., 3.], &[4]);
503    ///
504    /// assert!(a.eq_values(&b).unwrap());
505    /// ```
506    #[deprecated(since = "0.26.0", note = "use `eq_values` for a Rust boolean")]
507    pub fn array_eq(
508        &self,
509        other: impl AsRef<Array>,
510        equal_nan: impl Into<Option<bool>>,
511    ) -> Result<Array> {
512        self.array_eq_result(other.as_ref(), equal_nan.into().unwrap_or(false))
513    }
514
515    fn array_eq_result(&self, other: &Array, equal_nan: bool) -> Result<Array> {
516        let stream = Stream::thread_local_or_default();
517        Array::try_from_op(|res| unsafe {
518            mlx_sys::mlx_array_equal(
519                res,
520                self.as_ptr(),
521                other.as_ptr(),
522                equal_nan,
523                stream.as_ref().as_ptr(),
524            )
525        })
526    }
527
528    /// Compare shape, dtype, and values exactly.
529    ///
530    /// A dtype or shape mismatch returns `false` without evaluating either array. Otherwise this
531    /// evaluates the equality result before returning a Rust `bool`. NaNs compare unequal.
532    pub fn eq_exact(&self, other: impl AsRef<Array>) -> Result<bool> {
533        let other = other.as_ref();
534        if self.dtype() != other.dtype() || self.shape() != other.shape() {
535            return Ok(false);
536        }
537        let result = self.array_eq_result(other, false)?;
538        result.eval()?;
539        bool::array_item(&result)
540    }
541
542    /// Compare shape and values exactly while allowing different dtypes.
543    ///
544    /// This evaluates the equality result before returning a Rust `bool`. NaNs compare unequal.
545    pub fn eq_values(&self, other: impl AsRef<Array>) -> Result<bool> {
546        let result = self.array_eq_result(other.as_ref(), false)?;
547        result.eval()?;
548        bool::array_item(&result)
549    }
550
551    /// Compatibility shim for [`array_eq`].
552    #[deprecated(
553        since = "0.26.0",
554        note = "use `with_stream` or `with_device` around `array_eq`"
555    )]
556    pub fn array_eq_device(
557        &self,
558        other: impl AsRef<Array>,
559        equal_nan: impl Into<Option<bool>>,
560        stream: impl AsRef<Stream>,
561    ) -> Result<Array> {
562        let other = other.as_ref();
563        let equal_nan = equal_nan.into().unwrap_or(false);
564        crate::with_stream(stream.as_ref(), || self.array_eq_result(other, equal_nan))
565    }
566
567    /// An `or` reduction over the given axes returning an error if the axes are invalid.
568    ///
569    /// # Params
570    ///
571    /// - axes: axes to reduce over -- defaults to all axes if not provided
572    /// - keep_dims: if `true` keep reduced axis as singleton dimension -- defaults to false if not provided
573    ///
574    ///  # Example
575    ///
576    /// ```rust
577    /// use mlx_rs::Array;
578    ///
579    /// let array = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], &[3, 4]);
580    ///
581    /// // will produce a scalar Array with true -- some of the values are non-zero
582    /// let all = array.any(None).unwrap();
583    ///
584    /// // produces an Array([true, true, true, true]) -- all rows have non-zeros
585    /// let all_rows = array.any_axes(&[0], None).unwrap();
586    /// ```
587    pub fn any_axes(&self, axes: &[i32], keep_dims: impl Into<Option<bool>>) -> Result<Array> {
588        let stream = Stream::thread_local_or_default();
589        Array::try_from_op(|res| unsafe {
590            mlx_sys::mlx_any_axes(
591                res,
592                self.as_ptr(),
593                axes.as_ptr(),
594                axes.len(),
595                keep_dims.into().unwrap_or(false),
596                stream.as_ref().as_ptr(),
597            )
598        })
599    }
600
601    /// Compatibility shim for [`any_axes`].
602    #[deprecated(
603        since = "0.26.0",
604        note = "use `with_stream` or `with_device` around `any_axes`"
605    )]
606    pub fn any_axes_device(
607        &self,
608        axes: &[i32],
609        keep_dims: impl Into<Option<bool>>,
610        stream: impl AsRef<Stream>,
611    ) -> Result<Array> {
612        crate::with_stream(stream.as_ref(), || self.any_axes(axes, keep_dims))
613    }
614
615    /// Similar to [`any_axes`] but defaults to all axes.
616    pub fn any_axis(&self, axis: i32, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
617        let stream = Stream::thread_local_or_default();
618        Array::try_from_op(|res| unsafe {
619            mlx_sys::mlx_any_axis(
620                res,
621                self.as_ptr(),
622                axis,
623                keep_dims.into().unwrap_or(false),
624                stream.as_ref().as_ptr(),
625            )
626        })
627    }
628
629    /// Compatibility shim for [`any_axis`].
630    #[deprecated(
631        since = "0.26.0",
632        note = "use `with_stream` or `with_device` around `any_axis`"
633    )]
634    pub fn any_axis_device(
635        &self,
636        axis: i32,
637        keep_dims: impl Into<Option<bool>>,
638        stream: impl AsRef<Stream>,
639    ) -> Result<Array> {
640        crate::with_stream(stream.as_ref(), || self.any_axis(axis, keep_dims))
641    }
642
643    /// Similar to [`any_axes`] but defaults to all axes.
644    pub fn any(&self, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
645        let stream = Stream::thread_local_or_default();
646        Array::try_from_op(|res| unsafe {
647            mlx_sys::mlx_any(
648                res,
649                self.as_ptr(),
650                keep_dims.into().unwrap_or(false),
651                stream.as_ref().as_ptr(),
652            )
653        })
654    }
655
656    /// Compatibility shim for [`any`].
657    #[deprecated(
658        since = "0.26.0",
659        note = "use `with_stream` or `with_device` around `any`"
660    )]
661    pub fn any_device(
662        &self,
663        keep_dims: impl Into<Option<bool>>,
664        stream: impl AsRef<Stream>,
665    ) -> Result<Array> {
666        crate::with_stream(stream.as_ref(), || self.any(keep_dims))
667    }
668}
669
670/// Cast both operands to boolean, broadcast them, and compare for inequality.
671///
672/// ```rust
673/// use mlx_rs::{array, ops::logical_xor};
674///
675/// let output = logical_xor(array!([true, false]), array!([false, false])).unwrap();
676/// assert_eq!(output.shape(), &[2]);
677/// ```
678pub fn logical_xor(lhs: impl AsRef<Array>, rhs: impl AsRef<Array>) -> Result<Array> {
679    let stream = Stream::thread_local_or_default();
680    Array::try_from_op(|res| unsafe {
681        mlx_sys::mlx_logical_xor(
682            res,
683            lhs.as_ref().as_ptr(),
684            rhs.as_ref().as_ptr(),
685            stream.as_ref().as_ptr(),
686        )
687    })
688}
689
690/// See [`Array::any`]
691pub fn any_axes(
692    array: impl AsRef<Array>,
693    axes: &[i32],
694    keep_dims: impl Into<Option<bool>>,
695) -> Result<Array> {
696    array.as_ref().any_axes(axes, keep_dims)
697}
698
699/// Compatibility shim for [`any_axes`].
700#[generate_macro(customize(forwarding_shim = true))]
701#[deprecated(
702    since = "0.26.0",
703    note = "use `with_stream` or `with_device` around `any_axes`"
704)]
705pub fn any_axes_device(
706    array: impl AsRef<Array>,
707    axes: &[i32],
708    #[optional] keep_dims: impl Into<Option<bool>>,
709    #[optional] stream: impl AsRef<Stream>,
710) -> Result<Array> {
711    crate::with_stream(stream.as_ref(), || any_axes(array, axes, keep_dims))
712}
713
714/// See [`Array::any`]
715pub fn any_axis(
716    array: impl AsRef<Array>,
717    axis: i32,
718    keep_dims: impl Into<Option<bool>>,
719) -> Result<Array> {
720    array.as_ref().any_axis(axis, keep_dims)
721}
722
723/// Compatibility shim for [`any_axis`].
724#[generate_macro(customize(forwarding_shim = true))]
725#[deprecated(
726    since = "0.26.0",
727    note = "use `with_stream` or `with_device` around `any_axis`"
728)]
729pub fn any_axis_device(
730    array: impl AsRef<Array>,
731    axis: i32,
732    #[optional] keep_dims: impl Into<Option<bool>>,
733    #[optional] stream: impl AsRef<Stream>,
734) -> Result<Array> {
735    crate::with_stream(stream.as_ref(), || any_axis(array, axis, keep_dims))
736}
737
738/// See [`Array::any`]
739pub fn any(array: impl AsRef<Array>, keep_dims: impl Into<Option<bool>>) -> Result<Array> {
740    array.as_ref().any(keep_dims)
741}
742
743/// Compatibility shim for [`any`].
744#[generate_macro(customize(forwarding_shim = true))]
745#[deprecated(
746    since = "0.26.0",
747    note = "use `with_stream` or `with_device` around `any`"
748)]
749pub fn any_device(
750    array: impl AsRef<Array>,
751    #[optional] keep_dims: impl Into<Option<bool>>,
752    #[optional] stream: impl AsRef<Stream>,
753) -> Result<Array> {
754    crate::with_stream(stream.as_ref(), || any(array, keep_dims))
755}
756
757/// See [`Array::logical_and`]
758pub fn logical_and(a: impl AsRef<Array>, b: impl AsRef<Array>) -> Result<Array> {
759    a.as_ref().logical_and(b)
760}
761
762/// Compatibility shim for [`logical_and`].
763#[generate_macro(customize(forwarding_shim = true))]
764#[deprecated(
765    since = "0.26.0",
766    note = "use `with_stream` or `with_device` around `logical_and`"
767)]
768pub fn logical_and_device(
769    a: impl AsRef<Array>,
770    b: impl AsRef<Array>,
771    #[optional] stream: impl AsRef<Stream>,
772) -> Result<Array> {
773    crate::with_stream(stream.as_ref(), || logical_and(a, b))
774}
775
776/// See [`Array::logical_or`]
777pub fn logical_or(a: impl AsRef<Array>, b: impl AsRef<Array>) -> Result<Array> {
778    a.as_ref().logical_or(b)
779}
780
781/// Compatibility shim for [`logical_or`].
782#[generate_macro(customize(forwarding_shim = true))]
783#[deprecated(
784    since = "0.26.0",
785    note = "use `with_stream` or `with_device` around `logical_or`"
786)]
787pub fn logical_or_device(
788    a: impl AsRef<Array>,
789    b: impl AsRef<Array>,
790    #[optional] stream: impl AsRef<Stream>,
791) -> Result<Array> {
792    crate::with_stream(stream.as_ref(), || logical_or(a, b))
793}
794
795/// See [`Array::logical_not`]
796pub fn logical_not(a: impl AsRef<Array>) -> Result<Array> {
797    a.as_ref().logical_not()
798}
799
800/// Compatibility shim for [`logical_not`].
801#[generate_macro(customize(forwarding_shim = true))]
802#[deprecated(
803    since = "0.26.0",
804    note = "use `with_stream` or `with_device` around `logical_not`"
805)]
806pub fn logical_not_device(
807    a: impl AsRef<Array>,
808    #[optional] stream: impl AsRef<Stream>,
809) -> Result<Array> {
810    crate::with_stream(stream.as_ref(), || logical_not(a))
811}
812
813/// See [`Array::all_close`]
814pub fn all_close(
815    a: impl AsRef<Array>,
816    b: impl AsRef<Array>,
817    rtol: impl Into<Option<f64>>,
818    atol: impl Into<Option<f64>>,
819    equal_nan: impl Into<Option<bool>>,
820) -> Result<bool> {
821    a.as_ref().all_close(b, rtol, atol, equal_nan)
822}
823
824/// Compatibility shim for [`all_close`].
825#[generate_macro(customize(forwarding_shim = true))]
826#[deprecated(
827    since = "0.26.0",
828    note = "use `with_stream` or `with_device` around `all_close`"
829)]
830pub fn all_close_device(
831    a: impl AsRef<Array>,
832    b: impl AsRef<Array>,
833    #[optional] rtol: impl Into<Option<f64>>,
834    #[optional] atol: impl Into<Option<f64>>,
835    #[optional] equal_nan: impl Into<Option<bool>>,
836    #[optional] stream: impl AsRef<Stream>,
837) -> Result<bool> {
838    crate::with_stream(stream.as_ref(), || all_close(a, b, rtol, atol, equal_nan))
839}
840
841/// See [`Array::is_close`]
842pub fn is_close(
843    a: impl AsRef<Array>,
844    b: impl AsRef<Array>,
845    rtol: impl Into<Option<f64>>,
846    atol: impl Into<Option<f64>>,
847    equal_nan: impl Into<Option<bool>>,
848) -> Result<Array> {
849    a.as_ref().is_close(b, rtol, atol, equal_nan)
850}
851
852/// Compatibility shim for [`is_close`].
853#[generate_macro(customize(forwarding_shim = true))]
854#[deprecated(
855    since = "0.26.0",
856    note = "use `with_stream` or `with_device` around `is_close`"
857)]
858pub fn is_close_device(
859    a: impl AsRef<Array>,
860    b: impl AsRef<Array>,
861    #[optional] rtol: impl Into<Option<f64>>,
862    #[optional] atol: impl Into<Option<f64>>,
863    #[optional] equal_nan: impl Into<Option<bool>>,
864    #[optional] stream: impl AsRef<Stream>,
865) -> Result<Array> {
866    crate::with_stream(stream.as_ref(), || is_close(a, b, rtol, atol, equal_nan))
867}
868
869/// See [`Array::array_eq`]
870pub fn array_eq(
871    a: impl AsRef<Array>,
872    b: impl AsRef<Array>,
873    equal_nan: impl Into<Option<bool>>,
874) -> Result<Array> {
875    a.as_ref()
876        .array_eq_result(b.as_ref(), equal_nan.into().unwrap_or(false))
877}
878
879/// Compatibility shim for [`array_eq`].
880#[generate_macro(customize(forwarding_shim = true))]
881#[deprecated(
882    since = "0.26.0",
883    note = "use `with_stream` or `with_device` around `array_eq`"
884)]
885pub fn array_eq_device(
886    a: impl AsRef<Array>,
887    b: impl AsRef<Array>,
888    #[optional] equal_nan: impl Into<Option<bool>>,
889    #[optional] stream: impl AsRef<Stream>,
890) -> Result<Array> {
891    crate::with_stream(stream.as_ref(), || array_eq(a, b, equal_nan))
892}
893
894/// See [`Array::eq`]
895pub fn eq(a: impl AsRef<Array>, b: impl AsRef<Array>) -> Result<Array> {
896    a.as_ref().eq(b)
897}
898
899/// Compatibility shim for [`eq`].
900#[generate_macro(customize(forwarding_shim = true))]
901#[deprecated(
902    since = "0.26.0",
903    note = "use `with_stream` or `with_device` around `eq`"
904)]
905pub fn eq_device(
906    a: impl AsRef<Array>,
907    b: impl AsRef<Array>,
908    #[optional] stream: impl AsRef<Stream>,
909) -> Result<Array> {
910    crate::with_stream(stream.as_ref(), || eq(a, b))
911}
912
913/// See [`Array::le`]
914pub fn le(a: impl AsRef<Array>, b: impl AsRef<Array>) -> Result<Array> {
915    a.as_ref().le(b)
916}
917
918/// Compatibility shim for [`le`].
919#[generate_macro(customize(forwarding_shim = true))]
920#[deprecated(
921    since = "0.26.0",
922    note = "use `with_stream` or `with_device` around `le`"
923)]
924pub fn le_device(
925    a: impl AsRef<Array>,
926    b: impl AsRef<Array>,
927    #[optional] stream: impl AsRef<Stream>,
928) -> Result<Array> {
929    crate::with_stream(stream.as_ref(), || le(a, b))
930}
931
932/// See [`Array::ge`]
933pub fn ge(a: impl AsRef<Array>, b: impl AsRef<Array>) -> Result<Array> {
934    a.as_ref().ge(b)
935}
936
937/// Compatibility shim for [`ge`].
938#[generate_macro(customize(forwarding_shim = true))]
939#[deprecated(
940    since = "0.26.0",
941    note = "use `with_stream` or `with_device` around `ge`"
942)]
943pub fn ge_device(
944    a: impl AsRef<Array>,
945    b: impl AsRef<Array>,
946    #[optional] stream: impl AsRef<Stream>,
947) -> Result<Array> {
948    crate::with_stream(stream.as_ref(), || ge(a, b))
949}
950
951/// See [`Array::ne`]
952pub fn ne(a: impl AsRef<Array>, b: impl AsRef<Array>) -> Result<Array> {
953    a.as_ref().ne(b)
954}
955
956/// Compatibility shim for [`ne`].
957#[generate_macro(customize(forwarding_shim = true))]
958#[deprecated(
959    since = "0.26.0",
960    note = "use `with_stream` or `with_device` around `ne`"
961)]
962pub fn ne_device(
963    a: impl AsRef<Array>,
964    b: impl AsRef<Array>,
965    #[optional] stream: impl AsRef<Stream>,
966) -> Result<Array> {
967    crate::with_stream(stream.as_ref(), || ne(a, b))
968}
969
970/// See [`Array::lt`]
971pub fn lt(a: impl AsRef<Array>, b: impl AsRef<Array>) -> Result<Array> {
972    a.as_ref().lt(b)
973}
974
975/// Compatibility shim for [`lt`].
976#[generate_macro(customize(forwarding_shim = true))]
977#[deprecated(
978    since = "0.26.0",
979    note = "use `with_stream` or `with_device` around `lt`"
980)]
981pub fn lt_device(
982    a: impl AsRef<Array>,
983    b: impl AsRef<Array>,
984    #[optional] stream: impl AsRef<Stream>,
985) -> Result<Array> {
986    crate::with_stream(stream.as_ref(), || lt(a, b))
987}
988
989/// See [`Array::gt`]
990pub fn gt(a: impl AsRef<Array>, b: impl AsRef<Array>) -> Result<Array> {
991    a.as_ref().gt(b)
992}
993
994/// Compatibility shim for [`gt`].
995#[generate_macro(customize(forwarding_shim = true))]
996#[deprecated(
997    since = "0.26.0",
998    note = "use `with_stream` or `with_device` around `gt`"
999)]
1000pub fn gt_device(
1001    a: impl AsRef<Array>,
1002    b: impl AsRef<Array>,
1003    #[optional] stream: impl AsRef<Stream>,
1004) -> Result<Array> {
1005    crate::with_stream(stream.as_ref(), || gt(a, b))
1006}
1007// TODO: check if the functions below could throw an exception.
1008
1009/// Return a boolean array indicating which elements are NaN.
1010pub fn is_nan(array: impl AsRef<Array>) -> Result<Array> {
1011    let stream = Stream::thread_local_or_default();
1012    Array::try_from_op(|res| unsafe {
1013        mlx_sys::mlx_isnan(res, array.as_ref().as_ptr(), stream.as_ref().as_ptr())
1014    })
1015}
1016
1017/// Compatibility shim for [`is_nan`].
1018#[generate_macro(customize(forwarding_shim = true))]
1019#[deprecated(
1020    since = "0.26.0",
1021    note = "use `with_stream` or `with_device` around `is_nan`"
1022)]
1023pub fn is_nan_device(
1024    array: impl AsRef<Array>,
1025    #[optional] stream: impl AsRef<Stream>,
1026) -> Result<Array> {
1027    crate::with_stream(stream.as_ref(), || is_nan(array))
1028}
1029
1030/// Return a boolean array indicating which elements are +/- inifnity.
1031pub fn is_inf(array: impl AsRef<Array>) -> Result<Array> {
1032    let stream = Stream::thread_local_or_default();
1033    Array::try_from_op(|res| unsafe {
1034        mlx_sys::mlx_isinf(res, array.as_ref().as_ptr(), stream.as_ref().as_ptr())
1035    })
1036}
1037
1038/// Compatibility shim for [`is_inf`].
1039#[generate_macro(customize(forwarding_shim = true))]
1040#[deprecated(
1041    since = "0.26.0",
1042    note = "use `with_stream` or `with_device` around `is_inf`"
1043)]
1044pub fn is_inf_device(
1045    array: impl AsRef<Array>,
1046    #[optional] stream: impl AsRef<Stream>,
1047) -> Result<Array> {
1048    crate::with_stream(stream.as_ref(), || is_inf(array))
1049}
1050
1051/// Return a boolean array indicating which elements are positive infinity.
1052pub fn is_pos_inf(array: impl AsRef<Array>) -> Result<Array> {
1053    let stream = Stream::thread_local_or_default();
1054    Array::try_from_op(|res| unsafe {
1055        mlx_sys::mlx_isposinf(res, array.as_ref().as_ptr(), stream.as_ref().as_ptr())
1056    })
1057}
1058
1059/// Compatibility shim for [`is_pos_inf`].
1060#[generate_macro(customize(forwarding_shim = true))]
1061#[deprecated(
1062    since = "0.26.0",
1063    note = "use `with_stream` or `with_device` around `is_pos_inf`"
1064)]
1065pub fn is_pos_inf_device(
1066    array: impl AsRef<Array>,
1067    #[optional] stream: impl AsRef<Stream>,
1068) -> Result<Array> {
1069    crate::with_stream(stream.as_ref(), || is_pos_inf(array))
1070}
1071
1072/// Return a boolean array indicating which elements are negative infinity.
1073pub fn is_neg_inf(array: impl AsRef<Array>) -> Result<Array> {
1074    let stream = Stream::thread_local_or_default();
1075    Array::try_from_op(|res| unsafe {
1076        mlx_sys::mlx_isneginf(res, array.as_ref().as_ptr(), stream.as_ref().as_ptr())
1077    })
1078}
1079
1080/// Compatibility shim for [`is_neg_inf`].
1081#[generate_macro(customize(forwarding_shim = true))]
1082#[deprecated(
1083    since = "0.26.0",
1084    note = "use `with_stream` or `with_device` around `is_neg_inf`"
1085)]
1086pub fn is_neg_inf_device(
1087    array: impl AsRef<Array>,
1088    #[optional] stream: impl AsRef<Stream>,
1089) -> Result<Array> {
1090    crate::with_stream(stream.as_ref(), || is_neg_inf(array))
1091}
1092
1093/// Select from `a` or `b` according to `condition` returning an error if the arrays are not
1094/// broadcastable.
1095///
1096/// The condition and input arrays must be the same shape or
1097/// [broadcasting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/broadcasting)
1098/// with each another.
1099///
1100/// # Params
1101///
1102/// - condition: condition array
1103/// - a: input selected from where condition is non-zero or `true`
1104/// - b: input selected from where condition is zero or `false`
1105pub fn select(
1106    condition: impl AsRef<Array>,
1107    a: impl AsRef<Array>,
1108    b: impl AsRef<Array>,
1109) -> Result<Array> {
1110    let stream = Stream::thread_local_or_default();
1111    Array::try_from_op(|res| unsafe {
1112        mlx_sys::mlx_where(
1113            res,
1114            condition.as_ref().as_ptr(),
1115            a.as_ref().as_ptr(),
1116            b.as_ref().as_ptr(),
1117            stream.as_ref().as_ptr(),
1118        )
1119    })
1120}
1121
1122/// Compatibility alias for [`select`].
1123#[deprecated(since = "0.26.0", note = "renamed to `select`")]
1124pub fn r#where(
1125    condition: impl AsRef<Array>,
1126    a: impl AsRef<Array>,
1127    b: impl AsRef<Array>,
1128) -> Result<Array> {
1129    select(condition, a, b)
1130}
1131
1132/// Compatibility shim for [`select`].
1133#[deprecated(
1134    since = "0.26.0",
1135    note = "use `with_stream` or `with_device` around `select`"
1136)]
1137pub fn r#where_device(
1138    condition: impl AsRef<Array>,
1139    a: impl AsRef<Array>,
1140    b: impl AsRef<Array>,
1141    stream: impl AsRef<Stream>,
1142) -> Result<Array> {
1143    crate::with_stream(stream.as_ref(), || select(condition, a, b))
1144}
1145
1146/// Compatibility alias for [`select`].
1147#[deprecated(since = "0.26.0", note = "renamed to `select`")]
1148pub fn which(
1149    condition: impl AsRef<Array>,
1150    a: impl AsRef<Array>,
1151    b: impl AsRef<Array>,
1152) -> Result<Array> {
1153    select(condition, a, b)
1154}
1155
1156/// Compatibility shim for [`select`].
1157#[generate_macro(customize(forwarding_shim = true))]
1158#[deprecated(
1159    since = "0.26.0",
1160    note = "use `with_stream` or `with_device` around `select`"
1161)]
1162pub fn which_device(
1163    condition: impl AsRef<Array>,
1164    a: impl AsRef<Array>,
1165    b: impl AsRef<Array>,
1166    #[optional] stream: impl AsRef<Stream>,
1167) -> Result<Array> {
1168    crate::with_stream(stream.as_ref(), || select(condition, a, b))
1169}
1170
1171#[cfg(test)]
1172mod tests {
1173    use crate::{array, Dtype};
1174
1175    use super::*;
1176
1177    #[test]
1178    fn test_eq() {
1179        let a = Array::from_slice(&[1, 2, 3], &[3]);
1180        let b = Array::from_slice(&[1, 2, 3], &[3]);
1181        let c = a.eq(&b).unwrap();
1182
1183        let c_data: &[bool] = c.as_slice();
1184        assert_eq!(c_data, [true, true, true]);
1185
1186        // check a and b are not modified
1187        let a_data: &[i32] = a.as_slice();
1188        assert_eq!(a_data, [1, 2, 3]);
1189
1190        let b_data: &[i32] = b.as_slice();
1191        assert_eq!(b_data, [1, 2, 3]);
1192    }
1193
1194    #[test]
1195    fn test_eq_invalid_broadcast() {
1196        let a = Array::from_slice(&[1, 2, 3], &[3]);
1197        let b = Array::from_slice(&[1, 2, 3, 4], &[4]);
1198        let c = a.eq(&b);
1199        assert!(c.is_err());
1200    }
1201
1202    #[test]
1203    fn test_le() {
1204        let a = Array::from_slice(&[1, 2, 3], &[3]);
1205        let b = Array::from_slice(&[1, 2, 3], &[3]);
1206        let c = a.le(&b).unwrap();
1207
1208        let c_data: &[bool] = c.as_slice();
1209        assert_eq!(c_data, [true, true, true]);
1210
1211        // check a and b are not modified
1212        let a_data: &[i32] = a.as_slice();
1213        assert_eq!(a_data, [1, 2, 3]);
1214
1215        let b_data: &[i32] = b.as_slice();
1216        assert_eq!(b_data, [1, 2, 3]);
1217    }
1218
1219    #[test]
1220    fn test_le_invalid_broadcast() {
1221        let a = Array::from_slice(&[1, 2, 3], &[3]);
1222        let b = Array::from_slice(&[1, 2, 3, 4], &[4]);
1223        let c = a.le(&b);
1224        assert!(c.is_err());
1225    }
1226
1227    #[test]
1228    fn test_ge() {
1229        let a = Array::from_slice(&[1, 2, 3], &[3]);
1230        let b = Array::from_slice(&[1, 2, 3], &[3]);
1231        let c = a.ge(&b).unwrap();
1232
1233        let c_data: &[bool] = c.as_slice();
1234        assert_eq!(c_data, [true, true, true]);
1235
1236        // check a and b are not modified
1237        let a_data: &[i32] = a.as_slice();
1238        assert_eq!(a_data, [1, 2, 3]);
1239
1240        let b_data: &[i32] = b.as_slice();
1241        assert_eq!(b_data, [1, 2, 3]);
1242    }
1243
1244    #[test]
1245    fn test_ge_invalid_broadcast() {
1246        let a = Array::from_slice(&[1, 2, 3], &[3]);
1247        let b = Array::from_slice(&[1, 2, 3, 4], &[4]);
1248        let c = a.ge(&b);
1249        assert!(c.is_err());
1250    }
1251
1252    #[test]
1253    fn test_ne() {
1254        let a = Array::from_slice(&[1, 2, 3], &[3]);
1255        let b = Array::from_slice(&[1, 2, 3], &[3]);
1256        let c = a.ne(&b).unwrap();
1257
1258        let c_data: &[bool] = c.as_slice();
1259        assert_eq!(c_data, [false, false, false]);
1260
1261        // check a and b are not modified
1262        let a_data: &[i32] = a.as_slice();
1263        assert_eq!(a_data, [1, 2, 3]);
1264
1265        let b_data: &[i32] = b.as_slice();
1266        assert_eq!(b_data, [1, 2, 3]);
1267    }
1268
1269    #[test]
1270    fn test_ne_invalid_broadcast() {
1271        let a = Array::from_slice(&[1, 2, 3], &[3]);
1272        let b = Array::from_slice(&[1, 2, 3, 4], &[4]);
1273        let c = a.ne(&b);
1274        assert!(c.is_err());
1275    }
1276
1277    #[test]
1278    fn test_lt() {
1279        let a = Array::from_slice(&[1, 0, 3], &[3]);
1280        let b = Array::from_slice(&[1, 2, 3], &[3]);
1281        let c = a.lt(&b).unwrap();
1282
1283        let c_data: &[bool] = c.as_slice();
1284        assert_eq!(c_data, [false, true, false]);
1285
1286        // check a and b are not modified
1287        let a_data: &[i32] = a.as_slice();
1288        assert_eq!(a_data, [1, 0, 3]);
1289
1290        let b_data: &[i32] = b.as_slice();
1291        assert_eq!(b_data, [1, 2, 3]);
1292    }
1293
1294    #[test]
1295    fn test_lt_invalid_broadcast() {
1296        let a = Array::from_slice(&[1, 2, 3], &[3]);
1297        let b = Array::from_slice(&[1, 2, 3, 4], &[4]);
1298        let c = a.lt(&b);
1299        assert!(c.is_err());
1300    }
1301
1302    #[test]
1303    fn test_gt() {
1304        let a = Array::from_slice(&[1, 4, 3], &[3]);
1305        let b = Array::from_slice(&[1, 2, 3], &[3]);
1306        let c = a.gt(&b).unwrap();
1307
1308        let c_data: &[bool] = c.as_slice();
1309        assert_eq!(c_data, [false, true, false]);
1310
1311        // check a and b are not modified
1312        let a_data: &[i32] = a.as_slice();
1313        assert_eq!(a_data, [1, 4, 3]);
1314
1315        let b_data: &[i32] = b.as_slice();
1316        assert_eq!(b_data, [1, 2, 3]);
1317    }
1318
1319    #[test]
1320    fn test_gt_invalid_broadcast() {
1321        let a = Array::from_slice(&[1, 2, 3], &[3]);
1322        let b = Array::from_slice(&[1, 2, 3, 4], &[4]);
1323        let c = a.gt(&b);
1324        assert!(c.is_err());
1325    }
1326
1327    #[test]
1328    fn test_logical_and() {
1329        let a = Array::from_slice(&[true, false, true], &[3]);
1330        let b = Array::from_slice(&[true, true, false], &[3]);
1331        let c = a.logical_and(&b).unwrap();
1332
1333        let c_data: &[bool] = c.as_slice();
1334        assert_eq!(c_data, [true, false, false]);
1335
1336        // check a and b are not modified
1337        let a_data: &[bool] = a.as_slice();
1338        assert_eq!(a_data, [true, false, true]);
1339
1340        let b_data: &[bool] = b.as_slice();
1341        assert_eq!(b_data, [true, true, false]);
1342    }
1343
1344    #[test]
1345    fn test_logical_and_invalid_broadcast() {
1346        let a = Array::from_slice(&[true, false, true], &[3]);
1347        let b = Array::from_slice(&[true, true, false, true], &[4]);
1348        let c = a.logical_and(&b);
1349        assert!(c.is_err());
1350    }
1351
1352    #[test]
1353    fn test_logical_or() {
1354        let a = Array::from_slice(&[true, false, true], &[3]);
1355        let b = Array::from_slice(&[true, true, false], &[3]);
1356        let c = a.logical_or(&b).unwrap();
1357
1358        let c_data: &[bool] = c.as_slice();
1359        assert_eq!(c_data, [true, true, true]);
1360
1361        // check a and b are not modified
1362        let a_data: &[bool] = a.as_slice();
1363        assert_eq!(a_data, [true, false, true]);
1364
1365        let b_data: &[bool] = b.as_slice();
1366        assert_eq!(b_data, [true, true, false]);
1367    }
1368
1369    #[test]
1370    fn test_logical_or_invalid_broadcast() {
1371        let a = Array::from_slice(&[true, false, true], &[3]);
1372        let b = Array::from_slice(&[true, true, false, true], &[4]);
1373        let c = a.logical_or(&b);
1374        assert!(c.is_err());
1375    }
1376
1377    #[test]
1378    fn test_all_close() {
1379        let a = Array::from_slice(&[0., 1., 2., 3.], &[4]).sqrt().unwrap();
1380        let b = Array::from_slice(&[0., 1., 2., 3.], &[4])
1381            .power(array!(0.5))
1382            .unwrap();
1383        let c = a.all_close(&b, 1e-5, None, None).unwrap();
1384        assert!(c);
1385    }
1386
1387    #[test]
1388    fn test_all_close_invalid_broadcast() {
1389        let a = Array::from_slice(&[0., 1., 2., 3.], &[4]);
1390        let b = Array::from_slice(&[0., 1., 2., 3., 4.], &[5]);
1391        let c = a.all_close(&b, 1e-5, None, None);
1392        assert!(c.is_err());
1393    }
1394
1395    #[test]
1396    fn test_is_close_false() {
1397        let a = Array::from_slice(&[1., 2., 3.], &[3]);
1398        let b = Array::from_slice(&[1.1, 2.2, 3.3], &[3]);
1399        let c = a.is_close(&b, None, None, false).unwrap();
1400
1401        let c_data: &[bool] = c.as_slice();
1402        assert_eq!(c_data, [false, false, false]);
1403    }
1404
1405    #[test]
1406    fn test_is_close_true() {
1407        let a = Array::from_slice(&[1., 2., 3.], &[3]);
1408        let b = Array::from_slice(&[1.1, 2.2, 3.3], &[3]);
1409        let c = a.is_close(&b, 0.1, 0.2, true).unwrap();
1410
1411        let c_data: &[bool] = c.as_slice();
1412        assert_eq!(c_data, [true, true, true]);
1413    }
1414
1415    #[test]
1416    fn test_is_close_invalid_broadcast() {
1417        let a = Array::from_slice(&[1., 2., 3.], &[3]);
1418        let b = Array::from_slice(&[1.1, 2.2, 3.3, 4.4], &[4]);
1419        let c = a.is_close(&b, None, None, false);
1420        assert!(c.is_err());
1421    }
1422
1423    #[test]
1424    fn test_array_eq() {
1425        let a = Array::from_slice(&[0, 1, 2, 3], &[4]);
1426        let b = Array::from_slice(&[0., 1., 2., 3.], &[4]);
1427        assert!(a.eq_values(&b).unwrap());
1428    }
1429
1430    #[test]
1431    fn test_any() {
1432        let array = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], &[3, 4]);
1433        let all = array.any_axes(&[0][..], None).unwrap();
1434
1435        let results: &[bool] = all.as_slice();
1436        assert_eq!(results, &[true, true, true, true]);
1437    }
1438
1439    #[test]
1440    fn test_any_empty_axes() {
1441        let array = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], &[3, 4]);
1442        let all = array.any_axes(&[][..], None).unwrap();
1443
1444        let results: &[bool] = all.as_slice();
1445        assert_eq!(
1446            results,
1447            &[false, true, true, true, true, true, true, true, true, true, true, true]
1448        );
1449    }
1450
1451    #[test]
1452    fn test_any_out_of_bounds() {
1453        let array = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], &[12]);
1454        let result = array.any_axes(&[1][..], None);
1455        assert!(result.is_err());
1456    }
1457
1458    #[test]
1459    fn test_any_duplicate_axes() {
1460        let array = Array::from_slice(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], &[3, 4]);
1461        let result = array.any_axes(&[0, 0][..], None);
1462        assert!(result.is_err());
1463    }
1464
1465    #[test]
1466    fn test_select() {
1467        let condition = Array::from_slice(&[true, false, true], &[3]);
1468        let a = Array::from_slice(&[1, 2, 3], &[3]);
1469        let b = Array::from_slice(&[4, 5, 6], &[3]);
1470        let c = select(&condition, &a, &b).unwrap();
1471
1472        let c_data: &[i32] = c.as_slice();
1473        assert_eq!(c_data, [1, 5, 3]);
1474    }
1475
1476    #[test]
1477    fn test_select_invalid_broadcast() {
1478        let condition = Array::from_slice(&[true, false, true], &[3]);
1479        let a = Array::from_slice(&[1, 2, 3], &[3]);
1480        let b = Array::from_slice(&[4, 5, 6, 7], &[4]);
1481        let c = select(&condition, &a, &b);
1482        assert!(c.is_err());
1483    }
1484
1485    // The unit tests below are adapted from the mlx c++ codebase
1486
1487    #[test]
1488    fn test_unary_logical_not() {
1489        let x = array!(false);
1490        assert!(logical_not(&x).unwrap().item_exact::<bool>());
1491
1492        let x = array!(1.0);
1493        let y = logical_not(&x).unwrap();
1494        assert_eq!(y.dtype(), Dtype::Bool);
1495        assert!(!y.item_exact::<bool>());
1496
1497        let x = array!(0);
1498        let y = logical_not(&x).unwrap();
1499        assert_eq!(y.dtype(), Dtype::Bool);
1500        assert!(y.item_exact::<bool>());
1501    }
1502
1503    #[test]
1504    fn test_unary_logical_and() {
1505        let x = array!(true);
1506        let y = array!(true);
1507        assert!(logical_and(&x, &y).unwrap().item_exact::<bool>());
1508
1509        let x = array!(1.0);
1510        let y = array!(1.0);
1511        let z = logical_and(&x, &y).unwrap();
1512        assert_eq!(z.dtype(), Dtype::Bool);
1513        assert!(z.item_exact::<bool>());
1514
1515        let x = array!(0);
1516        let y = array!(1.0);
1517        let z = logical_and(&x, &y).unwrap();
1518        assert_eq!(z.dtype(), Dtype::Bool);
1519        assert!(!z.item_exact::<bool>());
1520    }
1521
1522    #[test]
1523    fn test_unary_logical_or() {
1524        let a = array!(false);
1525        let b = array!(false);
1526        assert!(!logical_or(&a, &b).unwrap().item_exact::<bool>());
1527
1528        let a = array!(1.0);
1529        let b = array!(1.0);
1530        let c = logical_or(&a, &b).unwrap();
1531        assert_eq!(c.dtype(), Dtype::Bool);
1532        assert!(c.item_exact::<bool>());
1533
1534        let a = array!(0);
1535        let b = array!(1.0);
1536        let c = logical_or(&a, &b).unwrap();
1537        assert_eq!(c.dtype(), Dtype::Bool);
1538        assert!(c.item_exact::<bool>());
1539    }
1540}
mlx_rs/ops/logical.rs

mlx_rs/ops/
logical.rs
