gradcheck 0.1.0

Finite-difference gradient checking for Rust ML frameworks. Verifies an autodiff engine against an independent numerical oracle, with a negative control that must fail.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

// 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

//! Adapter for the [burn](https://burn.dev) framework.
//!
//! Enable with one of the `burn-*` features, which also picks the compute backend:
//! `burn-ndarray`, `burn-flex` or `burn-cpu`.
//!
//! ```toml
//! [dev-dependencies]
//! gradcheck = { version = "0.1", features = ["burn-ndarray"] }
//! ```
//!
//! The adapter is generic over tensor rank, so [`Burn<2>`] checks rank-2 functions,
//! [`Burn<4>`] checks rank-4 (convolutions, pooling), and so on.
//!
//! ```ignore
//! use gradcheck::{adapters::burn::Burn, assert_detects_wrong_gradient, gradcheck, Config};
//!
//! assert_detects_wrong_gradient::<Burn<2>>();
//!
//! gradcheck::<Burn<2>, _>(
//!     "tanh",
//!     &[0.7, -1.3, 2.1, -0.4],
//!     &[2, 2],
//!     |t| t.tanh(),
//!     &Config::default(),
//! )
//! .assert_pass();
//! ```
//!
//! # Note on precision
//!
//! burn computes in `f32` by default. Values cross the adapter boundary as `f64` and are
//! narrowed on the way in, so keep [`Config::step`](crate::Config::step) at `1e-3` or larger —
//! smaller steps are swamped by `f32` cancellation and produce spurious mismatches.

use burn_tensor::{Tensor, TensorData};

use crate::Backend;

/// The burn framework, at tensor rank `D`.
///
/// The concrete compute backend (ndarray, flex, cpu, metal, wgpu, …) is selected by
/// `burn-tensor`'s own feature flags, not by this type — so the same check runs against
/// whichever backend the test binary was compiled with. That is what makes it useful for
/// finding backend-specific defects: compile the same suite twice and compare.
#[derive(Debug, Clone, Copy)]
pub struct Burn<const D: usize>;

impl<const D: usize> Burn<D> {
    fn device() -> burn_tensor::Device {
        burn_tensor::Device::default().autodiff()
    }
}

impl<const D: usize> Backend for Burn<D> {
    type Tensor = Tensor<D>;

    fn name() -> &'static str {
        "burn"
    }

    fn from_slice(data: &[f64], shape: &[usize]) -> Self::Tensor {
        assert_eq!(
            shape.len(),
            D,
            "gradcheck/burn: shape {shape:?} has rank {} but this adapter is Burn<{D}>",
            shape.len()
        );
        let values: Vec<f32> = data.iter().map(|&v| v as f32).collect();
        Tensor::from_data(TensorData::new(values, shape.to_vec()), &Self::device())
    }

    fn to_vec(t: &Self::Tensor) -> Vec<f64> {
        t.clone()
            .to_data()
            .to_vec::<f32>()
            .expect("gradcheck/burn: could not read tensor data as f32")
            .into_iter()
            .map(|v| v as f64)
            .collect()
    }

    fn forward_sum(f: &dyn Fn(Self::Tensor) -> Self::Tensor, x: &Self::Tensor) -> f64 {
        // Deliberately no `require_grad` and no `backward` here: the numerical oracle must
        // not touch the autodiff machinery it is meant to be checking.
        let out: Vec<f32> = f(x.clone())
            .sum()
            .to_data()
            .to_vec::<f32>()
            .expect("gradcheck/burn: could not read the forward result");
        out[0] as f64
    }

    fn analytic_grad(f: &dyn Fn(Self::Tensor) -> Self::Tensor, x: &Self::Tensor) -> Vec<f64> {
        let tracked = x.clone().require_grad();
        let grads = f(tracked.clone()).sum().backward();
        let g = tracked.grad(&grads).expect(
            "gradcheck/burn: no gradient was recorded for the input. The tensor was marked \
             require_grad, so this usually means the function under test detached it (for \
             example by going through `into_data` or a non-differentiable op).",
        );
        g.to_data()
            .to_vec::<f32>()
            .expect("gradcheck/burn: could not read the gradient as f32")
            .into_iter()
            .map(|v| v as f64)
            .collect()
    }
}