feat(tensor): add grad field, create <set|zero>_grad methods

src/linalg/tensor.rs

@@ -14,21 +14,24 @@ mod macros;
 mod manipulation;
 mod ops;
 mod sparse;
+mod state;
 mod transformation;
 mod validation;
 
-/// A one-dimensional matrix of floating point values.
+/// A one-dimensional tensor of floating point values.
 pub type Tensor1D = Vec<f64>;
-/// A two-dimensional matrix of floating point values.
+/// A two-dimensional tensor of floating point values.
 pub type Tensor2D = Vec<Tensor1D>;
 
-/// A matrix of floating point values, represented as a two-dimensional vector.
+/// A tensor of floating point values.
 #[derive(Clone, Debug)]
 pub struct Tensor {
-    /// The number of rows in the matrix.
+    /// The number of rows in the tensor.
     pub rows: usize,
-    /// The number of columns in the matrix.
+    /// The number of columns in the tensor.
     pub cols: usize,
-    /// The data in the matrix, represented as a two-dimensional vector.
+    /// The data in the tensor, represented as a two-dimensional vector.
     pub data: Tensor2D,
+    /// Gradient of the tensor.
+    pub grad: Option<Tensor2D>,
 }

src/linalg/tensor/computation.rs

@@ -49,6 +49,7 @@ impl Tensor {
             rows: self.rows,
             cols: self.cols,
             data,
+            grad: self.grad.clone(),
         }
     }
 

src/linalg/tensor/creation.rs

@@ -16,6 +16,7 @@ impl Tensor {
             rows,
             cols,
             data: vec![vec![0.0; cols]; rows],
+            grad: None,
         }
     }
 
@@ -64,6 +65,7 @@ impl Tensor {
             rows,
             cols,
             data: vec![vec![1.0; cols]; rows],
+            grad: None,
         }
     }
 

src/linalg/tensor/from.rs

@@ -6,6 +6,18 @@ impl From<Tensor2D> for Tensor {
             rows: data.len(),
             cols: data[0].len(),
             data,
+            grad: None,
+        }
+    }
+}
+
+impl From<&Tensor2D> for Tensor {
+    fn from(data: &Tensor2D) -> Self {
+        Tensor {
+            rows: data.len(),
+            cols: data[0].len(),
+            data: data.to_owned(),
+            grad: None,
         }
     }
 }

src/linalg/tensor/iter.rs

@@ -42,6 +42,11 @@ impl FromIterator<Tensor1D> for Tensor {
         }
         let rows = data.len();
         let cols = data[0].len();
-        Tensor { rows, cols, data }
+        Tensor {
+            rows,
+            cols,
+            data,
+            grad: None,
+        }
     }
 }

src/linalg/tensor/manipulation.rs

@@ -123,6 +123,7 @@ impl Tensor {
             rows: end - start,
             cols: self.cols,
             data: self.data[start..end].to_vec(),
+            grad: None,
         }
     }
 }

src/linalg/tensor/state.rs

@@ -0,0 +1,42 @@
+use crate::Tensor;
+
+impl Tensor {
+    /// Sets the gradient of the tensor.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use engram::*;
+    /// let mut a = tensor![[1.0, 2.0], [3.0, 4.0]];
+    /// a.set_grad(vec![vec![1.0, 2.0], vec![3.0, 4.0]]);
+    /// assert_eq!(a.grad, Some(vec![vec![1.0, 2.0], vec![3.0, 4.0]]));
+    /// ```
+    pub fn set_grad(&mut self, grad: Vec<Vec<f64>>) {
+        self.grad = Some(grad);
+    }
+
+    /// Zeros out the gradient of the tensor.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use engram::*;
+    /// let mut a = tensor![[1.0, 2.0], [3.0, 4.0]];
+    /// a.zero_grad();
+    /// assert_eq!(a.grad, Some(vec![vec![0.0, 0.0], vec![0.0, 0.0]]));
+    /// ```
+    pub fn zero_grad(&mut self) {
+        match &mut self.grad {
+            Some(grad) => {
+                for row in grad.iter_mut() {
+                    for val in row.iter_mut() {
+                        *val = 0.0;
+                    }
+                }
+            }
+            None => {
+                self.grad = Some(vec![vec![0.0; self.cols]; self.rows]);
+            }
+        }
+    }
+}