diff --git a/src/tensor/mod.rs b/src/tensor/mod.rs index cfad14f..e1fcc7d 100644 --- a/src/tensor/mod.rs +++ b/src/tensor/mod.rs @@ -2,6 +2,7 @@ pub mod ops; mod utils; use ops::elementwise::{Add, Div, Mul, Sub}; +use ops::reductions::Sum; use std::ops::Index; use utils::{ are_broadcast_compatible, broadcast_shape, compute_strides, multi_dim_iter, recursive_slice, @@ -364,3 +365,55 @@ impl Div for Tensor { Ok(Box::new(Self::new(new_data, self.shape.clone()))) } } + +impl Sum for Tensor { + fn sum(&self, dim: Option) -> Self { + if dim.is_none() { + let total_sum = self.data.iter().sum(); + return Self::new(vec![total_sum], vec![1]); + } + + let dim = dim.unwrap(); + let mut result_shape = self.shape.clone(); + result_shape[dim] = 1; + + let mut result_data = vec![0.0; result_shape.iter().product()]; + let result_strides = compute_strides(&result_shape); + + // TODO: Might need to revisit as it will introduced overhead + // overhead because we're repeatedly converting between + // flat and multi-dimensional indices + for i in 0..self.shape[dim] { + for (j, result_val) in result_data.iter_mut().enumerate() { + let mut idx = flat_to_multi_index(&self.shape, &self.strides, j); + idx[dim] = i; + + let flat_index = compute_flat_index(&self.shape, &self.strides, &idx); + *result_val += self.data[flat_index]; + } + } + + Self { + data: result_data, + shape: result_shape, + strides: result_strides, + } + } +} + +fn compute_flat_index(shape: &[usize], strides: &[usize], idx: &[usize]) -> usize { + assert_eq!(shape.len(), idx.len()); + idx.iter().zip(strides).fold(0, |acc, (&i, &s)| acc + i * s) +} + +fn flat_to_multi_index(shape: &[usize], strides: &[usize], flat_index: usize) -> Vec { + let mut idx = vec![0; shape.len()]; + let mut remainder = flat_index; + + for i in (0..shape.len()).rev() { + idx[i] = remainder / strides[i]; + remainder %= strides[i]; + } + + idx +} diff --git a/src/tensor/ops/mod.rs b/src/tensor/ops/mod.rs index 4878130..ed3b228 100644 --- a/src/tensor/ops/mod.rs +++ b/src/tensor/ops/mod.rs @@ -1 +1,2 @@ pub mod elementwise; +pub mod reductions; diff --git a/src/tensor/ops/reductions.rs b/src/tensor/ops/reductions.rs new file mode 100644 index 0000000..f657d5e --- /dev/null +++ b/src/tensor/ops/reductions.rs @@ -0,0 +1,15 @@ +pub trait Sum { + fn sum(&self, dim: Option) -> Self; +} + +pub trait Average { + fn mean(&self, dim: Option) -> Self; +} + +pub trait Max { + fn max(&self, dim: Option) -> Self; +} + +pub trait Min { + fn min(&self, dim: Option) -> Self; +}