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

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

//! B5 — is second-order differentiation expressible at all in burn 0.22's API?
//!
//! Note the architectural change that matters here: autodiff is now a property of the
//! DEVICE (`Device::default().autodiff()`), not a type wrapper (`Autodiff<B>`). The
//! tensor type is rank-only (`Tensor<D>`) with no backend parameter, so `Autodiff<Autodiff<B>>`
//! is not expressible as a type. The question becomes whether the device can be nested.
//!
//! Ground truth: f(x) = x^3  ->  f'(x) = 3x^2  ->  f''(x) = 6x.  At x = 2: f'=12, f''=12.
//!
//! ```text
//! cargo test --features burn-ndarray --test second_order --release -- --nocapture
//! ```
#![cfg(feature = "burn")]

use burn_tensor::{Tensor, TensorData};

#[test]
fn first_derivative_is_the_baseline() {
    let dev = burn_tensor::Device::default().autodiff();
    let x = Tensor::<1>::from_data(TensorData::new(vec![2.0f32], vec![1]), &dev).require_grad();
    let y = x.clone() * x.clone() * x.clone();
    let g = y.sum().backward();
    let d: Vec<f32> = x.grad(&g).unwrap().to_data().to_vec().unwrap();
    println!("B5 f(x)=x^3 at x=2 -> f'  = {:?}   (expected [12.0])", d);
    assert!((d[0] - 12.0).abs() < 1e-3);
}

/// Answer, on burn 0.22.0-pre.1: no. `Device::autodiff()` panics when applied twice.
///
/// Asserted rather than merely observed, so that this starts failing the day burn gains
/// second-order support — at which point gradcheck can grow a second-order checker.
#[test]
#[should_panic(expected = "Only first-order autodiff is supported")]
fn can_the_autodiff_device_be_nested() {
    let once = burn_tensor::Device::default().autodiff();
    let twice = once.clone().autodiff();
    println!("B5 device.autodiff()            = {once:?}");
    println!("B5 device.autodiff().autodiff() = {twice:?}");
    println!(
        "B5 nesting changes the device   = {}",
        format!("{once:?}") != format!("{twice:?}")
    );
}

/// Can a gradient be obtained and then differentiated again?
///
/// Answer, on burn 0.22.0-pre.1: no — it panics before reaching the second backward, for the
/// same reason as above. Asserted so the limitation is documented by a passing test rather
/// than by a failing one.
#[test]
#[should_panic(expected = "Only first-order autodiff is supported")]
fn second_derivative_attempt() {
    let dev = burn_tensor::Device::default().autodiff().autodiff();
    let x = Tensor::<1>::from_data(TensorData::new(vec![2.0f32], vec![1]), &dev).require_grad();

    let y = x.clone() * x.clone() * x.clone();
    let g1 = y.sum().backward();
    let d1 = x.grad(&g1).expect("first gradient must exist");

    let d1v: Vec<f32> = d1.clone().to_data().to_vec().unwrap();
    println!("B5 first grad = {d1v:?}");

    // Is the returned gradient itself tracked? If it is, we can differentiate it.
    println!("B5 grad.require_grad? -> attempting second backward");
    let g2 = d1.sum().backward();
    match x.grad(&g2) {
        Some(t) => {
            let v: Vec<f32> = t.to_data().to_vec().unwrap();
            println!("B5 SECOND DERIVATIVE = {v:?}   (expected [12.0])");
        }
        None => println!("B5 SECOND DERIVATIVE = None -- the returned gradient is not tracked"),
    }
}