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

//! Finite-difference gradient checking for Rust machine-learning frameworks.
//!
//! PyTorch has `torch.autograd.gradcheck`. JAX has `check_grads`. Rust has had nothing like
//! it, which is how a silent wrong-answer bug in `matmul` lived through a framework's whole
//! test suite ([burn#5304]).
//!
//! # Why compare against numbers, not another backend
//!
//! Test two backends against each other and you only catch the bugs they don't share. A
//! central difference doesn't have that problem, because the reference comes from the
//! definition of a derivative rather than from somebody's implementation:
//!
//! ```text
//! df/dx_i  ≈  ( f(x + h·e_i) - f(x - h·e_i) ) / 2h
//! ```
//!
//! There's no code in there that can be broken the same way your autodiff is.
//!
//! # Run the negative control first
//!
//! A harness that silently computes nothing will report everything as passing, and a clean
//! run out of one of those is worse than no run. So the control is an actual function rather
//! than something you're trusted to remember: [`assert_detects_wrong_gradient`] corrupts a
//! gradient on purpose and **panics if the checker still reports a pass**.
//!
//! Put it in the same test binary as your real checks. If it ever fails, ignore every other
//! result in that run.
//!
//! # Example
//!
//! ```ignore
//! use gradcheck::{gradcheck, Config, assert_detects_wrong_gradient};
//!
//! // 1. Prove the instrument works.
//! assert_detects_wrong_gradient::<MyBackend>();
//!
//! // 2. Then trust its verdicts.
//! let report = gradcheck::<MyBackend, _>(
//!     "tanh",
//!     &[0.7, -1.3, 2.1, -0.4],
//!     &[2, 2],
//!     |t| t.tanh(),
//!     &Config::default(),
//! );
//! report.assert_pass();
//! ```
//!
//! [burn#5304]: https://github.com/tracel-ai/burn/issues/5304

#![warn(missing_docs)]
#![warn(clippy::all)]

pub mod adapters;

mod backend;
mod config;
mod numeric;
mod report;
mod sweep;

pub use backend::Backend;
pub use config::Config;
pub use numeric::{numerical_gradient, realized_step, realized_step_with, Precision};
pub use report::{ComponentRecord, InvalidReason, Report, Stage, UncheckedReason, Verdict};
pub use sweep::{shape_sweep, SweepReport};

/// Check one function's gradient at one input.
///
/// Computes the framework's analytic gradient of `sum(f(x))` and compares it against a
/// central-difference approximation of the same quantity.
///
/// `name` is a free-form label that appears in the [`Report`].
pub fn gradcheck<B, F>(name: &str, data: &[f64], shape: &[usize], f: F, cfg: &Config) -> Report
where
    B: Backend,
    F: Fn(B::Tensor) -> B::Tensor,
{
    let n: usize = shape.iter().product();
    assert_eq!(
        data.len(),
        n,
        "gradcheck: data length {} does not match shape {:?} (= {} elements)",
        data.len(),
        shape,
        n
    );

    let x = B::from_slice(data, shape);
    let analytic = B::analytic_grad(&f, &x);
    assert_eq!(
        analytic.len(),
        n,
        "gradcheck: backend returned a gradient of length {} for an input of {} elements",
        analytic.len(),
        n
    );

    let numeric = numerical_gradient::<B, _>(&f, data, shape, cfg.step, cfg.precision);

    build_report::<B>(name, shape, analytic, numeric, cfg)
}

/// Same as [`gradcheck`], but scales the analytic gradient by `corrupt` before comparing.
///
/// This exists so the negative control uses the *identical* comparison path as a real check —
/// a control that runs through different code proves nothing about the real one.
pub fn gradcheck_corrupted<B, F>(
    name: &str,
    data: &[f64],
    shape: &[usize],
    f: F,
    cfg: &Config,
    corrupt: f64,
) -> Report
where
    B: Backend,
    F: Fn(B::Tensor) -> B::Tensor,
{
    let n: usize = shape.iter().product();
    let x = B::from_slice(data, shape);
    let analytic: Vec<f64> = B::analytic_grad(&f, &x)
        .into_iter()
        .map(|v| v * corrupt)
        .collect();
    let numeric = numerical_gradient::<B, _>(&f, data, shape, cfg.step, cfg.precision);
    let _ = n;
    build_report::<B>(name, shape, analytic, numeric, cfg)
}

/// Compare two gradient vectors directly, without running a framework.
///
/// The same three-way rule and the same gates as [`gradcheck`], but with both sides
/// supplied by the caller. Two uses:
///
/// 1. **Cross-backend differential.** Dump analytic gradients from two backends and
///    adjudicate them here, so a differential comparison gets the same NaN handling and
///    the same abstention semantics as a finite-difference check. Comparing two backends
///    is a *consistency* oracle — it stays silent when both share a bug — so a `Pass`
///    here means "these agree", never "this is correct".
/// 2. **Control cells.** Exact `(analytic, numeric)` pairs with exact expected verdicts.
///
/// `backend` is a free-form label recorded in the report.
pub fn compare(
    name: &str,
    backend: &'static str,
    shape: &[usize],
    analytic: Vec<f64>,
    numeric: Vec<f64>,
    cfg: &Config,
) -> Report {
    // The label travels through LABEL below, not through a field — `Backend::name` is an
    // associated function with no `self` to read one from.
    struct Labelled;
    // A tiny shim so `build_report` can stay generic over `Backend` for its name.
    thread_local! {
        static LABEL: core::cell::Cell<&'static str> = const { core::cell::Cell::new("compare") };
    }
    impl Backend for Labelled {
        type Tensor = ();
        fn name() -> &'static str {
            LABEL.with(|l| l.get())
        }
        fn from_slice(_: &[f64], _: &[usize]) -> Self::Tensor {}
        fn to_vec(_: &Self::Tensor) -> Vec<f64> {
            Vec::new()
        }
        fn forward_sum(_: &dyn Fn(Self::Tensor) -> Self::Tensor, _: &Self::Tensor) -> f64 {
            0.0
        }
        fn analytic_grad(_: &dyn Fn(Self::Tensor) -> Self::Tensor, _: &Self::Tensor) -> Vec<f64> {
            Vec::new()
        }
    }
    LABEL.with(|l| l.set(backend));
    build_report::<Labelled>(name, shape, analytic, numeric, cfg)
}

/// Adjudicate one component under the three-way rule.
///
/// ```text
/// M = max(|a|, |n|)
/// if |a - n| >  abs_tol + rel_tol*M      -> Mismatch   (at EVERY magnitude)
/// if M > floor and |a - n| <= rel_tol*M  -> Pass
/// otherwise                              -> Unchecked(AmbiguousAtFloor)
/// ```
///
/// The additive tolerance appears only on the rejection side. It can excuse a
/// discrepancy; it can never certify agreement.
pub(crate) fn adjudicate(a: f64, n: f64, cfg: &Config) -> ComponentRecord {
    let signed_zero_discrepancy =
        a == 0.0 && n == 0.0 && a.is_sign_negative() != n.is_sign_negative();

    // Non-finite on either side is never adjudicated numerically.
    if !a.is_finite() || !n.is_finite() {
        let stage = if a.is_finite() {
            Stage::Probe
        } else {
            Stage::Analytic
        };
        return ComponentRecord {
            a,
            n,
            diff: f64::NAN,
            tol_used: f64::NAN,
            outcome: Verdict::Invalid {
                stage,
                reason: InvalidReason::NonFinite,
            },
            signed_zero_discrepancy,
        };
    }

    let m = a.abs().max(n.abs());
    let diff = (a - n).abs();
    let reject_bound = cfg.abs_tol + cfg.rel_tol * m;
    let certify_bound = cfg.rel_tol * m;

    let outcome = if diff > reject_bound {
        Verdict::Mismatch
    } else if m > cfg.floor && diff <= certify_bound {
        Verdict::Pass
    } else {
        Verdict::Unchecked(UncheckedReason::AmbiguousAtFloor)
    };

    ComponentRecord {
        a,
        n,
        diff,
        tol_used: reject_bound,
        outcome,
        signed_zero_discrepancy,
    }
}

fn build_report<B: Backend>(
    name: &str,
    shape: &[usize],
    analytic: Vec<f64>,
    numeric: Vec<f64>,
    cfg: &Config,
) -> Report {
    let mk = |verdict: Verdict, components: Vec<ComponentRecord>| Report {
        name: name.to_string(),
        backend: B::name(),
        shape: shape.to_vec(),
        analytic: analytic.clone(),
        numeric: numeric.clone(),
        components,
        worst_index: 0,
        worst_rel_error: 0.0,
        verdict,
        nondeterministic: false,
    };

    // Gate 1/2 — structural. An empty gradient is structural, never a pass: an empty
    // fold is exactly the "worst stays zero" mechanism this design exists to remove.
    if analytic.len() != numeric.len() || analytic.is_empty() {
        return mk(Verdict::StructuralMismatch, Vec::new());
    }

    // Gate 3 — the comparator's own arithmetic must be usable.
    if !cfg.is_finite() {
        return mk(
            Verdict::Invalid {
                stage: Stage::Comparator,
                reason: InvalidReason::NonFinite,
            },
            Vec::new(),
        );
    }

    // Components are always recorded, whatever the case verdict turns out to be.
    let components: Vec<ComponentRecord> = analytic
        .iter()
        .zip(numeric.iter())
        .map(|(&a, &n)| adjudicate(a, n, cfg))
        .collect();

    // Case verdict by precedence over component outcomes.
    let verdict = components
        .iter()
        .fold(Verdict::Pass, |acc, c| acc.combine(c.outcome.clone()));

    // If everything was adjudicable but some abstained, report the partial count.
    let verdict = match verdict {
        Verdict::Unchecked(_) => {
            let (checked, total) = (
                components
                    .iter()
                    .filter(|c| matches!(c.outcome, Verdict::Pass | Verdict::Mismatch))
                    .count() as u32,
                components.len() as u32,
            );
            if checked == 0 {
                Verdict::Unchecked(UncheckedReason::AmbiguousAtFloor)
            } else {
                Verdict::Unchecked(UncheckedReason::Partial { checked, total })
            }
        }
        v => v,
    };

    // Worst = largest diff among components that were actually adjudicated.
    let (worst_index, worst_rel_error) = components
        .iter()
        .enumerate()
        .filter(|(_, c)| matches!(c.outcome, Verdict::Pass | Verdict::Mismatch))
        .fold((0usize, 0.0f64), |(bi, bv), (i, c)| {
            let m = c.a.abs().max(c.n.abs()).max(f64::MIN_POSITIVE);
            let rel = c.diff / m;
            if rel > bv {
                (i, rel)
            } else {
                (bi, bv)
            }
        });

    Report {
        name: name.to_string(),
        backend: B::name(),
        shape: shape.to_vec(),
        analytic,
        numeric,
        components,
        worst_index,
        worst_rel_error,
        verdict,
        nondeterministic: false,
    }
}

/// **Run this before trusting any result from this crate.**
///
/// Feeds the checker a deliberately wrong gradient (scaled by 1.5) and panics unless the
/// checker reports a mismatch. Then feeds it the same gradient uncorrupted and panics unless
/// that passes — catching the opposite failure, a checker that flags everything.
///
/// If this function returns, the instrument demonstrably distinguishes a correct gradient
/// from an incorrect one on this backend.
pub fn assert_detects_wrong_gradient<B>()
where
    B: Backend,
{
    let data = [0.7_f64, -1.3, 2.1, -0.4, 1.9, -3.2];
    let shape = [2usize, 3];
    let cfg = Config::default();

    let corrupted =
        gradcheck_corrupted::<B, _>("negative-control(x1.5)", &data, &shape, |t| t, &cfg, 1.5);
    assert!(
        !corrupted.passed(),
        "NEGATIVE CONTROL FAILED: gradcheck reported a pass for a gradient scaled by 1.5. \
         The harness cannot detect a wrong gradient on backend '{}', so every other result \
         it produces is meaningless. Report: {corrupted}",
        B::name()
    );

    let clean = gradcheck::<B, _>("negative-control(x1.0)", &data, &shape, |t| t, &cfg);
    assert!(
        clean.passed(),
        "NEGATIVE CONTROL FAILED IN THE OTHER DIRECTION: gradcheck reported a mismatch for a \
         correct identity gradient on backend '{}'. Tolerances or the adapter are wrong. \
         Report: {clean}",
        B::name()
    );
}