gradcheck

Finite-difference gradient checking for Rust ML frameworks.

By Henos D (@4ktLuffy) · henosd19@gmail.com

PyTorch has torch.autograd.gradcheck. JAX has check_grads. Rust has had nothing like it, which is roughly why I wrote this.

Here's the bug that convinced me to package it up properly:

burn#5304 — on burn's CPU backend, a.transpose().matmul(b) returns wrong numbers for certain shapes. Nothing crashes, nothing warns.

cpu                 -> [2.8, 2.8, -4.5, -4.5,  4.0,  4.0]
ndarray/flex/metal  -> [0.3, 0.3,  0.6,  0.6, -1.1, -1.1]   <- correct

burn computes weight gradients as matmul(transpose(lhs), grad), so this quietly corrupts training. 33 of 512 shape combinations were affected, all sharing a parity condition. The maintainers confirmed it. It had been living in a framework with a large test suite.

Why compare against numbers, not another backend

Test CPU against GPU and you only catch the bugs they don't share. When both are wrong the same way, the diff comes back clean and you walk away happy.

A central difference doesn't have that problem. The reference comes from the definition of a derivative rather than from somebody's implementation:

∂f/∂xᵢ  ≈  ( f(x + h·eᵢ) − f(x − h·eᵢ) ) / 2h

There's no code in there that can be broken the same way your autodiff is. The whole crate is that one comparison, plus the discipline of admitting when it proves nothing:

                        your function f, at input x
                                     │
                 ┌───────────────────┴───────────────────┐
                 │                                       │
        the framework's gradient              a central difference
         (its own backward pass)          ( f(x+h) − f(x−h) ) / 2h
                 │                                       │
          via the Backend trait                 no framework involved,
          — burn, or your own —                    just arithmetic
                 │                                       │
                 └───────────────────┬───────────────────┘
                                     │
                          compared component by component
                                     │
             ┌───────────────────────┼───────────────────────┐
             │                       │                       │
           Pass                  Mismatch                Unchecked
    agrees within the      exceeds the tolerance     the gradient is too
    relative tolerance      at any magnitude          small to certify —
                                                     absence of evidence,
                                                     reported as such

The left branch is the only part that touches a framework. That separation is what stops the oracle from inheriting the bug it's meant to catch, and it's why adding a framework means implementing one trait rather than modifying the checker.

Run the negative control first

I have written harnesses that computed nothing at all and cheerfully reported everything as passing. A clean run out of one of those is worse than no run, so the control is an actual function in the API instead of something you're trusted to remember:

gradcheck::assert_detects_wrong_gradient::<MyBackend>();

It hands the checker a deliberately corrupted gradient and panics if the checker doesn't notice, then hands it a correct one and panics if that gets flagged. Both directions, because a checker that complains about everything is just as useless as one that complains about nothing.

Put it in the same test binary as your real checks. If it ever fails, ignore every other result in that run.

Quick start

cargo add gradcheck --dev

[dev-dependencies]
gradcheck = "0.1"

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

#[test]
fn gradients_are_correct() {
    // 1. Prove the instrument works.
    assert_detects_wrong_gradient::<MyBackend>();

    // 2. Then trust its verdicts.
    gradcheck::<MyBackend, _>(
        "tanh",
        &[0.7, -1.3, 2.1, -0.4],
        &[2, 2],
        |t| t.tanh(),
        &Config::default(),
    )
    .assert_pass();
}

Check more than one shape

The burn bug above showed up in 33 of 512 shape combinations. Pick a single shape and you had roughly a 94% chance of seeing nothing wrong, so sweep instead:

use gradcheck::shape_sweep;

let shapes = vec![vec![2, 3], vec![3, 2], vec![4, 4], vec![8, 2], vec![5, 7]];
shape_sweep::<MyBackend, _, _>(
    "matmul_self_transpose",
    &shapes,
    |n| (0..n).map(|i| ((i as f64) * 0.37).sin() * 1.7 + 0.15).collect(),
    |t| t.clone().matmul(t.transpose()),
    &Config::default(),
)
.assert_all_pass();

Using it with burn

burn already has an adapter, so you don't need to write one. Pick the compute backend with a feature flag:

[dev-dependencies]
gradcheck = { version = "0.1", features = ["burn-ndarray"] }

burn-ndarray, burn-flex and burn-cpu select which backend burn-tensor compiles against. The adapter itself is generic over tensor rank: Burn<2> checks rank-2 functions, Burn<4> checks rank-4 ones like convolution and pooling.

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

#[test]
fn burn_gradients_are_correct() {
    assert_detects_wrong_gradient::<Burn<2>>();

    let cfg = Config::f32_defaults();   // burn computes in f32
    gradcheck::<Burn<2>, _>("tanh", &[0.7, -1.3, 2.1, -0.4], &[2, 2], |t| t.tanh(), &cfg)
        .assert_pass();
}

Run the crate's own burn suite:

cargo test --features burn-ndarray --test burn_adapter

The same checks against a different backend is the whole trick — compile twice, compare. A gradient bug that only affects one backend is invisible until you do:

cargo test --features burn-flex --test burn_adapter
cargo test --features burn-cpu   --test burn_adapter

Two things worth knowing before you read a result. burn computes in f32, so keep step at 1e-3 or larger — smaller steps drown in cancellation. And quadratic functions like t.transpose().matmul(t) want 1e-1; there's a worked measurement in the header of tests/burn_adapter.rs showing why.

A project you can copy

example-project/ is a complete, standalone crate that depends on gradcheck over git the way yours would — no path dependency, nothing reaching back into this repo. Copy the directory, or just run it:

cd example-project
cargo test -- --nocapture

It covers the negative control, a single check, a shape sweep, and one case that finds a real defect when you point it at burn-cpu.

Adding a framework

Implement [Backend] once: build a tensor, read it back, run the function forward, and hand back the framework's own gradient. The core never sees framework types. Adapters convert at the boundary, which is what stops the oracle from quietly depending on the thing it is supposed to be testing.

impl Backend for MyBackend {
    type Tensor = my_framework::Tensor;
    fn name() -> &'static str { "my_framework" }
    fn from_slice(data: &[f64], shape: &[usize]) -> Self::Tensor { /* ... */ }
    fn to_vec(t: &Self::Tensor) -> Vec<f64> { /* ... */ }
    fn forward_sum(f: &dyn Fn(Self::Tensor) -> Self::Tensor, x: &Self::Tensor) -> f64 { /* ... */ }
    fn analytic_grad(f: &dyn Fn(Self::Tensor) -> Self::Tensor, x: &Self::Tensor) -> Vec<f64> { /* ... */ }
}

Picking a step size

Two errors pull against each other. Truncation error shrinks with h (as O(h²)), while floating-point cancellation grows. 1e-3 is a sane default for f32 frameworks, 1e-5 for f64.

If a check fails by a small margin, vary h and look again before you believe it. A real bug won't care about the step size. A tolerance artifact will move.

Before reporting an f16/bf16 mismatch

Check both sides against f32 first. The fast path is not automatically the wrong one.

I nearly filed a bug against Apple's MLX over three configurations where mx.fast.rope disagreed with MLX's own reference implementation. Measured against f32, the kernel turned out to be about 60× more accurate than the reference. The reference was the broken side.

Status

Early, and I'd rather say so than oversell it. The core, the negative control, the shape sweep and the reporting all work, and they're self-tested against a reference backend built on forward-mode dual numbers. Dual numbers are a completely different algorithm from central differences, so when the two agree it actually means something. cargo test runs with no ML framework installed.

Adapters are where I'd most like help. burn is in, everything else is open.

Author

Henos D — GitHub @4ktLuffy · henosd19@gmail.com

Licence

MIT OR Apache-2.0