mlx_rs/array/
operators.rs

1use crate::{utils::ScalarOrArray, Array};
2use num_traits::Pow;
3use std::{
4    iter::Product,
5    ops::{
6        Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Not, Rem, RemAssign, Sub, SubAssign,
7    },
8};
9
10macro_rules! impl_binary_op {
11    ($trait:ident, $method:ident, $c_method:ident) => {
12        impl<'a, T> $trait<T> for Array
13        where
14            T: ScalarOrArray<'a>,
15        {
16            type Output = Array;
17
18            fn $method(self, rhs: T) -> Self::Output {
19                Array::$c_method(&self, rhs.into_owned_or_ref_array()).unwrap()
20            }
21        }
22
23        impl<'a, 't: 'a, T> $trait<T> for &'a Array
24        where
25            T: ScalarOrArray<'t>,
26        {
27            type Output = Array;
28
29            fn $method(self, rhs: T) -> Self::Output {
30                Array::$c_method(self, rhs.into_owned_or_ref_array()).unwrap()
31            }
32        }
33    };
34}
35
36macro_rules! impl_binary_op_assign {
37    ($trait:ident, $method:ident, $c_method:ident) => {
38        impl<T: Into<Array>> $trait<T> for Array {
39            fn $method(&mut self, rhs: T) {
40                let new_array = Array::$c_method(self, &rhs.into()).unwrap();
41                *self = new_array;
42            }
43        }
44
45        impl $trait<&Array> for Array {
46            fn $method(&mut self, rhs: &Self) {
47                let new_array = Array::$c_method(self, rhs).unwrap();
48                *self = new_array;
49            }
50        }
51    };
52}
53
54impl_binary_op!(Add, add, add);
55impl_binary_op_assign!(AddAssign, add_assign, add);
56impl_binary_op!(Sub, sub, subtract);
57impl_binary_op_assign!(SubAssign, sub_assign, subtract);
58impl_binary_op!(Mul, mul, multiply);
59impl_binary_op_assign!(MulAssign, mul_assign, multiply);
60impl_binary_op!(Div, div, divide);
61impl_binary_op_assign!(DivAssign, div_assign, divide);
62impl_binary_op!(Rem, rem, remainder);
63impl_binary_op_assign!(RemAssign, rem_assign, remainder);
64impl_binary_op!(Pow, pow, power);
65
66impl Neg for &Array {
67    type Output = Array;
68    fn neg(self) -> Self::Output {
69        Array::negative(self).unwrap()
70    }
71}
72impl Neg for Array {
73    type Output = Array;
74    fn neg(self) -> Self::Output {
75        Array::negative(&self).unwrap()
76    }
77}
78
79impl Not for &Array {
80    type Output = Array;
81    fn not(self) -> Self::Output {
82        Array::logical_not(self).unwrap()
83    }
84}
85impl Not for Array {
86    type Output = Array;
87    fn not(self) -> Self::Output {
88        Array::logical_not(&self).unwrap()
89    }
90}
91
92impl Product<Array> for Array {
93    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
94        iter.fold(1.0.into(), |acc, x| acc * x)
95    }
96}
97
98impl<'a> Product<&'a Array> for Array {
99    fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
100        iter.fold(1.0.into(), |acc, x| acc * x)
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use pretty_assertions::assert_eq;
108
109    #[test]
110    fn test_add_assign() {
111        let mut a = Array::from_slice(&[1.0, 2.0, 3.0], &[3]);
112        let b = Array::from_slice(&[4.0, 5.0, 6.0], &[3]);
113        a += &b;
114
115        assert_eq!(a.as_slice::<f32>(), &[5.0, 7.0, 9.0]);
116    }
117
118    #[test]
119    fn test_sub_assign() {
120        let mut a = Array::from_slice(&[1.0, 2.0, 3.0], &[3]);
121        let b = Array::from_slice(&[4.0, 5.0, 6.0], &[3]);
122        a -= &b;
123
124        assert_eq!(a.as_slice::<f32>(), &[-3.0, -3.0, -3.0]);
125    }
126
127    #[test]
128    fn test_mul_assign() {
129        let mut a = Array::from_slice(&[1.0, 2.0, 3.0], &[3]);
130        let b = Array::from_slice(&[4.0, 5.0, 6.0], &[3]);
131        a *= &b;
132
133        assert_eq!(a.as_slice::<f32>(), &[4.0, 10.0, 18.0]);
134    }
135
136    #[test]
137    fn test_div_assign() {
138        let mut a = Array::from_slice(&[1.0, 2.0, 3.0], &[3]);
139        let b = Array::from_slice(&[4.0, 5.0, 6.0], &[3]);
140        a /= &b;
141
142        assert_eq!(a.as_slice::<f32>(), &[0.25, 0.4, 0.5]);
143    }
144}
mlx_rs/array/operators.rs

mlx_rs/array/
operators.rs
