# gradcheck
**Finite-difference gradient checking for Rust ML frameworks.**
[](https://github.com/4ktLuffy/gradcheck/actions/workflows/ci.yml)
[](#licence)
By **Henos D** ([@4ktLuffy](https://github.com/4ktLuffy)) · <henosd19@gmail.com>
PyTorch has [`torch.autograd.gradcheck`](https://pytorch.org/docs/stable/generated/torch.autograd.gradcheck.html).
JAX has [`check_grads`](https://docs.jax.dev/en/latest/_autosummary/jax.test_util.check_grads.html).
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](https://github.com/tracel-ai/burn/issues/5304)** — on burn's CPU backend,
> `a.transpose().matmul(b)` returns wrong numbers for certain shapes. Nothing crashes,
> nothing warns.
>
> ```text
> 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:
```text
∂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:
```text
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:
```rust
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
```bash
cargo add gradcheck --dev
```
```toml
[dev-dependencies]
gradcheck = "0.1"
```
```rust
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:
```rust
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:
```toml
[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.
```rust
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:
```bash
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:
```bash
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/`](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:
```bash
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.
```rust
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](https://github.com/4ktLuffy) · <henosd19@gmail.com>
## Licence
MIT OR Apache-2.0