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

//! Step 2a — does the weighted objective actually restore observability?
//!
//! `docs/objective-observability.md` argues that `J(x) = Σ softmax(x)` is constant, so the
//! gradient is identically zero and no comparison against it can detect a fault. The fix
//! proposed there is `J_w(x) = ⟨w, f(x)⟩`.
//!
//! This file tests both halves of that claim rather than assuming them:
//!
//!   1. reproduce the defect — the summed objective must produce an all-zero gradient,
//!   2. verify the fix — the weighted objective must produce non-zero gradients that the
//!      comparator actually adjudicates,
//!   3. and confirm the fix is *load-bearing* — a deliberately corrupted gradient must be
//!      REJECTED under the weighted objective, having been invisible under the summed one.
//!
//! Step 3 is the important one. A weighted objective that produces non-zero numbers but
//! still cannot reject a wrong answer would be theatre.
//!
//! Run on ndarray, which is known correct, so any rejection is attributable to the
//! deliberate corruption rather than to a backend defect.
//!
//! ```text
//! cargo test --features burn-ndarray --test observability --release -- --nocapture
//! ```
#![cfg(feature = "burn")]

use burn_tensor::activation;
use burn_tensor::{Tensor, TensorData};
use gradcheck::{adapters::burn::Burn, gradcheck, gradcheck_corrupted, Config, Verdict};

fn device() -> burn_tensor::Device {
    burn_tensor::Device::default().autodiff()
}

/// Deterministic weights. `w = ones` is deliberately included because it IS the degenerate
/// case — the summed objective is the weighted objective at its worst possible weight.
fn weights(kind: &str, n: usize) -> Vec<f32> {
    match kind {
        "ones" => vec![1.0; n],
        // ramp, L-infinity normalised so tolerance behaviour does not scale with length
        "ramp" => (0..n).map(|i| (i + 1) as f32 / n as f32).collect(),
        // deterministic +/-1, no zeros
        "rademacher" => (0..n)
            .map(|i| {
                if (i * 2654435761usize) % 2 == 0 {
                    1.0
                } else {
                    -1.0
                }
            })
            .collect(),
        // structured zeros: reaches zero-compaction paths that dense weights cannot
        "zeros3" => (0..n)
            .map(|i| {
                if i % 3 == 0 {
                    0.0
                } else if i % 2 == 0 {
                    1.0
                } else {
                    -1.0
                }
            })
            .collect(),
        _ => unreachable!(),
    }
}

fn ramp_data(n: usize) -> Vec<f64> {
    (0..n)
        .map(|i| ((i as f64) * 0.613).sin() * 1.3 + ((i as f64) * 0.271).cos() * 0.7 + 0.11)
        .collect()
}

/// `⟨w, softmax(x)⟩` expressed as a plain function, so `sum(f_w(x))` is the weighted
/// objective and no API change is needed.
fn weighted_softmax(kind: &'static str, shape: [usize; 2]) -> impl Fn(Tensor<2>) -> Tensor<2> {
    move |t: Tensor<2>| {
        let n = shape[0] * shape[1];
        let w =
            Tensor::<2>::from_data(TensorData::new(weights(kind, n), shape.to_vec()), &device());
        activation::softmax(t, 1) * w
    }
}

#[test]
fn step_1_reproduce_the_defect() {
    let shape = [4usize, 8];
    let d = ramp_data(32);
    let cfg = Config::f32_defaults();

    // the historical objective: plain sum
    let r = gradcheck::<Burn<2>, _>(
        "softmax_summed",
        &d,
        &shape,
        |t| activation::softmax(t, 1),
        &cfg,
    );
    let (checked, total) = r.checked_fraction();
    let max_abs = r.analytic.iter().fold(0.0f64, |a, v| a.max(v.abs()));

    println!(
        "OBS summed      verdict={:?} checked={}/{} max|analytic|={:.3e}",
        r.verdict, checked, total, max_abs
    );
    assert!(
        max_abs < 1e-6,
        "the summed objective should give an identically-zero gradient, got {max_abs:e}"
    );
    assert!(
        !r.passed(),
        "an unobservable objective must not certify; got {:?}",
        r.verdict
    );
}

#[test]
fn step_2_weighted_objective_restores_observability() {
    let shape = [4usize, 8];
    let d = ramp_data(32);
    let cfg = Config::f32_defaults();

    for kind in ["ones", "ramp", "rademacher", "zeros3"] {
        let r = gradcheck::<Burn<2>, _>(
            &format!("softmax_w_{kind}"),
            &d,
            &shape,
            weighted_softmax(kind, shape),
            &cfg,
        );
        let (checked, total) = r.checked_fraction();
        let max_abs = r.analytic.iter().fold(0.0f64, |a, v| a.max(v.abs()));
        println!(
            "OBS w={:<11} verdict={:?} checked={}/{} max|analytic|={:.3e}",
            kind, r.verdict, checked, total, max_abs
        );

        if kind == "ones" {
            // The control: ones reproduces the degenerate case exactly.
            assert!(
                max_abs < 1e-6,
                "w=ones IS the summed objective; it must stay unobservable"
            );
        } else {
            assert!(
                max_abs > 1e-3,
                "w={kind} should produce a non-trivial gradient, got {max_abs:e}"
            );
            assert!(
                checked > 0,
                "w={kind} should adjudicate at least one component"
            );
        }
    }
}

#[test]
fn step_3_the_weighted_objective_can_actually_reject() {
    let shape = [4usize, 8];
    let d = ramp_data(32);
    let cfg = Config::f32_defaults();

    // A x1.5 corruption. Under the summed objective this is invisible: 1.5 * 0 == 0.
    let summed = gradcheck_corrupted::<Burn<2>, _>(
        "softmax_summed_corrupted",
        &d,
        &shape,
        |t| activation::softmax(t, 1),
        &cfg,
        1.5,
    );
    println!("OBS corrupt summed   verdict={:?}", summed.verdict);
    assert_ne!(
        summed.verdict,
        Verdict::Mismatch,
        "a corruption SHOULD be invisible under the unobservable objective -- \
         if this now rejects, the premise of the design note is wrong"
    );

    // Under a non-degenerate weight the same corruption must be caught.
    for kind in ["ramp", "rademacher", "zeros3"] {
        let r = gradcheck_corrupted::<Burn<2>, _>(
            &format!("softmax_w_{kind}_corrupted"),
            &d,
            &shape,
            weighted_softmax(kind, shape),
            &cfg,
            1.5,
        );
        println!("OBS corrupt w={kind:<11} verdict={:?}", r.verdict);
        assert_eq!(
            r.verdict,
            Verdict::Mismatch,
            "w={kind}: the weighted objective must REJECT a 1.5x corruption, else it \
             produces numbers without producing detection"
        );
    }
}