add trait implementations for binary assign operations for tensor and f32

Author: PaytonWebber

tensor/src/lib.rs

@@ -711,6 +711,14 @@ impl AddAssign<&Tensor> for Tensor {
     }
 }
 
+impl AddAssign<f32> for Tensor {
+    fn add_assign(&mut self, rhs: f32) {
+        self.data.iter_mut().for_each(|a| {
+            *a += rhs;
+        });
+    }
+}
+
 impl Sub<&Tensor> for &Tensor {
     type Output = Tensor;
 
@@ -754,6 +762,14 @@ impl SubAssign<&Tensor> for Tensor {
     }
 }
 
+impl SubAssign<f32> for Tensor {
+    fn sub_assign(&mut self, rhs: f32) {
+        self.data.iter_mut().for_each(|a| {
+            *a -= rhs;
+        });
+    }
+}
+
 impl Mul<&Tensor> for &Tensor {
     type Output = Tensor;
 
@@ -797,6 +813,14 @@ impl MulAssign<&Tensor> for Tensor {
     }
 }
 
+impl MulAssign<f32> for Tensor {
+    fn mul_assign(&mut self, rhs: f32) {
+        self.data.iter_mut().for_each(|a| {
+            *a *= rhs;
+        });
+    }
+}
+
 impl Div<&Tensor> for &Tensor {
     type Output = Tensor;
 
@@ -830,3 +854,11 @@ impl DivAssign<&Tensor> for Tensor {
             });
     }
 }
+
+impl DivAssign<f32> for Tensor {
+    fn div_assign(&mut self, rhs: f32) {
+        self.data.iter_mut().for_each(|a| {
+            *a /= rhs;
+        });
+    }
+}

tensor/tests/tensor_core_test.rs

@@ -387,10 +387,14 @@ fn tensor_add_inplace_operator() {
 
     let mut a = Tensor::new(a_shape, a_data).unwrap();
     let b = Tensor::new(b_shape, b_data).unwrap();
-    a += &b;
 
+    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());
+
+    a += 4.0;
+    let expected_data: Vec<f32> = vec![6.0, 8.0, 10.0, 12.0, 14.0, 16.0];
+    assert_eq!(expected_data, *a.data());
 }
 
 #[test]
@@ -487,10 +491,14 @@ fn tensor_sub_inplace_operator() {
 
     let mut a = Tensor::new(a_shape, a_data).unwrap();
     let b = Tensor::new(b_shape, b_data).unwrap();
-    a -= &b;
 
+    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());
+
+    a -= 2.0;
+    let expected_data: Vec<f32> = vec![-3.0, -2.0, -1.0, 0.0, 1.0, 2.0];
+    assert_eq!(expected_data, *a.data());
 }
 
 #[test]
@@ -581,10 +589,14 @@ fn tensor_mul_inplace_operator() {
 
     let mut a = Tensor::new(a_shape, a_data).unwrap();
     let b = Tensor::new(b_shape, b_data).unwrap();
-    a *= &b;
 
+    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());
+
+    a *= 2.0;
+    let expected_data: Vec<f32> = vec![4.0, 8.0, 12.0, 16.0, 20.0, 24.0];
+    assert_eq!(expected_data, *a.data());
 }
 
 #[test]
@@ -672,10 +684,14 @@ fn tensor_div_inplace_operator() {
 
     let mut a = Tensor::new(a_shape, a_data).unwrap();
     let b = Tensor::new(b_shape, b_data).unwrap();
-    a /= &b;
 
+    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());
+
+    a /= 0.5;
+    let expected_data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+    assert_eq!(expected_data, *a.data());
 }
 
 #[test]