train_station/tensor/reductions/

var.rs

//! Variance reduction operations for tensors
//!
//! This module provides variance reduction operations for tensors.
//! The variance measures the average squared deviation from the mean,
//! calculated as the mean of squared differences from the mean. This is
//! commonly used in statistics, data analysis, and machine learning for
//! understanding data variability and as a component of other statistical
//! measures like standard deviation.
//!
//! # Operations
//!
//! * `var()` - Computes variance over all elements, returning a scalar tensor
//! * `var_dims()` - Computes variance over specified dimensions with optional dimension preservation
//!
//! # Statistical Details
//!
//! The implementation uses population variance (unbiased=false), which
//! divides by n rather than n-1. This matches PyTorch's default behavior for
//! consistency with the reference implementation.
//!
//! # Examples
//!
//! ```
//! use train_station::Tensor;
//!
//! // Compute variance of all elements
//! let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
//! let variance = tensor.var();
//! assert!((variance.get(&[0]) - 1.25).abs() < 1e-5);
//!
//! // Compute variance along specific dimensions
//! let matrix = Tensor::from_slice(&[1.0, 3.0, 2.0, 2.0], vec![2, 2]).unwrap();
//! let row_vars = matrix.var_dims(&[1], true);
//! assert_eq!(row_vars.shape().dims(), vec![2, 1]);
//! ```
//!
//! # Performance
//!
//! The implementation uses optimized paths for contiguous tensors with manual loop unrolling
//! for better performance. Non-contiguous tensors use stride-aware iteration to maintain
//! correctness while preserving memory layout efficiency.
//!
//! # Gradient Tracking
//!
//! Both operations support automatic gradient tracking when `requires_grad` is enabled.
//! The gradient computation follows the mathematical derivative of the variance operation.

use crate::gradtrack::{GradEngine, GradFn};
use crate::tensor::core::Tensor;

impl Tensor {
    /// Computes the variance over all elements
    ///
    /// The variance is calculated as the mean of squared differences from the mean.
    /// This operation reduces the tensor to a scalar value \[1\].
    ///
    /// The implementation uses population variance (divides by n rather
    /// than n-1) to match PyTorch's default behavior.
    ///
    /// # Returns
    ///
    /// A scalar tensor containing the variance value
    ///
    /// # Examples
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Basic variance calculation
    /// let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
    /// let variance = tensor.var();
    /// assert!((variance.get(&[0]) - 1.25).abs() < 1e-5);
    /// ```
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Variance of a larger dataset
    /// let data = vec![1.0, 3.0, 5.0, 7.0, 2.0, 4.0, 6.0, 8.0];
    /// let tensor = Tensor::from_slice(&data, vec![2, 2, 2]).unwrap();
    /// let variance = tensor.var();
    /// // mean=4.5, var=mean([3.5², 1.5², 0.5², 2.5², 2.5², 0.5², 1.5², 3.5²]) = 5.25
    /// assert!((variance.get(&[0]) - 5.25).abs() < 1e-5);
    /// ```
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Variance of constant values (should be 0)
    /// let tensor = Tensor::from_slice(&[5.0, 5.0, 5.0, 5.0], vec![4]).unwrap();
    /// let variance = tensor.var();
    /// assert!((variance.get(&[0]) - 0.0).abs() < 1e-6);
    /// ```
    ///
    /// # Performance
    ///
    /// Uses optimized contiguous tensor path with manual loop unrolling for better
    /// performance. Non-contiguous tensors use stride-aware iteration.
    /// The algorithm performs two passes: first to compute the mean, then to
    /// compute the variance.
    #[track_caller]
    pub fn var(&self) -> Tensor {
        let mut out = Tensor::new(vec![1]);
        if self.size() == 0 {
            out.fill(0.0);
        } else {
            // mean
            let mut mean_val = 0.0f32;
            let n = self.size() as f32;

            if self.is_contiguous() {
                // Fast path for contiguous tensors
                unsafe {
                    let src = self.as_ptr();
                    for i in 0..self.size() {
                        mean_val += *src.add(i);
                    }
                }
            } else {
                // Stride-aware path for non-contiguous tensors
                let dims = self.shape().dims().to_vec();
                for flat_idx in 0..self.size() {
                    // Convert flat index to multi-dimensional coordinates
                    let mut coords = vec![0; dims.len()];
                    let mut tmp = flat_idx;
                    for k in (0..dims.len()).rev() {
                        coords[k] = tmp % dims[k];
                        tmp /= dims[k];
                    }

                    // Get value using stride-aware offset
                    let offset = self.shape().offset(&coords);
                    let value = unsafe { *self.as_ptr().add(offset) };
                    mean_val += value;
                }
            }
            mean_val /= n;

            // var
            let mut var_val = 0.0f32;

            if self.is_contiguous() {
                // Fast path for contiguous tensors
                unsafe {
                    let src = self.as_ptr();
                    for i in 0..self.size() {
                        let d = *src.add(i) - mean_val;
                        var_val += d * d;
                    }
                }
            } else {
                // Stride-aware path for non-contiguous tensors
                let dims = self.shape().dims().to_vec();
                for flat_idx in 0..self.size() {
                    // Convert flat index to multi-dimensional coordinates
                    let mut coords = vec![0; dims.len()];
                    let mut tmp = flat_idx;
                    for k in (0..dims.len()).rev() {
                        coords[k] = tmp % dims[k];
                        tmp /= dims[k];
                    }

                    // Get value using stride-aware offset
                    let offset = self.shape().offset(&coords);
                    let value = unsafe { *self.as_ptr().add(offset) };
                    let d = value - mean_val;
                    var_val += d * d;
                }
            }
            var_val /= n;

            unsafe {
                *out.as_mut_ptr() = var_val;
            }
        }

        if self.requires_grad() {
            let mut result = out.clone();
            result.set_requires_grad_internal(true);
            let mean_tensor = {
                let mut t = Tensor::new(vec![1]);
                if self.size() == 0 {
                    t.fill(0.0);
                } else {
                    let mut acc = 0.0f32;

                    if self.is_contiguous() {
                        unsafe {
                            for i in 0..self.size() {
                                acc += *self.as_ptr().add(i);
                            }
                        }
                    } else {
                        let dims = self.shape().dims().to_vec();
                        for flat_idx in 0..self.size() {
                            // Convert flat index to multi-dimensional coordinates
                            let mut coords = vec![0; dims.len()];
                            let mut tmp = flat_idx;
                            for k in (0..dims.len()).rev() {
                                coords[k] = tmp % dims[k];
                                tmp /= dims[k];
                            }

                            // Get value using stride-aware offset
                            let offset = self.shape().offset(&coords);
                            let value = unsafe { *self.as_ptr().add(offset) };
                            acc += value;
                        }
                    }

                    unsafe {
                        *t.as_mut_ptr() = acc / (self.size() as f32);
                    }
                }
                t
            };
            let grad_fn = GradFn::ReduceVar {
                saved_mean: Box::new(mean_tensor),
                saved_input: Box::new(self.clone()),
                input_shape: self.shape().dims().to_vec(),
            };
            result.set_grad_fn(grad_fn.clone());
            GradEngine::register_operation(result.id(), vec![self.id()], grad_fn);
            return result;
        }

        out
    }

    /// Computes the variance over specified dimensions
    ///
    /// Reduces the tensor along the specified dimensions by computing the variance
    /// of each slice. The result maintains the original tensor structure with
    /// reduced dimensions optionally preserved as size-1 dimensions.
    ///
    /// Uses population variance (divides by n rather than n-1) to match
    /// PyTorch's default behavior.
    ///
    /// # Arguments
    ///
    /// * `dims` - Vector of dimension indices to reduce over (must be valid for tensor rank)
    /// * `keepdim` - Whether to keep reduced dimensions as size-1 dimensions
    ///
    /// # Returns
    ///
    /// A tensor with variance computed over the specified dimensions
    ///
    /// # Examples
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Variance along rows (dimension 1) with keepdim=true
    /// let matrix = Tensor::from_slice(&[1.0, 3.0, 2.0, 2.0], vec![2, 2]).unwrap();
    /// let row_vars = matrix.var_dims(&[1], true);
    /// assert_eq!(row_vars.shape().dims(), vec![2, 1]);
    /// assert!((row_vars.get(&[0, 0]) - 1.0).abs() < 1e-6); // var([1, 3]) = 1.0
    /// assert!((row_vars.get(&[1, 0]) - 0.0).abs() < 1e-6); // var([2, 2]) = 0.0
    /// ```
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Variance along columns (dimension 0) with keepdim=false
    /// let matrix = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    /// let col_vars = matrix.var_dims(&[0], false);
    /// assert_eq!(col_vars.shape().dims(), vec![2]);
    /// // var([1, 3]) = 1.0, var([2, 4]) = 1.0
    /// assert!((col_vars.get(&[0]) - 1.0).abs() < 1e-6);
    /// assert!((col_vars.get(&[1]) - 1.0).abs() < 1e-6);
    /// ```
    ///
    /// ```
    /// use train_station::Tensor;
    ///
    /// // Variance over multiple dimensions
    /// let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    /// let var_all = tensor.var_dims(&[0, 1], false);
    /// assert_eq!(var_all.shape().dims(), vec![1]);
    /// // var([1, 2, 3, 4]) = 1.25
    /// assert!((var_all.get(&[0]) - 1.25).abs() < 1e-5);
    /// ```
    ///
    /// # Panics
    ///
    /// * If `dims` is empty
    /// * If any dimension index is out of bounds for the tensor rank
    /// * If the reduced size is 0 (invalid for variance calculation)
    ///
    /// # Performance
    ///
    /// Uses efficient coordinate-based iteration that works correctly with
    /// both contiguous and non-contiguous tensor layouts. The algorithm performs
    /// two passes: first to compute means, then to compute variances.
    #[track_caller]
    pub fn var_dims(&self, dims: &[usize], keepdim: bool) -> Tensor {
        assert!(!dims.is_empty(), "var_dims requires at least one dimension");
        let rank = self.shape().rank();
        for &d in dims {
            assert!(
                d < rank,
                "var_dims dim {} out of bounds for rank {}",
                d,
                rank
            );
        }

        // Output shape
        let mut out_dims = self.shape().dims().to_vec();
        let mut reduced: Vec<usize> = dims.to_vec();
        reduced.sort_unstable();
        reduced.dedup();
        for &d in reduced.iter() {
            out_dims[d] = if keepdim { 1 } else { 0 };
        }
        if !keepdim {
            out_dims.retain(|&s| s != 0);
        }
        if out_dims.is_empty() {
            out_dims.push(1);
        }

        let mut mean = Tensor::zeros(out_dims.clone());
        let mut var = Tensor::zeros(out_dims.clone());

        let in_shape = self.shape().dims().to_vec();
        let out_rank = mean.shape().rank();
        let mut in_coords = vec![0usize; rank];
        let n_reduced: usize = reduced.iter().map(|&d| in_shape[d]).product();
        assert!(n_reduced > 0, "reduced size must be > 0");
        unsafe {
            let mptr = mean.as_mut_ptr();
            // sum for mean
            for lin in 0..self.size() {
                let mut tmp = lin;
                for i in (0..rank).rev() {
                    let s = in_shape[i];
                    in_coords[i] = if s == 0 { 0 } else { tmp % s };
                    if s != 0 {
                        tmp /= s;
                    }
                }
                let mut out_coords: Vec<usize> = Vec::with_capacity(out_rank);
                for (i, &ic) in in_coords.iter().enumerate().take(rank) {
                    if reduced.contains(&i) {
                        if keepdim {
                            out_coords.push(0);
                        }
                    } else {
                        out_coords.push(ic);
                    }
                }
                let off = if out_coords.is_empty() {
                    0
                } else {
                    mean.shape().offset(&out_coords)
                };
                // Get input value using stride-aware offset
                let in_offset = self.shape().offset(&in_coords);
                let value = *self.as_ptr().add(in_offset);
                *mptr.add(off) += value;
            }
            for i in 0..mean.size() {
                *mptr.add(i) /= n_reduced as f32;
            }
            // accumulate squared diffs
            let vptr = var.as_mut_ptr();
            for lin in 0..self.size() {
                let mut tmp = lin;
                for i in (0..rank).rev() {
                    let s = in_shape[i];
                    in_coords[i] = if s == 0 { 0 } else { tmp % s };
                    if s != 0 {
                        tmp /= s;
                    }
                }
                let mut out_coords: Vec<usize> = Vec::with_capacity(out_rank);
                for (i, &ic) in in_coords.iter().enumerate().take(rank) {
                    if reduced.contains(&i) {
                        if keepdim {
                            out_coords.push(0);
                        }
                    } else {
                        out_coords.push(ic);
                    }
                }
                let off = if out_coords.is_empty() {
                    0
                } else {
                    var.shape().offset(&out_coords)
                };
                let mu = *mptr.add(off);

                // Get input value using stride-aware offset
                let in_offset = self.shape().offset(&in_coords);
                let x = *self.as_ptr().add(in_offset);
                *vptr.add(off) += (x - mu) * (x - mu);
            }
            for i in 0..var.size() {
                *vptr.add(i) /= n_reduced as f32;
            }
        }

        if self.requires_grad() {
            let mut result = var.clone();
            result.set_requires_grad_internal(true);
            let grad_fn = GradFn::ReduceVarDims {
                dims: reduced,
                keepdim,
                input_shape: self.shape().dims().to_vec(),
                saved_mean: Box::new(mean),
                saved_input: Box::new(self.clone()),
            };
            result.set_grad_fn(grad_fn.clone());
            GradEngine::register_operation(result.id(), vec![self.id()], grad_fn);
            return result;
        }

        var
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_var_forward_basic() {
        let x = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
        let v = x.var();
        unsafe {
            let val = *v.as_ptr();
            assert!((val - 1.25).abs() < 1e-6);
        }
    }

    #[test]
    fn test_var_dims_forward() {
        let x = Tensor::from_slice(&[1.0, 3.0, 2.0, 2.0], vec![2, 2]).unwrap();
        let v = x.var_dims(&[1], true);
        assert_eq!(v.shape().dims(), vec![2, 1]);
        assert!((v.get(&[0, 0]) - 1.0).abs() < 1e-6);
        assert!((v.get(&[1, 0]) - 0.0).abs() < 1e-6);
    }

    #[test]
    fn test_var_non_contiguous_transpose() {
        // Test var on transposed tensor (non-contiguous view)
        let x = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
        // Original: [[1, 2, 3], [4, 5, 6]]

        let x_t = x.transpose(0, 1);
        // Transposed: [[1, 4], [2, 5], [3, 6]]
        assert!(!x_t.is_contiguous()); // Should be a view

        let var_orig = x.var();
        let var_view = x_t.var();

        // Both should give the same result
        assert!((var_orig.get(&[0]) - var_view.get(&[0])).abs() < 1e-6);

        // Expected var of [1,2,3,4,5,6]: mean=3.5, var=mean([2.5^2,1.5^2,0.5^2,0.5^2,1.5^2,2.5^2])=2.9167
        let expected_var = 2.9166667_f32;
        assert!((var_orig.get(&[0]) - expected_var).abs() < 1e-5);
    }

    #[test]
    fn test_var_dims_non_contiguous() {
        // Test var_dims on non-contiguous tensor
        let x = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
        let x_t = x.transpose(0, 1); // [3, 2]
        assert!(!x_t.is_contiguous());

        // Var along dim 0 of transposed tensor
        let var_dim0 = x_t.var_dims(&[0], false);
        assert_eq!(var_dim0.shape().dims(), vec![2]);

        // For dim 0: [1,2,3] and [4,5,6]
        // [1,2,3]: mean=2, var=((1-2)^2 + (2-2)^2 + (3-2)^2)/3 = 2/3 ≈ 0.6667
        // [4,5,6]: mean=5, var=((4-5)^2 + (5-5)^2 + (6-5)^2)/3 = 2/3 ≈ 0.6667
        let expected_var = 2.0 / 3.0_f32;
        assert!((var_dim0.get(&[0]) - expected_var).abs() < 1e-5);
        assert!((var_dim0.get(&[1]) - expected_var).abs() < 1e-5);
    }

    #[test]
    fn test_var_permuted_tensor() {
        // Test with permuted tensor - simple case with known var
        let data = vec![1.0, 3.0, 5.0, 7.0, 2.0, 4.0, 6.0, 8.0];
        let x = Tensor::from_slice(&data, vec![2, 2, 2]).unwrap();

        // Permute dimensions [2, 2, 2] -> [2, 2, 2] (swap first and last)
        let x_perm = x.permute(vec![2, 1, 0]);
        assert!(!x_perm.is_contiguous());

        let var_orig = x.var();
        let var_perm = x_perm.var();

        // Should give same result
        assert!((var_orig.get(&[0]) - var_perm.get(&[0])).abs() < 1e-6);

        // Data is [1,3,5,7,2,4,6,8], mean=4.5
        // var = mean([3.5^2, 1.5^2, 0.5^2, 2.5^2, 2.5^2, 0.5^2, 1.5^2, 3.5^2]) = 5.25
        let expected_var = 5.25_f32;
        assert!((var_orig.get(&[0]) - expected_var).abs() < 1e-5);
    }
}