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

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

//! The independent oracle: central-difference numerical gradients.

use crate::Backend;

/// The precision a backend actually computes in.
///
/// This matters because the *realized* step depends on it, and using the wrong one
/// injects error rather than measuring it: rounding an `f64` backend's perturbation
/// through `f32` produces a ~4e-4 relative discrepancy that looks exactly like a
/// gradient bug.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Precision {
    /// The framework stores and computes in `f32` (most GPU/accelerator backends).
    F32,
    /// The framework stores and computes in `f64`.
    F64,
}

/// The step that a perturbation actually achieves in the backend's working precision.
///
/// Requesting `x + h` does not guarantee moving by `h`. At `x = 1e6` in `f32` the ULP is
/// `0.0625`, so a step of `1e-3` rounds away entirely and `x + h == x`. The difference
/// quotient is then `0/0`-flavoured nonsense that looks like a perfectly clean zero
/// gradient — a false pass produced by arithmetic, not by the framework.
///
/// Callers should divide by the *realized* step, and treat a realized step of zero as
/// invalid rather than as a measurement.
pub fn realized_step_with(x: f64, h: f64, p: Precision) -> f64 {
    match p {
        Precision::F32 => {
            let up = (x + h) as f32 as f64;
            let down = (x - h) as f32 as f64;
            (up - down) / 2.0
        }
        Precision::F64 => ((x + h) - (x - h)) / 2.0,
    }
}

/// [`realized_step_with`] at `f32`, the common case.
pub fn realized_step(x: f64, h: f64) -> f64 {
    realized_step_with(x, h, Precision::F32)
}

/// Central-difference gradient of `sum(f(x))` with respect to every element of `x`.
///
/// For each component `i`:
///
/// ```text
/// g_i = ( sum(f(x + h·e_i)) - sum(f(x - h·e_i)) ) / 2h
/// ```
///
/// Central difference is used rather than forward difference because its truncation error
/// is `O(h²)` instead of `O(h)`.
///
/// The divisor is the **realized** step, not the requested one, so a perturbation that is
/// partly lost to rounding does not silently scale the result.
///
/// Cost is `2N` forward evaluations for `N` elements — keep test tensors small.
pub fn numerical_gradient<B, F>(
    f: &F,
    data: &[f64],
    shape: &[usize],
    step: f64,
    precision: Precision,
) -> Vec<f64>
where
    B: Backend,
    F: Fn(B::Tensor) -> B::Tensor,
{
    assert!(step > 0.0, "gradcheck: step must be positive, got {step}");

    let mut out = Vec::with_capacity(data.len());
    let mut probe = data.to_vec();

    for i in 0..data.len() {
        let original = probe[i];
        let realized = realized_step_with(original, step, precision);

        probe[i] = original + step;
        let up = B::forward_sum(f, &B::from_slice(&probe, shape));

        probe[i] = original - step;
        let down = B::forward_sum(f, &B::from_slice(&probe, shape));

        probe[i] = original;

        if realized == 0.0 {
            // Unrepresentable; the caller's gate turns this into Invalid{Probe, ...}.
            out.push(f64::NAN);
        } else {
            out.push((up - down) / (2.0 * realized));
        }
    }

    out
}