diff --git a/crates/kornia-core-ops/src/ops.rs b/crates/kornia-core-ops/src/ops.rs index 3fd61bf0..05a2d392 100644 --- a/crates/kornia-core-ops/src/ops.rs +++ b/crates/kornia-core-ops/src/ops.rs @@ -3,156 +3,201 @@ use num_traits::Float; use crate::error::TensorOpsError; -/// Compute the sum of the elements in the tensor along dimension `dim` -/// -/// # Arguments -/// -/// * `tensor` - The tensor to sum the elements of. -/// * `dim` - The index of the dimension/axis to perform the sum operation over. -/// -/// # Returns -/// -/// A new `Tensor` containing the element sums. -/// -/// # Errors -/// -/// If the requested dimension is greater than the number of axes of the tensor, an error is returned. -/// -/// # Example -/// -/// ``` -/// use kornia_core::{Tensor, CpuAllocator}; -/// use kornia_core_ops::ops::sum_elements; -/// -/// let data: [u8; 6] = [1, 1, 1, 1, 1, 1]; -/// let t = Tensor::::from_shape_slice([2, 3], &data, CpuAllocator).unwrap(); -/// let agg = sum_elements(&t, 1).unwrap(); -/// assert_eq!(agg.shape, [2, 1]); -/// assert_eq!(agg.as_slice(), [3, 3]); -/// ``` -pub fn sum_elements( - tensor: &Tensor, - dim: usize, -) -> Result, TensorOpsError> +/// A trait to define operationt that can be performed on a tensor. +pub trait TensorOps where - T: SafeTensorType + std::ops::Add, - A: TensorAllocator + Clone + 'static, + T: SafeTensorType, + A: TensorAllocator, { - if dim >= N { - return Err(TensorOpsError::DimOutOfBounds(dim, N - 1)); - } - - let mut out_shape = tensor.shape; - out_shape[dim] = 1; + /// Apply the power function to the tensor data. + /// + /// # Arguments + /// + /// * `n` - The power to raise the tensor data to. + /// + /// # Returns + /// + /// A new tensor with the data raised to the power. + fn powf(&self, n: T) -> Tensor + where + T: Float; - let mut out_strides = tensor.strides; - if dim > 0 { - out_strides - .iter_mut() - .take(dim) - .for_each(|s| *s /= tensor.shape[dim]); - } + /// Apply the square root function to the tensor elements. + /// + /// # Returns + /// + /// A new tensor with the square root values. + fn sqrt(&self) -> Tensor + where + T: Float; - let numel: usize = out_shape.iter().product(); - let mut data = vec![T::default(); numel]; - - for (i, v) in tensor.as_slice().iter().enumerate() { - let mut out_index = tensor.get_index_unchecked(i); - out_index[dim] = 0; - let out_offset = out_index - .iter() - .zip(out_strides.iter()) - .fold(0, |acc, (&idx, &stride)| acc + idx * stride); - let agg = unsafe { data.get_unchecked_mut(out_offset) }; - *agg = *agg + *v; - } + /// Compute the sum of the elements in the tensor along dimension `dim` + /// + /// # Arguments + /// + /// * `dim` - The index of the dimension/axis to perform the sum operation over. + /// + /// # Returns + /// + /// A new `Tensor` containing the element sums. + /// + /// # Errors + /// + /// If the requested dimension is greater than the number of axes of the tensor, an error is returned. + /// + /// # Example + /// + /// ``` + /// use kornia_core::{Tensor, CpuAllocator}; + /// use kornia_core_ops::ops::TensorOps; + /// + /// let data: [u8; 6] = [1, 1, 1, 1, 1, 1]; + /// let t = Tensor::::from_shape_slice([2, 3], &data, CpuAllocator).unwrap(); + /// let agg = t.sum_elements(1).unwrap(); + /// assert_eq!(agg.shape, [2, 1]); + /// assert_eq!(agg.as_slice(), [3, 3]); + /// ``` + fn sum_elements(&self, dim: usize) -> Result, TensorOpsError> + where + T: std::ops::Add; - let storage = TensorStorage::from_vec(data, tensor.storage.alloc().clone()); + /// Compute the p-norm of the tensor along some dimension. + /// + /// # Arguments + /// + /// * `p` - The order of the norm. + /// * `dim` - The index of the dimension/axis to use as vectors for the p-norm calculation. + /// + /// # Returns + /// + /// A new `Tensor` containing the p-norm values at all the vector locations. + /// + /// # Errors + /// + /// If the requested dimension is greater than the number of axes of the tensor, an error is returned. + /// + /// # Example + /// + /// ``` + /// use kornia_core::{Tensor, CpuAllocator}; + /// use kornia_core_ops::ops::TensorOps; + /// + /// let data: [f32; 5] = [3., 3., 2., 1., 1.]; + /// let t = Tensor::::from_shape_slice([5], &data, CpuAllocator).unwrap(); + /// let norm = t.pnorm(3., 0).unwrap(); + /// assert_eq!(norm.shape, [1]); + /// assert_eq!(norm.as_slice(), [4.]); + /// ``` + fn pnorm(&self, p: T, dim: usize) -> Result, TensorOpsError> + where + T: std::ops::Add + Float; - Ok(Tensor { - storage, - shape: out_shape, - strides: out_strides, - }) + /// Compute the euclidean norm of a tensor along some dimension. + /// + /// # Arguments + /// + /// * `dim` - The index of the dimension/axis to use as vectors for the norm calculation. + /// + /// # Returns + /// + /// A new `Tensor` containing the norm values at all the vector locations. + /// + /// # Errors + /// + /// If the requested dimension is greater than the number of axes of the tensor, an error is returned. + /// + /// # Example + /// + /// ``` + /// use kornia_core::{Tensor, CpuAllocator}; + /// use kornia_core_ops::ops::TensorOps; + /// + /// let data: [f32; 5] = [3., 2., 1., 1., 1.]; + /// let t = Tensor::::from_shape_slice([5], &data, CpuAllocator).unwrap(); + /// let norm = t.euclidean_norm(0).unwrap(); + /// assert_eq!(norm.shape, [1]); + /// assert_eq!(norm.as_slice(), [4.]); + /// ``` + fn euclidean_norm(&self, dim: usize) -> Result, TensorOpsError> + where + T: std::ops::Add + Float; } -/// Compute the p-norm of the tensor along some dimension. -/// -/// # Arguments -/// -/// * `tensor` - The tensor to calculate the p-norm from. -/// * `p` - The order of the norm. -/// * `dim` - The index of the dimension/axis to use as vectors for the p-norm calculation. -/// -/// # Returns -/// -/// A new `Tensor` containing the p-norm values at all the vector locations. -/// -/// # Errors -/// -/// If the requested dimension is greater than the number of axes of the tensor, an error is returned. -/// -/// # Example -/// -/// ``` -/// use kornia_core::{Tensor, CpuAllocator}; -/// use kornia_core_ops::ops::pnorm; -/// -/// let data: [f32; 5] = [3., 3., 2., 1., 1.]; -/// let t = Tensor::::from_shape_slice([5], &data, CpuAllocator).unwrap(); -/// let norm = pnorm(&t, 3., 0).unwrap(); -/// assert_eq!(norm.shape, [1]); -/// assert_eq!(norm.as_slice(), [4.]); -/// ``` -pub fn pnorm( - tensor: &Tensor, - p: T, - dim: usize, -) -> Result, TensorOpsError> +impl TensorOps for Tensor where - T: SafeTensorType + std::ops::Add + Float, - A: TensorAllocator + Clone + 'static, + T: SafeTensorType, + A: TensorAllocator + 'static, { - let p_inv = T::one() / p; - Ok(sum_elements(&tensor.powf(p), dim)?.powf(p_inv)) -} + fn powf(&self, n: T) -> Tensor + where + T: Float, + { + self.map(|x| x.powf(n)) + } -/// Compute the euclidean norm of a tensor along some dimension. -/// -/// # Arguments -/// -/// * `tensor` - The tensor to calculate the euclidean norm from. -/// * `dim` - The index of the dimension/axis to use as vectors for the norm calculation. -/// -/// # Returns -/// -/// A new `Tensor` containing the norm values at all the vector locations. -/// -/// # Errors -/// -/// If the requested dimension is greater than the number of axes of the tensor, an error is returned. -/// -/// # Example -/// -/// ``` -/// use kornia_core::{Tensor, CpuAllocator}; -/// use kornia_core_ops::ops::euclidean_norm; -/// -/// let data: [f32; 5] = [3., 2., 1., 1., 1.]; -/// let t = Tensor::::from_shape_slice([5], &data, CpuAllocator).unwrap(); -/// let norm = euclidean_norm(&t, 0).unwrap(); -/// assert_eq!(norm.shape, [1]); -/// assert_eq!(norm.as_slice(), [4.]); -/// ``` -pub fn euclidean_norm( - tensor: &Tensor, - dim: usize, -) -> Result, TensorOpsError> -where - T: SafeTensorType + std::ops::Add + Float, - A: TensorAllocator + Clone + 'static, -{ - Ok(sum_elements(&tensor.powi(2), dim)?.sqrt()) + fn sqrt(&self) -> Tensor + where + T: Float, + { + self.map(|x| x.sqrt()) + } + + fn sum_elements(&self, dim: usize) -> Result, TensorOpsError> + where + T: std::ops::Add, + { + if dim >= N { + return Err(TensorOpsError::DimOutOfBounds(dim, N - 1)); + } + + let mut out_shape = self.shape; + out_shape[dim] = 1; + + let mut out_strides = self.strides; + if dim > 0 { + out_strides + .iter_mut() + .take(dim) + .for_each(|s| *s /= self.shape[dim]); + } + + let numel: usize = out_shape.iter().product(); + let mut data = vec![T::default(); numel]; + + for (i, v) in self.as_slice().iter().enumerate() { + let mut out_index = self.get_index_unchecked(i); + out_index[dim] = 0; + let out_offset = out_index + .iter() + .zip(out_strides.iter()) + .fold(0, |acc, (&idx, &stride)| acc + idx * stride); + let agg = unsafe { data.get_unchecked_mut(out_offset) }; + *agg = *agg + *v; + } + + let storage = TensorStorage::from_vec(data, self.storage.alloc().clone()); + + Ok(Tensor { + storage, + shape: out_shape, + strides: out_strides, + }) + } + + fn pnorm(&self, p: T, dim: usize) -> Result, TensorOpsError> + where + T: std::ops::Add + Float, + { + let p_inv = T::one() / p; + Ok(self.powf(p).sum_elements(dim)?.powf(p_inv)) + } + + fn euclidean_norm(&self, dim: usize) -> Result, TensorOpsError> + where + T: std::ops::Add + Float, + { + Ok(self.powi(2).sum_elements(dim)?.sqrt()) + } } #[cfg(test)] @@ -165,7 +210,7 @@ mod tests { fn test_sum_dim_oob() -> Result<(), TensorError> { let data: [u8; 4] = [2, 2, 2, 2]; let t = Tensor::::from_shape_slice([2, 2], &data, CpuAllocator)?; - let res = sum_elements(&t, 2); + let res = t.sum_elements(2); assert!(res.is_err_and(|e| e == TensorOpsError::DimOutOfBounds(2, 1))); Ok(()) } @@ -174,7 +219,7 @@ mod tests { fn test_sum_1d() -> Result<(), TensorError> { let data: [u8; 4] = [1, 1, 1, 1]; let t = Tensor::::from_shape_slice([4], &data, CpuAllocator)?; - let res = sum_elements(&t, 0); + let res = t.sum_elements(0); assert!(res.is_ok_and(|v| v.as_slice() == [4])); Ok(()) @@ -187,19 +232,19 @@ mod tests { let t_f32 = t.cast::(); let t_i32 = t.cast::(); - let agg = sum_elements(&t, 1)?; + let agg = t.sum_elements(1)?; assert_eq!(agg.shape, [2, 1]); assert_eq!(agg.as_slice(), [6, 15]); - assert_eq!(sum_elements(&t_f32, 1)?.as_slice(), [6., 15.]); - assert_eq!(sum_elements(&t_i32, 1)?.as_slice(), [6, 15]); + assert_eq!(t_f32.sum_elements(1)?.as_slice(), [6., 15.]); + assert_eq!(t_i32.sum_elements(1)?.as_slice(), [6, 15]); - let agg = sum_elements(&t, 0)?; + let agg = t.sum_elements(0)?; assert_eq!(agg.shape, [1, 3]); assert_eq!(agg.as_slice(), [5, 7, 9]); - assert_eq!(sum_elements(&t_f32, 0)?.as_slice(), [5., 7., 9.]); - assert_eq!(sum_elements(&t_i32, 0)?.as_slice(), [5, 7, 9]); + assert_eq!(t_f32.sum_elements(0)?.as_slice(), [5., 7., 9.]); + assert_eq!(t_i32.sum_elements(0)?.as_slice(), [5, 7, 9]); Ok(()) } @@ -210,23 +255,23 @@ mod tests { let t_f32 = t.cast::(); let t_i32 = t.cast::(); - let agg = sum_elements(&t, 0)?; + let agg = t.sum_elements(0)?; assert_eq!(agg.shape, [1, 3, 4]); assert_eq!(agg.as_slice(), [2; 12]); - assert_eq!(sum_elements(&t_f32, 0)?.as_slice(), [2.; 12]); - assert_eq!(sum_elements(&t_i32, 0)?.as_slice(), [2; 12]); + assert_eq!(t_f32.sum_elements(0)?.as_slice(), [2.; 12]); + assert_eq!(t_i32.sum_elements(0)?.as_slice(), [2; 12]); - let agg = sum_elements(&t, 1)?; + let agg = t.sum_elements(1)?; assert_eq!(agg.shape, [2, 1, 4]); assert_eq!(agg.as_slice(), [3; 8]); - assert_eq!(sum_elements(&t_f32, 1)?.as_slice(), [3.; 8]); - assert_eq!(sum_elements(&t_i32, 1)?.as_slice(), [3; 8]); + assert_eq!(t_f32.sum_elements(1)?.as_slice(), [3.; 8]); + assert_eq!(t_i32.sum_elements(1)?.as_slice(), [3; 8]); - let agg = sum_elements(&t, 2)?; + let agg = t.sum_elements(2)?; assert_eq!(agg.shape, [2, 3, 1]); assert_eq!(agg.as_slice(), [4; 6]); - assert_eq!(sum_elements(&t_f32, 2)?.as_slice(), [4.; 6]); - assert_eq!(sum_elements(&t_i32, 2)?.as_slice(), [4; 6]); + assert_eq!(t_f32.sum_elements(2)?.as_slice(), [4.; 6]); + assert_eq!(t_i32.sum_elements(2)?.as_slice(), [4; 6]); Ok(()) } @@ -234,7 +279,7 @@ mod tests { fn test_pnorm_1d() -> Result<(), TensorOpsError> { let data: [f32; 5] = [3., 3., 2., 1., 1.]; let t = Tensor::::from_shape_slice([5], &data, CpuAllocator).unwrap(); - let norm = pnorm(&t, 3., 0)?; + let norm = t.pnorm(3., 0)?; assert_eq!(norm.shape, [1]); assert_eq!(norm.as_slice(), [4.]); Ok(()) @@ -244,7 +289,7 @@ mod tests { fn test_euclidean_norm_1d() -> Result<(), TensorOpsError> { let data: [f32; 5] = [3., 2., 1., 1., 1.]; let t = Tensor::::from_shape_slice([5], &data, CpuAllocator).unwrap(); - let norm = euclidean_norm(&t, 0)?; + let norm = t.euclidean_norm(0)?; assert_eq!(norm.shape, [1]); assert_eq!(norm.as_slice(), [4.]); Ok(()) diff --git a/crates/kornia-core/src/tensor.rs b/crates/kornia-core/src/tensor.rs index e46ebaf7..3720c01b 100644 --- a/crates/kornia-core/src/tensor.rs +++ b/crates/kornia-core/src/tensor.rs @@ -653,36 +653,6 @@ where self.map(|x| x.powi(n)) } - /// Apply the power function to the tensor data. - /// - /// # Arguments - /// - /// * `n` - The power to raise the tensor data to. - /// - /// # Returns - /// - /// A new tensor with the data raised to the power. - pub fn powf(&self, n: T) -> Tensor - where - T: Float, - A: Clone, - { - self.map(|x| x.powf(n)) - } - - /// Apply the square root function to the tensor elements. - /// - /// # Returns - /// - /// A new tensor with the square root values. - pub fn sqrt(&self) -> Tensor - where - T: Float, - A: Clone, - { - self.map(|x| x.sqrt()) - } - /// Compute absolute value of the pixel data. /// /// # Returns