add trait implementations for binary assign operations for tensors

Author: PaytonWebber

tensor/src/lib.rs

@@ -1,5 +1,5 @@
 use core::{f32, fmt};
-use std::ops::{Add, Div, Index, Mul, Sub};
+use std::ops::{Add, AddAssign, Div, DivAssign, Index, Mul, MulAssign, Sub, SubAssign};
 
 #[derive(Debug, Clone)]
 pub struct Tensor {
@@ -244,7 +244,7 @@ impl Tensor {
         let self_shape = self.shape();
         let other_shape = other.shape();
         if self_shape != other_shape {
-            return;
+            panic!("The tensor shape not compatible for inplace addition")
         }
         self.data
             .iter_mut()
@@ -303,7 +303,7 @@ impl Tensor {
         let self_shape = self.shape();
         let other_shape = other.shape();
         if self_shape != other_shape {
-            return;
+            panic!("The tensor shape not compatible for inplace subtraction")
         }
         self.data
             .iter_mut()
@@ -361,7 +361,7 @@ impl Tensor {
         let self_shape = self.shape();
         let other_shape = other.shape();
         if self_shape != other_shape {
-            return;
+            panic!("The tensor shape not compatible for inplace multiplication")
         }
         self.data
             .iter_mut()
@@ -419,7 +419,7 @@ impl Tensor {
         let self_shape = self.shape();
         let other_shape = other.shape();
         if self_shape != other_shape {
-            return;
+            panic!("The tensor shape not compatible for inplace division")
         }
         self.data
             .iter_mut()
@@ -697,6 +697,20 @@ impl Add<&Tensor> for f32 {
     }
 }
 
+impl AddAssign<&Tensor> for Tensor {
+    fn add_assign(&mut self, rhs: &Tensor) {
+        if *self.shape() != *rhs.shape() {
+            panic!("The tensor shape not compatible for inplace addition")
+        }
+        self.data
+            .iter_mut()
+            .zip(rhs.data().iter())
+            .for_each(|(a, b)| {
+                *a += b;
+            });
+    }
+}
+
 impl Sub<&Tensor> for &Tensor {
     type Output = Tensor;
 
@@ -726,6 +740,20 @@ impl Sub<&Tensor> for f32 {
     }
 }
 
+impl SubAssign<&Tensor> for Tensor {
+    fn sub_assign(&mut self, rhs: &Tensor) {
+        if *self.shape() != *rhs.shape() {
+            panic!("The tensor shape not compatible for inplace subtraction")
+        }
+        self.data
+            .iter_mut()
+            .zip(rhs.data().iter())
+            .for_each(|(a, b)| {
+                *a -= b;
+            });
+    }
+}
+
 impl Mul<&Tensor> for &Tensor {
     type Output = Tensor;
 
@@ -755,6 +783,20 @@ impl Mul<&Tensor> for f32 {
     }
 }
 
+impl MulAssign<&Tensor> for Tensor {
+    fn mul_assign(&mut self, rhs: &Tensor) {
+        if *self.shape() != *rhs.shape() {
+            panic!("The tensor shape not compatible for inplace subtraction")
+        }
+        self.data
+            .iter_mut()
+            .zip(rhs.data().iter())
+            .for_each(|(a, b)| {
+                *a *= b;
+            });
+    }
+}
+
 impl Div<&Tensor> for &Tensor {
     type Output = Tensor;
 
@@ -774,3 +816,17 @@ impl Div<f32> for &Tensor {
         Tensor::new(self.shape().clone(), result_data).unwrap()
     }
 }
+
+impl DivAssign<&Tensor> for Tensor {
+    fn div_assign(&mut self, rhs: &Tensor) {
+        if *self.shape() != *rhs.shape() {
+            panic!("The tensor shape not compatible for inplace subtraction")
+        }
+        self.data
+            .iter_mut()
+            .zip(rhs.data().iter())
+            .for_each(|(a, b)| {
+                *a /= b;
+            });
+    }
+}

tensor/tests/tensor_core_test.rs

@@ -365,7 +365,7 @@ fn tensor_broadcasted_addition_operator() {
 }
 
 #[test]
-fn tensor_add_inplace() {
+fn tensor_add_inplace_method() {
     let a_shape = vec![2, 3];
     let b_shape = vec![2, 3];
     let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
@@ -378,6 +378,21 @@ fn tensor_add_inplace() {
     assert_eq!(expected_data, *a.data());
 }
 
+#[test]
+fn tensor_add_inplace_operator() {
+    let a_shape = vec![2, 3];
+    let b_shape = vec![2, 3];
+    let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+    let b_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+
+    let mut a = Tensor::new(a_shape, a_data).unwrap();
+    let b = Tensor::new(b_shape, b_data).unwrap();
+    a += &b;
+
+    let expected_data: Vec<f32> = vec![2.0, 4.0, 6.0, 8.0, 10.0, 12.0];
+    assert_eq!(expected_data, *a.data());
+}
+
 #[test]
 fn tensor_subtraction_method() {
     let shape = vec![4, 2];
@@ -450,7 +465,7 @@ fn tensor_broadcasted_subtraction_operator() {
 }
 
 #[test]
-fn tensor_sub_inplace() {
+fn tensor_sub_inplace_method() {
     let a_shape = vec![2, 3];
     let b_shape = vec![2, 3];
     let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
@@ -463,6 +478,21 @@ fn tensor_sub_inplace() {
     assert_eq!(expected_data, *a.data());
 }
 
+#[test]
+fn tensor_sub_inplace_operator() {
+    let a_shape = vec![2, 3];
+    let b_shape = vec![2, 3];
+    let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+    let b_data: Vec<f32> = vec![2.0, 2.0, 2.0, 2.0, 2.0, 2.0];
+
+    let mut a = Tensor::new(a_shape, a_data).unwrap();
+    let b = Tensor::new(b_shape, b_data).unwrap();
+    a -= &b;
+
+    let expected_data: Vec<f32> = vec![-1.0, 0.0, 1.0, 2.0, 3.0, 4.0];
+    assert_eq!(expected_data, *a.data());
+}
+
 #[test]
 fn tensor_mul_method() {
     let a_shape = vec![1, 3];
@@ -528,7 +558,7 @@ fn tensor_broadcasted_mul_method() {
 }
 
 #[test]
-fn tensor_mul_inplace() {
+fn tensor_mul_inplace_method() {
     let a_shape = vec![2, 3];
     let b_shape = vec![2, 3];
     let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
@@ -542,6 +572,21 @@ fn tensor_mul_inplace() {
     assert_eq!(expected_data, *a.data());
 }
 
+#[test]
+fn tensor_mul_inplace_operator() {
+    let a_shape = vec![2, 3];
+    let b_shape = vec![2, 3];
+    let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+    let b_data: Vec<f32> = vec![2.0, 2.0, 2.0, 2.0, 2.0, 2.0];
+
+    let mut a = Tensor::new(a_shape, a_data).unwrap();
+    let b = Tensor::new(b_shape, b_data).unwrap();
+    a *= &b;
+
+    let expected_data: Vec<f32> = vec![2.0, 4.0, 6.0, 8.0, 10.0, 12.0];
+    assert_eq!(expected_data, *a.data());
+}
+
 #[test]
 fn tensor_div_method() {
     let a_shape = vec![1, 3];
@@ -604,7 +649,7 @@ fn tensor_broadcasted_div_method() {
 }
 
 #[test]
-fn tensor_div_inplace() {
+fn tensor_div_inplace_method() {
     let a_shape = vec![2, 3];
     let b_shape = vec![2, 3];
     let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
@@ -618,6 +663,21 @@ fn tensor_div_inplace() {
     assert_eq!(expected_data, *a.data());
 }
 
+#[test]
+fn tensor_div_inplace_operator() {
+    let a_shape = vec![2, 3];
+    let b_shape = vec![2, 3];
+    let a_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+    let b_data: Vec<f32> = vec![2.0, 2.0, 2.0, 2.0, 2.0, 2.0];
+
+    let mut a = Tensor::new(a_shape, a_data).unwrap();
+    let b = Tensor::new(b_shape, b_data).unwrap();
+    a /= &b;
+
+    let expected_data: Vec<f32> = vec![0.5, 1.0, 1.5, 2.0, 2.5, 3.0];
+    assert_eq!(expected_data, *a.data());
+}
+
 #[test]
 fn tensor_matmul() {
     // A is 2x3: