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

//! Gradient checks against the [burn](https://burn.dev) framework.
//!
//! Run with the default (ndarray) backend:
//!
//! ```text
//! cargo test --features burn-ndarray --test burn_adapter
//! ```
//!
//! The compute backend is chosen by `burn-tensor`'s feature flags, so the *same* checks can be
//! pointed at a different backend by changing features. That is the point: a gradient bug that
//! is backend-specific is invisible until you run the identical suite on both.
#![cfg(feature = "burn")]

use gradcheck::{
    adapters::burn::Burn, assert_detects_wrong_gradient, gradcheck, shape_sweep, Config,
};

/// burn computes in `f32`, so the step cannot be very small without cancellation dominating.
fn cfg() -> Config {
    Config::f32_defaults()
}

/// Deterministic, well-conditioned data. Values avoid 0 so that `abs`/`max`-style kinks and
/// division poles do not invalidate the finite-difference approximation.
fn ramp(n: usize, seed: f64) -> Vec<f64> {
    (0..n)
        .map(|i| ((i as f64) * 0.37 + seed).sin() * 1.7 + 0.15)
        .collect()
}

/// Config for the `t·tᵀ` sweeps, which need a larger step than the elementwise ops.
///
/// These functions are quadratic in the input, so `f32` cancellation dominates early — the
/// same reason [`Config`] recommends `1e-1` for `square`. Measured on ndarray at `[6, 3]`,
/// `worst_rel` moves monotonically with the step:
///
/// ```text
/// h=1e-1 -> 1.17e-5    h=1e-2 -> 1.47e-3
/// h=5e-2 -> 1.36e-4    h=5e-3 -> 4.15e-3
/// h=2e-2 -> 5.37e-4    h=1e-3 -> 1.70e-2
/// ```
///
/// Error that grows as the step shrinks is cancellation, not a defect (a real defect is
/// insensitive to `h`). At `1e-1` every shape is fully checked with two to three orders of
/// margin, and a gradient corrupted by 1.5x is still rejected at this step — verified before
/// adopting it, since a step large enough to hide the noise could also hide a real fault.
fn matmul_cfg() -> Config {
    Config::f32_defaults().with_step(1e-1)
}

/// **Run first.** If this fails, every other result in this file is meaningless.
#[test]
fn negative_control_detects_a_wrong_gradient() {
    assert_detects_wrong_gradient::<Burn<2>>();
}

#[test]
fn elementwise_gradients_are_correct() {
    let d = ramp(6, 0.3);
    let s = [2usize, 3];
    let c = cfg();

    gradcheck::<Burn<2>, _>("tanh", &d, &s, |t| t.tanh(), &c).assert_pass();
    gradcheck::<Burn<2>, _>("exp", &d, &s, |t| t.exp(), &c).assert_pass();
    gradcheck::<Burn<2>, _>("sin", &d, &s, |t| t.sin(), &c).assert_pass();
    gradcheck::<Burn<2>, _>("erf", &d, &s, |t| t.erf(), &c).assert_pass();
    gradcheck::<Burn<2>, _>("square", &d, &s, |t| t.clone() * t, &c).assert_pass();
    gradcheck::<Burn<2>, _>("log_positive", &d, &s, |t| (t.abs() + 0.5).log(), &c).assert_pass();
    gradcheck::<Burn<2>, _>("sqrt_positive", &d, &s, |t| (t.abs() + 0.5).sqrt(), &c).assert_pass();
    gradcheck::<Burn<2>, _>("chained", &d, &s, |t| t.tanh().exp(), &c).assert_pass();
}

#[test]
fn reduction_gradients_are_correct() {
    let d = ramp(6, 0.3);
    let s = [2usize, 3];
    let c = cfg();

    gradcheck::<Burn<2>, _>("sum_dim0", &d, &s, |t| t.sum_dim(0), &c).assert_pass();
    gradcheck::<Burn<2>, _>("sum_dim1", &d, &s, |t| t.sum_dim(1), &c).assert_pass();
    gradcheck::<Burn<2>, _>("mean_dim1", &d, &s, |t| t.mean_dim(1), &c).assert_pass();
    gradcheck::<Burn<2>, _>("cumsum", &d, &s, |t| t.cumsum(1), &c).assert_pass();
}

#[test]
fn broadcast_and_view_gradients_are_correct() {
    let d = ramp(6, 0.3);
    let s = [2usize, 3];
    let c = cfg();

    // Backward must sum over the broadcast axis.
    gradcheck::<Burn<2>, _>(
        "broadcast_by_colsum",
        &d,
        &s,
        |t| {
            let col = t.clone().sum_dim(1);
            t * col
        },
        &c,
    )
    .assert_pass();

    gradcheck::<Burn<2>, _>(
        "transpose_roundtrip",
        &d,
        &s,
        |t| t.transpose().transpose(),
        &c,
    )
    .assert_pass();

    // `slice` drops the last column, so 2 of 6 inputs cannot affect the output and their
    // gradient is zero by construction. Those components abstain rather than certify — the
    // behaviour `docs/objective-observability.md` argues for — so the honest assertion here
    // is `assert_no_mismatch`, not `assert_pass`. The 4 observable components are checked
    // exactly (worst_rel = 0.0); demanding a Pass would mean demanding certification of
    // components carrying no signal.
    gradcheck::<Burn<2>, _>("slice", &d, &s, |t| t.slice([0..2, 0..2]), &c).assert_no_mismatch();
}

/// Regression fixture for [burn#5304](https://github.com/tracel-ai/burn/issues/5304).
///
/// `a.transpose().matmul(b)` returns numerically wrong values on burn's **cpu** backend for a
/// subset of shapes — 33 of 512 combinations, all with `m` even, `n` even and `k != m`. Because
/// burn computes weight gradients as `matmul(transpose(lhs), grad)`, the defect silently
/// corrupts training.
///
/// This test **passes on ndarray and flex** (documenting that they are correct) and **fails on
/// the cpu backend**, catching the bug:
///
/// ```text
/// cargo test --features burn-ndarray --test burn_adapter transpose_matmul  # passes
/// cargo test --features burn-cpu     --test burn_adapter transpose_matmul  # the bug
/// ```
///
/// The shape list is deliberately wide. The defect is invisible at most shapes — checking one
/// shape would have found nothing, which is exactly why it survived burn's own test suite.
#[test]
fn transpose_matmul_gradient_across_shapes() {
    let c = matmul_cfg();
    let shapes: Vec<Vec<usize>> = vec![
        vec![2, 3],
        vec![3, 2],
        vec![2, 2],
        vec![4, 4],
        vec![4, 3],
        vec![6, 3],
        vec![8, 2],
        vec![2, 5],
        vec![5, 5],
        vec![3, 8],
    ];

    let sweep = shape_sweep::<Burn<2>, _, _>(
        "self_transpose_matmul",
        &shapes,
        |n| ramp(n, 0.3),
        |t| t.clone().transpose().matmul(t),
        &c,
    );

    for r in &sweep.reports {
        println!("{r}");
    }
    sweep.assert_all_pass();
}

/// The other orientation, which is what `dLHS` uses in burn's matmul backward.
#[test]
fn matmul_self_transpose_gradient_across_shapes() {
    let c = matmul_cfg();
    let shapes: Vec<Vec<usize>> = vec![
        vec![2, 3],
        vec![3, 2],
        vec![4, 4],
        vec![2, 8],
        vec![8, 2],
        vec![5, 7],
    ];

    shape_sweep::<Burn<2>, _, _>(
        "matmul_self_transpose",
        &shapes,
        |n| ramp(n, 0.3),
        |t| t.clone().matmul(t.transpose()),
        &c,
    )
    .assert_all_pass();
}