Merge pull request #12 from PaytonWebber/feature/tensor-inplace-ops

Feature/tensor inplace ops

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 {
@@ -240,6 +240,20 @@ impl Tensor {
         Tensor::new(bc_shape, result_data)
     }
 
+    pub fn add_inplace(&mut self, other: &Tensor) {
+        let self_shape = self.shape();
+        let other_shape = other.shape();
+        if self_shape != other_shape {
+            panic!("The tensor shape not compatible for inplace addition")
+        }
+        self.data
+            .iter_mut()
+            .zip(other.data().iter())
+            .for_each(|(a, &b)| {
+                *a += b;
+            });
+    }
+
     pub fn sub(&self, other: &Tensor) -> Result<Tensor, &'static str> {
         let self_shape = self.shape();
         let other_shape = other.shape();
@@ -285,6 +299,20 @@ impl Tensor {
         Tensor::new(bc_shape, result_data)
     }
 
+    pub fn sub_inplace(&mut self, other: &Tensor) {
+        let self_shape = self.shape();
+        let other_shape = other.shape();
+        if self_shape != other_shape {
+            panic!("The tensor shape not compatible for inplace subtraction")
+        }
+        self.data
+            .iter_mut()
+            .zip(other.data().iter())
+            .for_each(|(a, &b)| {
+                *a -= b;
+            });
+    }
+
     pub fn mul(&self, other: &Tensor) -> Result<Tensor, &'static str> {
         let self_shape = self.shape();
         let other_shape = other.shape();
@@ -329,6 +357,20 @@ impl Tensor {
         Tensor::new(bc_shape, result_data)
     }
 
+    pub fn mul_inplace(&mut self, other: &Tensor) {
+        let self_shape = self.shape();
+        let other_shape = other.shape();
+        if self_shape != other_shape {
+            panic!("The tensor shape not compatible for inplace multiplication")
+        }
+        self.data
+            .iter_mut()
+            .zip(other.data().iter())
+            .for_each(|(a, &b)| {
+                *a *= b;
+            });
+    }
+
     pub fn div(&self, other: &Tensor) -> Result<Tensor, &'static str> {
         let self_shape = self.shape();
         let other_shape = other.shape();
@@ -373,6 +415,20 @@ impl Tensor {
         Tensor::new(bc_shape, result_data)
     }
 
+    pub fn div_inplace(&mut self, other: &Tensor) {
+        let self_shape = self.shape();
+        let other_shape = other.shape();
+        if self_shape != other_shape {
+            panic!("The tensor shape not compatible for inplace division")
+        }
+        self.data
+            .iter_mut()
+            .zip(other.data().iter())
+            .for_each(|(a, &b)| {
+                *a /= b;
+            });
+    }
+
     pub fn matmul(&self, other: &Tensor) -> Result<Tensor, &'static str> {
         let lhs_shape: &Vec<usize> = self.shape();
         let rhs_shape: &Vec<usize> = other.shape();
@@ -447,7 +503,7 @@ fn unravel_index(mut i: usize, shape: &[usize]) -> Vec<usize> {
 
 pub fn is_broadcastable(a: &Vec<usize>, b: &Vec<usize>) -> bool {
     // This is based on NumPy's rules: https://numpy.org/doc/stable/user/basics.broadcasting.html
-    for (i, j) in a.into_iter().rev().zip(b.into_iter().rev()) {
+    for (i, j) in a.iter().rev().zip(b.iter().rev()) {
         if *i == 1 || *j == 1 {
             continue;
         }
@@ -480,26 +536,26 @@ pub fn compute_broadcast_shape_and_strides(
             if dim_a != dim_b {
                 a_bc_strides[ndims - i - 1] = match dim_a {
                     1 => 0,
-                    _ => a_dims[ndims - i..].into_iter().product(),
+                    _ => a_dims[ndims - i..].iter().product(),
                 };
                 b_bc_strides[ndims - i - 1] = match dim_b {
                     1 => 0,
-                    _ => b_dims[ndims - i..].into_iter().product(),
+                    _ => b_dims[ndims - i..].iter().product(),
                 };
             }
         } else {
             if dim_a != dim_b {
                 a_bc_strides[ndims - i - 1] = match dim_a {
                     1 => 0,
-                    _ => a_dims[ndims - i..].into_iter().product(),
+                    _ => a_dims[ndims - i..].iter().product(),
                 };
                 b_bc_strides[ndims - i - 1] = match dim_b {
                     1 => 0,
-                    _ => b_dims[ndims - i..].into_iter().product(),
+                    _ => b_dims[ndims - i..].iter().product(),
                 };
             } else {
-                a_bc_strides[ndims - i - 1] = a_dims[ndims - i..].into_iter().product();
-                b_bc_strides[ndims - i - 1] = b_dims[ndims - i..].into_iter().product();
+                a_bc_strides[ndims - i - 1] = a_dims[ndims - i..].iter().product();
+                b_bc_strides[ndims - i - 1] = b_dims[ndims - i..].iter().product();
             }
         }
         bc_shape[ndims - i - 1] = dim_a.max(dim_b);
@@ -641,6 +697,28 @@ 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 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;
 
@@ -670,6 +748,28 @@ 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 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;
 
@@ -699,6 +799,28 @@ 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 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;
 
@@ -718,3 +840,25 @@ 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;
+            });
+    }
+}
+
+impl DivAssign<f32> for Tensor {
+    fn div_assign(&mut self, rhs: f32) {
+        self.data.iter_mut().for_each(|a| {
+            *a /= rhs;
+        });
+    }
+}