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

//! Self-test using a reference backend with forward-mode automatic differentiation.
//!
//! This deliberately uses dual numbers rather than reverse mode, and certainly not finite
//! differences: if the "framework" under test computed its gradient the same way the oracle
//! does, agreement would prove nothing. Forward-mode duals are exact to machine precision and
//! algorithmically independent of central differences, so agreement here is real evidence
//! that the comparison machinery works.
//!
//! It also means `cargo test` verifies this crate with no ML framework installed.

use gradcheck::{assert_detects_wrong_gradient, gradcheck, shape_sweep, Backend, Config};

/// A dual number: value plus derivative with respect to one seeded input.
#[derive(Clone, Copy, Debug)]
struct Dual {
    v: f64,
    d: f64,
}

impl Dual {
    fn c(v: f64) -> Self {
        Self { v, d: 0.0 }
    }
    fn mul(self, o: Self) -> Self {
        Self {
            v: self.v * o.v,
            d: self.v * o.d + self.d * o.v,
        }
    }
    fn add(self, o: Self) -> Self {
        Self {
            v: self.v + o.v,
            d: self.d + o.d,
        }
    }
    fn tanh(self) -> Self {
        let t = self.v.tanh();
        Self {
            v: t,
            d: self.d * (1.0 - t * t),
        }
    }
    fn exp(self) -> Self {
        let e = self.v.exp();
        Self {
            v: e,
            d: self.d * e,
        }
    }
    fn sin(self) -> Self {
        Self {
            v: self.v.sin(),
            d: self.d * self.v.cos(),
        }
    }
    fn recip_shifted(self) -> Self {
        // 1 / (x + 2), kept away from the pole so finite differences stay valid
        let den = self.v + 2.0;
        Self {
            v: 1.0 / den,
            d: -self.d / (den * den),
        }
    }
}

struct RefBackend;

impl Backend for RefBackend {
    type Tensor = Vec<Dual>;

    fn name() -> &'static str {
        "reference/forward-dual"
    }

    fn from_slice(data: &[f64], _shape: &[usize]) -> Self::Tensor {
        data.iter().map(|&v| Dual::c(v)).collect()
    }

    fn to_vec(t: &Self::Tensor) -> Vec<f64> {
        t.iter().map(|d| d.v).collect()
    }

    fn forward_sum(f: &dyn Fn(Self::Tensor) -> Self::Tensor, x: &Self::Tensor) -> f64 {
        f(x.clone()).iter().map(|d| d.v).sum()
    }

    fn analytic_grad(f: &dyn Fn(Self::Tensor) -> Self::Tensor, x: &Self::Tensor) -> Vec<f64> {
        // Seed one input at a time; the derivative of the summed output is that component's
        // partial derivative. N passes, exact to machine precision.
        (0..x.len())
            .map(|i| {
                let mut seeded = x.clone();
                seeded[i].d = 1.0;
                f(seeded).iter().map(|d| d.d).sum::<f64>()
            })
            .collect()
    }
}

fn cfg() -> Config {
    // The reference backend is f64-exact, so tolerances can be tight; the step is still
    // limited by the truncation error of central differences.
    Config::f64_defaults()
        .with_step(1e-4)
        .with_rel_tol(1e-5)
        .with_floor(1e-8)
}

/// THE control. If this test ever fails, nothing else in this file means anything.
#[test]
fn negative_control_detects_a_wrong_gradient() {
    assert_detects_wrong_gradient::<RefBackend>();
}

#[test]
fn elementwise_ops_agree_with_the_oracle() {
    let data = vec![0.7, -1.3, 2.1, -0.4, 1.9, -3.2];
    let shape = [2usize, 3];
    let c = cfg();

    gradcheck::<RefBackend, _>("identity", &data, &shape, |t| t, &c).assert_pass();
    gradcheck::<RefBackend, _>(
        "tanh",
        &data,
        &shape,
        |t| t.into_iter().map(Dual::tanh).collect(),
        &c,
    )
    .assert_pass();
    gradcheck::<RefBackend, _>(
        "exp",
        &data,
        &shape,
        |t| t.into_iter().map(Dual::exp).collect(),
        &c,
    )
    .assert_pass();
    gradcheck::<RefBackend, _>(
        "sin",
        &data,
        &shape,
        |t| t.into_iter().map(Dual::sin).collect(),
        &c,
    )
    .assert_pass();
    gradcheck::<RefBackend, _>(
        "square",
        &data,
        &shape,
        |t| t.into_iter().map(|d| d.mul(d)).collect(),
        &c,
    )
    .assert_pass();
    gradcheck::<RefBackend, _>(
        "recip_shifted",
        &data,
        &shape,
        |t| t.into_iter().map(Dual::recip_shifted).collect(),
        &c,
    )
    .assert_pass();
    gradcheck::<RefBackend, _>(
        "chain_tanh_exp",
        &data,
        &shape,
        |t| t.into_iter().map(|d| d.exp().tanh()).collect(),
        &c,
    )
    .assert_pass();
}

#[test]
fn cross_element_op_agrees() {
    // Every output depends on every input, so a backward pass that drops or misroutes a
    // contribution shows up here even when elementwise ops look fine.
    let data = vec![0.7, -1.3, 2.1, -0.4];
    let c = cfg();
    let r = gradcheck::<RefBackend, _>(
        "sum_times_self",
        &data,
        &[2, 2],
        |t| {
            let total = t.iter().fold(Dual::c(0.0), |a, &b| a.add(b));
            t.into_iter().map(|d| d.mul(total)).collect()
        },
        &c,
    );
    r.assert_pass();
}

#[test]
fn sweeps_across_shapes() {
    let c = cfg();
    let shapes = vec![
        vec![1, 1],
        vec![2, 3],
        vec![3, 2],
        vec![4, 4],
        vec![1, 7],
        vec![5, 1],
    ];
    let sweep = shape_sweep::<RefBackend, _, _>(
        "tanh",
        &shapes,
        |n| {
            (0..n)
                .map(|i| ((i as f64) * 0.37).sin() * 1.7 + 0.15)
                .collect()
        },
        |t| t.into_iter().map(Dual::tanh).collect(),
        &c,
    );
    assert_eq!(sweep.len(), shapes.len());
    sweep.assert_all_pass();
}

#[test]
fn a_deliberately_wrong_backward_is_caught() {
    // Simulates the real failure mode: forward is right, backward is subtly wrong.
    // Here the "gradient" is scaled, exactly as a mis-accumulating kernel might behave.
    let data = vec![0.7, -1.3, 2.1, -0.4];
    let c = cfg();
    let bad = gradcheck::gradcheck_corrupted::<RefBackend, _>(
        "wrong_backward",
        &data,
        &[2, 2],
        |t| t.into_iter().map(Dual::tanh).collect(),
        &c,
        1.10, // only 10% off — the kind of error a loose tolerance would wave through
    );
    assert!(
        !bad.passed(),
        "a 10% wrong gradient must be caught, got: {bad}"
    );
}

#[test]
fn report_lists_every_failing_component() {
    let data = vec![0.7, -1.3, 2.1, -0.4];
    let c = cfg();
    let bad = gradcheck::gradcheck_corrupted::<RefBackend, _>(
        "all_wrong",
        &data,
        &[2, 2],
        |t| t,
        &c,
        2.0,
    );
    assert!(!bad.passed());
    assert_eq!(
        bad.failing_indices().len(),
        4,
        "all four components are wrong"
    );
}