Add generic numeric type so that tensors can be created with f32 and f64 (#21)

Author: PaytonWebber

tensor/examples/generic_test.rs

@@ -0,0 +1,18 @@
+use tensor::{Tensor, TensorF32, TensorF64};
+
+fn main() {
+    // Test f32 tensors
+    let a = TensorF32::ones(vec![2, 2]);
+    let b = TensorF32::zeros(vec![2, 2]);
+    println!("f32 tensor:\n{}", &a + &b);
+
+    // Test f64 tensors
+    let a_f64 = TensorF64::ones(vec![2, 2]);
+    let b_f64 = TensorF64::zeros(vec![2, 2]);
+    println!("f64 tensor:\n{}", &a_f64 + &b_f64);
+
+    // Test with explicit generic syntax
+    let a_explicit: Tensor<f32> = Tensor::ones(vec![2, 2]);
+    let b_explicit: Tensor<f32> = Tensor::new(vec![2, 2], vec![2.5, 3.0, 1.5, 4.0]).unwrap();
+    println!("Explicit f32 tensor:\n{}", &a_explicit + &b_explicit);
+}

tensor/src/core/mod.rs

@@ -4,16 +4,97 @@ mod traits;
 mod utils;
 
 use errors::TensorError;
+use std::fmt;
 use utils::calculate_strides;
 
-pub struct Tensor {
+pub trait Numeric:
+    Copy
+    + fmt::Display
+    + fmt::Debug
+    + std::ops::Add<Output = Self>
+    + std::ops::Sub<Output = Self>
+    + std::ops::Mul<Output = Self>
+    + std::ops::Div<Output = Self>
+    + std::ops::AddAssign
+    + std::ops::SubAssign
+    + std::ops::MulAssign
+    + std::ops::DivAssign
+    + PartialEq
+    + PartialOrd
+    + 'static
+{
+    fn zero() -> Self;
+    fn one() -> Self;
+    fn exp(self) -> Self;
+    fn ln(self) -> Self;
+    fn neg_infinity() -> Self;
+    fn nan() -> Self;
+    fn is_sign_negative(self) -> bool;
+    fn from_usize(val: usize) -> Self;
+}
+
+impl Numeric for f32 {
+    fn zero() -> Self {
+        0.0
+    }
+    fn one() -> Self {
+        1.0
+    }
+    fn exp(self) -> Self {
+        self.exp()
+    }
+    fn ln(self) -> Self {
+        self.ln()
+    }
+    fn neg_infinity() -> Self {
+        f32::NEG_INFINITY
+    }
+    fn nan() -> Self {
+        f32::NAN
+    }
+    fn is_sign_negative(self) -> bool {
+        self.is_sign_negative()
+    }
+    fn from_usize(val: usize) -> Self {
+        val as f32
+    }
+}
+
+impl Numeric for f64 {
+    fn zero() -> Self {
+        0.0
+    }
+    fn one() -> Self {
+        1.0
+    }
+    fn exp(self) -> Self {
+        self.exp()
+    }
+    fn ln(self) -> Self {
+        self.ln()
+    }
+    fn neg_infinity() -> Self {
+        f64::NEG_INFINITY
+    }
+    fn nan() -> Self {
+        f64::NAN
+    }
+    fn is_sign_negative(self) -> bool {
+        self.is_sign_negative()
+    }
+    fn from_usize(val: usize) -> Self {
+        val as f64
+    }
+}
+
+pub struct Tensor<T: Numeric = f32> {
     shape: Vec<usize>,
     strides: Vec<usize>,
-    data: Vec<f32>,
+    data: Vec<T>,
 }
 
-impl Tensor {
-    pub fn new<S>(shape: S, data: Vec<f32>) -> Result<Self, TensorError>
+impl<T: Numeric> Tensor<T> {
+    pub fn new<S>(shape: S, data: Vec<T>) -> Result<Self, TensorError>
     where
         S: Into<Vec<usize>>,
     {
@@ -42,7 +123,7 @@ impl Tensor {
         Tensor {
             shape: shape_vec,
             strides,
-            data: vec![0.0; num_elements],
+            data: vec![T::zero(); num_elements],
         }
     }
 
@@ -56,11 +137,11 @@ impl Tensor {
         Tensor {
             shape: shape_vec,
             strides,
-            data: vec![1.0; num_elements],
+            data: vec![T::one(); num_elements],
         }
     }
 
-    pub fn get(&self, indices: &[usize]) -> Option<&f32> {
+    pub fn get(&self, indices: &[usize]) -> Option<&T> {
         if indices.len() != self.shape.len() {
             return None;
         }
@@ -83,11 +164,11 @@ impl Tensor {
         &self.strides
     }
 
-    pub fn data(&self) -> &[f32] {
+    pub fn data(&self) -> &[T] {
         &self.data
     }
 
-    pub fn data_mut(&mut self) -> &mut Vec<f32> {
+    pub fn data_mut(&mut self) -> &mut Vec<T> {
         &mut self.data
     }
 }

tensor/src/core/ops/binary.rs

@@ -1,12 +1,12 @@
 use super::broadcast::{compute_broadcast_shape_and_strides, is_broadcastable};
 use crate::core::utils::unravel_index;
-use crate::core::TensorError;
+use crate::core::{TensorError, Numeric};
 use crate::Tensor;
 
-impl Tensor {
-    fn binary_op<F>(&self, other: &Tensor, op: F) -> Result<Tensor, TensorError>
+impl<T: Numeric> Tensor<T> {
+    fn binary_op<F>(&self, other: &Tensor<T>, op: F) -> Result<Tensor<T>, TensorError>
     where
-        F: Fn(f32, f32) -> f32,
+        F: Fn(T, T) -> T,
     {
         let self_shape = self.shape();
         let other_shape = other.shape();
@@ -17,7 +17,7 @@ impl Tensor {
         }
 
         if self_shape == other_shape {
-            let result_data: Vec<f32> = self
+            let result_data: Vec<T> = self
                 .data
                 .iter()
                 .zip(other.data.iter())
@@ -32,7 +32,7 @@ impl Tensor {
         let self_data = self.data();
         let other_data = other.data();
         let result_size: usize = bc_shape.iter().product();
-        let mut result_data: Vec<f32> = Vec::with_capacity(result_size);
+        let mut result_data: Vec<T> = Vec::with_capacity(result_size);
 
         for i in 0..result_size {
             let multi_idx = unravel_index(i, &bc_shape);
@@ -50,9 +50,9 @@ impl Tensor {
         Tensor::new(bc_shape, result_data)
     }
 
-    fn binary_op_inplace<F>(&mut self, other: &Tensor, op: F)
+    fn binary_op_inplace<F>(&mut self, other: &Tensor<T>, op: F)
     where
-        F: Fn(&mut f32, f32),
+        F: Fn(&mut T, T),
     {
         let self_shape = self.shape();
         let other_shape = other.shape();
@@ -69,39 +69,39 @@ impl Tensor {
             });
     }
 
-    pub fn add(&self, other: &Tensor) -> Result<Tensor, TensorError> {
+    pub fn add(&self, other: &Tensor<T>) -> Result<Tensor<T>, TensorError> {
         self.binary_op(other, |a, b| a + b)
     }
 
-    pub fn add_inplace(&mut self, other: &Tensor) {
+    pub fn add_inplace(&mut self, other: &Tensor<T>) {
         self.binary_op_inplace(other, |a, b| *a += b);
     }
 
-    pub fn sub(&self, other: &Tensor) -> Result<Tensor, TensorError> {
+    pub fn sub(&self, other: &Tensor<T>) -> Result<Tensor<T>, TensorError> {
         self.binary_op(other, |a, b| a - b)
     }
 
-    pub fn sub_inplace(&mut self, other: &Tensor) {
+    pub fn sub_inplace(&mut self, other: &Tensor<T>) {
         self.binary_op_inplace(other, |a, b| *a -= b);
     }
 
-    pub fn mul(&self, other: &Tensor) -> Result<Tensor, TensorError> {
+    pub fn mul(&self, other: &Tensor<T>) -> Result<Tensor<T>, TensorError> {
         self.binary_op(other, |a, b| a * b)
     }
 
-    pub fn mul_inplace(&mut self, other: &Tensor) {
+    pub fn mul_inplace(&mut self, other: &Tensor<T>) {
         self.binary_op_inplace(other, |a, b| *a *= b);
     }
 
-    pub fn div(&self, other: &Tensor) -> Result<Tensor, TensorError> {
+    pub fn div(&self, other: &Tensor<T>) -> Result<Tensor<T>, TensorError> {
         self.binary_op(other, |a, b| a / b)
     }
 
-    pub fn div_inplace(&mut self, other: &Tensor) {
+    pub fn div_inplace(&mut self, other: &Tensor<T>) {
         self.binary_op_inplace(other, |a, b| *a /= b);
     }
 
-    pub fn matmul(&self, other: &Tensor) -> Result<Tensor, TensorError> {
+    pub fn matmul(&self, other: &Tensor<T>) -> Result<Tensor<T>, TensorError> {
         let lhs_shape = self.shape();
         let rhs_shape = other.shape();
         if lhs_shape.len() != 2 || rhs_shape.len() != 2 {
@@ -120,7 +120,7 @@ impl Tensor {
 
         let lhs_data = self.data();
         let rhs_data = other.data();
-        let mut result_data: Vec<f32> = vec![0.0_f32; rows_left * cols_right];
+        let mut result_data: Vec<T> = vec![T::zero(); rows_left * cols_right];
         for i in 0..rows_left {
             for k in 0..cols_left {
                 for j in 0..cols_right {
@@ -129,6 +129,6 @@ impl Tensor {
                 }
             }
         }
-        Ok(Tensor::new(vec![rows_left, cols_right], result_data).unwrap())
+        Tensor::new(vec![rows_left, cols_right], result_data)
     }
 }

tensor/src/core/ops/movement.rs

@@ -1,7 +1,7 @@
-use crate::core::{calculate_strides, TensorError};
+use crate::core::{calculate_strides, TensorError, Numeric};
 use crate::Tensor;
 
-impl Tensor {
+impl<T: Numeric> Tensor<T> {
     pub fn reshape(&mut self, shape: &[usize]) -> Result<(), TensorError> {
         let new_length: usize = shape.iter().product();
         let current_length: usize = self.shape.iter().product();
@@ -51,7 +51,7 @@ impl Tensor {
         }
 
         let (m, n) = (self.shape[0], self.shape[1]);
-        let mut new_data = vec![0.0_f32; self.data.len()];
+        let mut new_data = vec![T::zero(); self.data.len()];
         for i in 0..m {
             for j in 0..n {
                 let old_idx = i * n + j;

tensor/src/core/ops/reduce.rs

@@ -1,14 +1,14 @@
 use crate::core::utils::unravel_index;
-use crate::core::TensorError;
+use crate::core::{TensorError, Numeric};
 use crate::Tensor;
 
-impl Tensor {
-    pub fn sum(&self) -> Tensor {
-        let sum: f32 = self.data().iter().sum();
+impl<T: Numeric> Tensor<T> {
+    pub fn sum(&self) -> Tensor<T> {
+        let sum: T = self.data().iter().fold(T::zero(), |acc, &x| acc + x);
         Tensor::new(vec![1], vec![sum]).unwrap()
     }
 
-    pub fn sum_dim(&self, dim: usize) -> Result<Tensor, TensorError> {
+    pub fn sum_dim(&self, dim: usize) -> Result<Tensor<T>, TensorError> {
         let self_data = self.data();
         let self_shape = self.shape();
         let self_strides = self.strides();
@@ -22,7 +22,7 @@ impl Tensor {
         let dim_size = result_shape.remove(dim);
 
         let result_size: usize = result_shape.iter().product();
-        let mut result_data = vec![0.0_f32; result_size];
+        let mut result_data = vec![T::zero(); result_size];
 
         for i in 0..result_size {
             let result_multi_idx = unravel_index(i, &result_shape);
@@ -36,7 +36,7 @@ impl Tensor {
                     j += 1;
                 }
             }
-            let mut sum = 0.0_f32;
+            let mut sum = T::zero();
             for k in 0..dim_size {
                 full_multi_idx[dim] = k;
                 let mut offset = 0;
@@ -50,18 +50,18 @@ impl Tensor {
         Tensor::new(result_shape, result_data)
     }
 
-    pub fn mean(&self) -> Tensor {
-        let sum: Tensor = self.sum();
-        &sum / self.shape().iter().product::<usize>() as f32
+    pub fn mean(&self) -> Tensor<T> {
+        let sum: Tensor<T> = self.sum();
+        &sum / T::from_usize(self.shape().iter().product::<usize>())
     }
 
-    pub fn mean_dim(&self, dim: usize) -> Result<Tensor, TensorError> {
+    pub fn mean_dim(&self, dim: usize) -> Result<Tensor<T>, TensorError> {
         if self.shape().len() < dim {
             return Err(TensorError::IndexError(
                 "Dimension out of range for the tensor".to_string(),
             ));
         }
-        let sum: Tensor = self.sum_dim(dim).unwrap();
-        Ok(&sum / self.shape()[dim] as f32)
+        let sum: Tensor<T> = self.sum_dim(dim).unwrap();
+        Ok(&sum / T::from_usize(self.shape()[dim]))
     }
 }

tensor/src/core/ops/unary.rs

@@ -1,31 +1,31 @@
-use crate::Tensor;
+use crate::{Tensor, core::Numeric};
 
-impl Tensor {
-    fn unary_op<F>(&self, op: F) -> Tensor
+impl<T: Numeric> Tensor<T> {
+    fn unary_op<F>(&self, op: F) -> Tensor<T>
     where
-        F: Fn(f32) -> f32,
+        F: Fn(T) -> T,
     {
-        let result_data: Vec<f32> = self.data().iter().map(|x| op(*x)).collect();
-        return Tensor::new(self.shape(), result_data).unwrap();
+        let result_data: Vec<T> = self.data().iter().map(|x| op(*x)).collect();
+        Tensor::new(self.shape(), result_data).unwrap()
     }
 
-    pub fn exp(&self) -> Tensor {
+    pub fn exp(&self) -> Tensor<T> {
         self.unary_op(|x| x.exp())
     }
 
-    pub fn log(&self) -> Tensor {
+    pub fn log(&self) -> Tensor<T> {
         self.unary_op(|x| {
-            if x == 0.0 {
-                f32::NEG_INFINITY // log(0) -> -inf
-            } else if x < 0.0 {
-                f32::NAN // log of negative numbers is undefined
+            if x == T::zero() {
+                T::neg_infinity()
+            } else if x.is_sign_negative() {
+                T::nan()
             } else {
                 x.ln()
             }
         })
     }
 
-    pub fn relu(&self) -> Tensor {
-        self.unary_op(|x| if x > 0.0_f32 { x } else { 0.0_f32 })
+    pub fn relu(&self) -> Tensor<T> {
+        self.unary_op(|x| if x > T::zero() { x } else { T::zero() })
     }
 }