gradcheck 0.1.0

// gradcheck — finite-difference gradient checking for Rust ML frameworks.
// Copyright (c) 2026 Henos D <henosd19@gmail.com> (GitHub: @4ktLuffy)
// Repository: https://github.com/4ktLuffy/gradcheck
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Sweeping a check across many shapes.
//!
//! Motivation, from a real defect: burn's CPU `matmul` returned wrong values for 33 of 512
//! shape combinations, and every failing case shared a parity condition. Checking one shape
//! would have found nothing. **Checking a single shape is not evidence of correctness.**

use crate::{gradcheck, Backend, Config, Report};

/// Result of running the same check across a set of shapes.
#[derive(Debug, Clone)]
pub struct SweepReport {
    /// Every individual report, in the order the shapes were supplied.
    pub reports: Vec<Report>,
}

impl SweepReport {
    /// Shapes whose check reported a mismatch.
    pub fn failing(&self) -> Vec<&Report> {
        self.reports.iter().filter(|r| !r.passed()).collect()
    }

    /// True when every shape passed.
    pub fn all_passed(&self) -> bool {
        self.reports.iter().all(|r| r.passed())
    }

    /// Number of shapes checked.
    pub fn len(&self) -> usize {
        self.reports.len()
    }

    /// True when no shapes were checked.
    pub fn is_empty(&self) -> bool {
        self.reports.is_empty()
    }

    /// Panic with every failing shape listed, if any failed.
    pub fn assert_all_pass(&self) {
        let bad = self.failing();
        if !bad.is_empty() {
            let detail = bad
                .iter()
                .map(|r| r.to_string())
                .collect::<Vec<_>>()
                .join("\n  ");
            panic!(
                "gradcheck sweep: {} of {} shapes mismatched:\n  {detail}",
                bad.len(),
                self.reports.len()
            );
        }
    }
}

/// Run the same function across many shapes, generating deterministic input for each.
///
/// `data_for` receives the element count and returns that many values, so callers control
/// the numerical regime (positive-only for `log`/`sqrt`, away from kinks for `abs`/`max`).
pub fn shape_sweep<B, F, D>(
    name: &str,
    shapes: &[Vec<usize>],
    data_for: D,
    f: F,
    cfg: &Config,
) -> SweepReport
where
    B: Backend,
    F: Fn(B::Tensor) -> B::Tensor + Copy,
    D: Fn(usize) -> Vec<f64>,
{
    let mut reports = Vec::with_capacity(shapes.len());
    for shape in shapes {
        let n: usize = shape.iter().product();
        let data = data_for(n);
        let label = format!("{name}{shape:?}");
        reports.push(gradcheck::<B, _>(&label, &data, shape, f, cfg));
    }
    SweepReport { reports }
}