oxmera-tensor 0.5.2

//! Small batched linear algebra on the last two dimensions: identity and
//! diagonal constructors, trace, Cholesky, log-determinant and the
//! symmetric eigen-decomposition.
//!
//! `eye`, `diag`, `diag_embed`, `trace`, `logdet` and `det` are composites
//! of recorded ops and therefore run and differentiate on every backend.
//! `cholesky` and `eigh` are backend primitives; a backend that declines
//! them falls back to the CPU implementation through a device round-trip,
//! the same contract `index_select` uses. Sizes are meant to be small
//! (covariance blocks, determinantal kernels), where the round-trip is
//! cheap and exactness matters more than throughput.

use oxmera_core::{DType, Device, Error, Result, Shape};

use crate::autograd::GradFn;
use crate::autograd::is_recording;
use crate::backend::{Backend, backend_for};
use crate::tensor::Tensor;

fn square_matrix_dims(t: &Tensor, op: &'static str) -> Result<(usize, usize)> {
    let d = t.dims();
    if d.len() < 2 || d[d.len() - 1] != d[d.len() - 2] {
        return Err(Error::InvalidArgument {
            op,
            detail: format!("needs a [.., n, n] tensor, got shape {:?}", t.shape()),
        });
    }
    let n = d[d.len() - 1];
    let batch: usize = d[..d.len() - 2].iter().product();
    Ok((batch, n))
}

/// Run a linear-algebra primitive on the tensor's backend, falling back
/// to a CPU round-trip when the backend declines.
fn dispatch_cpu_fallback<T>(
    t: &Tensor,
    f: impl Fn(&dyn Backend, &Tensor) -> Result<T>,
    upload: impl Fn(&dyn Backend, T) -> Result<T>,
) -> Result<T> {
    let backend = backend_for(t.device())?;
    match f(backend.as_ref(), t) {
        Err(Error::NotImplemented { .. }) if t.device() != Device::Cpu => {
            let cpu = backend.download(t)?;
            let out = f(backend_for(Device::Cpu)?.as_ref(), &cpu)?;
            upload(backend.as_ref(), out)
        }
        other => other,
    }
}

impl Tensor {
    /// The `n × n` identity matrix on the CPU.
    /// # Panics
    /// Panics if `n * n` overflows `usize`. Unchecked, this wrapped to a
    /// short allocation and surfaced as an index-out-of-bounds below.
    pub fn eye(n: usize) -> Tensor {
        let cells = n
            .checked_mul(n)
            .expect("eye: n * n overflows usize — the identity is too large to build");
        let mut v = vec![0.0f32; cells];
        for i in 0..n {
            v[i * n + i] = 1.0;
        }
        Tensor::from_vec_f32(v, Shape::from([n, n])).expect("lengths match by construction")
    }

    /// The `n × n` identity matrix on `device`.
    pub fn eye_on(n: usize, device: Device) -> Result<Tensor> {
        Tensor::eye(n).to_device(device)
    }

    /// The identity with the dtype and device of `like` (VJP plumbing).
    fn eye_like(n: usize, like: &Tensor) -> Result<Tensor> {
        Tensor::eye(n)
            .to_dtype(like.dtype())?
            .to_device(like.device())
    }

    /// The diagonal of every matrix in a `[.., n, n]` tensor, as `[.., n]`.
    /// Differentiable.
    pub fn diag(&self) -> Result<Tensor> {
        let (_, n) = square_matrix_dims(self, "diag")?;
        let eye = Tensor::eye_like(n, self)?;
        self.mul(&eye)?.sum(&[self.ndim() - 1])
    }

    /// Matrices with the vectors of a `[.., n]` tensor on their diagonals,
    /// as `[.., n, n]`. Differentiable.
    pub fn diag_embed(&self) -> Result<Tensor> {
        let d = self.dims();
        let Some(&n) = d.last() else {
            return Err(Error::InvalidArgument {
                op: "diag_embed",
                detail: "needs rank >= 1".into(),
            });
        };
        let eye = Tensor::eye_like(n, self)?;
        self.unsqueeze(self.ndim())?.mul(&eye)
    }

    /// The trace of every matrix in a `[.., n, n]` tensor, as `[..]`.
    /// Differentiable.
    pub fn trace(&self) -> Result<Tensor> {
        // Report `trace`: the caller never wrote `diag`, and an error naming
        // it sends them looking for a call that does not exist.
        let diag = self.diag().map_err(|e| match e {
            Error::InvalidArgument { detail, .. } => Error::InvalidArgument {
                op: "trace",
                detail,
            },
            other => other,
        })?;
        diag.sum(&[diag.ndim() - 1])
    }

    /// Lower-triangular Cholesky factor `L` of every symmetric
    /// positive-definite matrix in a `[.., n, n]` tensor, `L Lᵀ = A`.
    ///
    /// Reads the lower triangle. A matrix that is not positive definite
    /// is a typed [`Error::InvalidArgument`] naming the batch index and
    /// pivot. The backward pass runs on the host and returns to the input's
    /// device; its intermediates are `f64`, but it takes its inputs as
    /// `f32`, so an `f64` gradient carries `f32` precision.
    ///
    /// # Gradient convention
    ///
    /// The gradient (Murray 2016) is taken with respect to **symmetric
    /// perturbations** of the input: `d logdet/dA = A⁻¹`, which is the
    /// standard result and what PyTorch returns. Because the forward reads
    /// only the lower triangle, an *elementwise* finite difference — which
    /// perturbs one entry and so breaks symmetry — does not agree with it,
    /// and `oxmera_autograd::gradcheck` cannot be used on `cholesky`,
    /// `logdet` or `det` directly. Check them through a symmetrizer
    /// (`(X + Xᵀ)/2`), as `tests/gradcheck.rs` does, or against `A⁻¹`.
    ///
    /// The practical consequence: feed these ops a symmetric matrix. Given
    /// an asymmetric one the forward silently uses the lower triangle while
    /// the gradient describes a symmetric matrix, and the two disagree.
    pub fn cholesky(&self) -> Result<Tensor> {
        square_matrix_dims(self, "cholesky")?;
        let out = dispatch_cpu_fallback(self, |be, t| be.cholesky(t), |be, l| be.upload(&l))?;
        if !(is_recording() && self.is_tracked()) {
            return Ok(out);
        }
        let l = out.clone();
        let shape = self.shape().clone();
        let device = self.device();
        Ok(out.with_grad_fn(GradFn {
            inputs: vec![self.clone()],
            vjp: Box::new(move |g: &Tensor| {
                let (batch, n) = square_matrix_dims(&l, "cholesky backward")?;
                let dtype = l.dtype();
                let lv = l
                    .to_device(Device::Cpu)?
                    .to_dtype(DType::F32)?
                    .to_vec_f32()?;
                let gv = g
                    .to_device(Device::Cpu)?
                    .to_dtype(DType::F32)?
                    .to_vec_f32()?;
                let grad = crate::cpu_linalg::cholesky_backward(&lv, &gv, batch, n);
                let grad = Tensor::from_vec_f32(grad, shape.clone())?
                    .to_dtype(dtype)?
                    .to_device(device)?;
                Ok(vec![Some(grad)])
            }),
        }))
    }

    /// `ln det A` of every SPD matrix in a `[.., n, n]` tensor, as `[..]`,
    /// through the Cholesky factor: `2 Σ ln diag(L)`. Differentiable
    /// (the gradient is `A⁻¹`, symmetrized).
    pub fn logdet(&self) -> Result<Tensor> {
        let l = self.cholesky()?;
        let d = l.diag()?;
        d.ln()?.sum(&[d.ndim() - 1])?.mul_scalar(2.0)
    }

    /// `det A` of every SPD matrix in a `[.., n, n]` tensor, as `[..]`,
    /// through the Cholesky factor. Differentiable. For an indefinite
    /// matrix use `eigh` — this is the SPD determinant.
    pub fn det(&self) -> Result<Tensor> {
        self.logdet()?.exp()
    }

    /// Eigen-decomposition of every symmetric matrix in a `[.., n, n]`
    /// tensor: eigenvalues ascending as `[.., n]` and orthonormal
    /// eigenvectors as the columns of `[.., n, n]` (`A V = V Λ`). Not
    /// differentiable.
    ///
    /// The input must be symmetric to within [`EIGH_SYMMETRY_TOL`]
    /// (relative); anything further is a typed [`Error::InvalidArgument`]
    /// naming the batch index and the worst offending pair.
    ///
    /// Before 0.4.0 the full matrix was read and silently symmetrized, so
    /// a non-symmetric input returned the eigenpairs of `(A + Aᵀ)/2` — a
    /// different matrix — with no error: `[[1, 2], [5, 1]]` answered
    /// `[-2.5, 4.5]` where the true eigenvalues are `1 ± √10`, and the
    /// returned pair did not satisfy `A v = λ v` for the `A` that was
    /// passed. The tolerance keeps the case the check exists to permit —
    /// a covariance or Gram matrix assembled as `XᵀX / n` in `f32`, which
    /// is symmetric in intent and asymmetric in the last few bits — while
    /// refusing a transpose that was actually missed.
    pub fn eigh(&self) -> Result<(Tensor, Tensor)> {
        square_matrix_dims(self, "eigh")?;
        check_symmetric(self, "eigh")?;
        dispatch_cpu_fallback(
            self,
            |be, t| be.eigh(t),
            |be, (w, v)| Ok((be.upload(&w)?, be.upload(&v)?)),
        )
    }
}

/// How far from symmetric an [`Tensor::eigh`] input may be, relative to
/// its own largest magnitude.
///
/// `1e-5` is the bound the CPU↔GPU parity suite already uses for
/// elementwise disagreement, so it is the project's existing answer to
/// "how much floating-point drift is not a bug". A matrix that has
/// accumulated more asymmetry than the backends disagree by has a real
/// problem, not a rounding one.
pub const EIGH_SYMMETRY_TOL: f32 = 1e-5;

/// Refuse a matrix that is not symmetric to within [`EIGH_SYMMETRY_TOL`].
///
/// Reported like every other precondition in this module: the batch index
/// so that a `[.., n, n]` input names *which* matrix is wrong rather than
/// the first one, and the worst pair so the caller can see how far off it
/// is rather than only that it is off.
fn check_symmetric(t: &Tensor, op: &'static str) -> Result<()> {
    let (batch, n) = square_matrix_dims(t, op)?;
    if n < 2 {
        // A 0x0 or 1x1 matrix is symmetric by construction, and the loop
        // below would read nothing. Say so rather than relying on it.
        return Ok(());
    }
    // Host-side and in f32: this reads the values once to compare them,
    // and a GPU tensor has to come down for the Jacobi sweep anyway.
    let a = t
        .to_device(Device::Cpu)?
        .to_dtype(DType::F32)?
        .to_vec_f32()?;
    let rank = t.dims().len();
    for b in 0..batch {
        let m = &a[b * n * n..(b + 1) * n * n];
        let scale = m.iter().fold(0.0f32, |acc, v| acc.max(v.abs()));
        let mut worst = (0usize, 0usize, 0.0f32);
        for i in 0..n {
            for j in (i + 1)..n {
                let d = (m[i * n + j] - m[j * n + i]).abs();
                if d > worst.2 {
                    worst = (i, j, d);
                }
            }
        }
        // Relative to the matrix's own magnitude. An absolute bound would
        // refuse a well-formed matrix scaled up and accept a badly-formed
        // one scaled down.
        let bound = EIGH_SYMMETRY_TOL * scale.max(f32::MIN_POSITIVE);
        if worst.2 > bound {
            let (i, j, d) = worst;
            return Err(Error::InvalidArgument {
                op,
                detail: format!(
                    "matrix {b} is not symmetric (|a[{i}][{j}] - a[{j}][{i}]| = {d:e}, \
                     tolerance {bound:e}); {op} needs a symmetric input — if that is \
                     what you meant, symmetrize it explicitly with \
                     a.add(&a.transpose({d0}, {d1})?)?.mul_scalar(0.5)?",
                    d0 = rank - 2,
                    d1 = rank - 1,
                ),
            });
        }
    }
    Ok(())
}