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

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

//! Version-boundary probe for the burn-cpu pooling defect.
//!
//! burn `main` pins cubecl by git revision (`b2703c28`), while the PUBLISHED
//! `burn-tensor 0.22.0-pre.1` resolves `cubecl 0.11.0-pre.1` from the registry.
//! For burn#5304 that distinction mattered: the defect exists only against the
//! pinned revision. This file asks the same question about the pooling defect.
//!
//! The check is the sentinel technique: fill a large allocation with an
//! unmistakable value and drop it, so the allocator is likely to hand that block
//! to the pooling kernel. If the sentinel comes back inside a gradient, the
//! output buffer was never written.
//!
//! ```text
//! cargo test --features burn-ndarray --test pool_version_boundary -- --nocapture
//! cargo test --features burn-cpu     --test pool_version_boundary -- --nocapture
//! ```
#![cfg(feature = "burn")]

use burn_tensor::module::avg_pool1d;
use burn_tensor::{Tensor, TensorData};

const SENTINEL: f32 = 1234.5;

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

/// Churn the allocator with a recognisable value, then let it go.
fn poison() -> f32 {
    let big = vec![SENTINEL; 8192];
    let t = Tensor::<2>::from_data(TensorData::new(big, [64, 128]), &device());
    let s: Vec<f32> = (t.clone() * t).sum().to_data().to_vec().unwrap();
    s[0]
}

fn pool_grad(c: usize) -> Vec<f32> {
    let len = 6usize;
    let n = c * len;
    let d: Vec<f32> = (0..n).map(|i| (i as f32) * 0.1 + 0.5).collect();
    let x = Tensor::<3>::from_data(TensorData::new(d, [1, c, len]), &device()).require_grad();
    let g = avg_pool1d(x.clone(), 3, 2, 1, true, false).sum().backward();
    x.grad(&g).unwrap().to_data().to_vec::<f32>().unwrap()
}

#[test]
fn published_burn_pooling_does_not_leak_uninitialised_memory() {
    // truth for this configuration, identical for every channel:
    // padded length 8, kernel 3, stride 2 -> windows at 0, 2, 4
    let expected = [
        1.0 / 3.0,
        2.0 / 3.0,
        1.0 / 3.0,
        2.0 / 3.0,
        1.0 / 3.0,
        1.0 / 3.0,
    ];

    let mut leaked = vec![];
    let mut wrong = vec![];

    for c in [1usize, 2, 3, 4] {
        let _ = poison();
        let g = pool_grad(c);

        let has_sentinel = g.iter().any(|v| (v - SENTINEL).abs() < 1e-3);
        let matches_truth = g
            .iter()
            .enumerate()
            .all(|(i, v)| (v - expected[i % 6]).abs() < 1e-4);

        println!(
            "BOUNDARY c={c} sentinel_leaked={has_sentinel} correct={matches_truth} grad={:?}",
            &g[..g.len().min(6)]
        );
        if has_sentinel {
            leaked.push(c);
        }
        if !matches_truth {
            wrong.push(c);
        }
    }

    println!("BOUNDARY summary leaked_at={leaked:?} wrong_at={wrong:?}");
    assert!(
        leaked.is_empty(),
        "uninitialised memory leaked into the gradient at channel counts {leaked:?}"
    );
    assert!(
        wrong.is_empty(),
        "gradient incorrect at channel counts {wrong:?}"
    );
}