tenflowers_core/
gradient_executor.rs

1//! Gradient Executor Trait — Dependency-Inversion Bridge
2//!
3//! `tenflowers-core`'s operation registry only stores gradient function names as
4//! metadata (e.g. `grad_fn: Option<String>`); it never stores callable
5//! forward/backward closures. Real forward+backward execution lives one layer up,
6//! in `tenflowers-autograd`'s `GradientTape` -- which depends on `tenflowers-core`,
7//! not the other way around, so `tenflowers-core` cannot call into it directly.
8//!
9//! This module defines a small trait, [`GradientExecutor`], that lets
10//! `tenflowers-core` call into a real gradient implementation without depending on
11//! `tenflowers-autograd`: a higher-level crate registers an implementation via
12//! [`register_gradient_executor`] at an appropriate initialization point (for
13//! example, `tenflowers_autograd::init()`). [`crate::gradient_validation_framework`]
14//! consults [`get_gradient_executor`] and, if nothing has been registered, returns
15//! an honest "cannot verify" status rather than fabricating a `passed: true`
16//! result.
17
18use crate::{Result, Tensor};
19use std::sync::OnceLock;
20
21/// Bridges `tenflowers-core`'s gradient validation framework to a real
22/// forward+backward implementation living in a higher-level crate (normally
23/// `tenflowers-autograd`).
24///
25/// Implementations always operate in `f64`, regardless of the dtype under test in
26/// [`crate::gradient_validation_framework::GradientTestCase`]: `f64` has enough
27/// precision to be a safe universal choice for a validation/testing framework,
28/// and keeping the trait monomorphic avoids needing a type-erased/generic trait
29/// object design for this bridge.
30pub trait GradientExecutor {
31    /// Compute the gradient(s) of `op` applied to `inputs`, with respect to each
32    /// element of `inputs`, in order.
33    ///
34    /// Returns one gradient tensor per input, each with the same shape as its
35    /// corresponding input tensor. Returns `Err` if `op` is not recognized by
36    /// this executor, or if the gradient computation itself fails -- never
37    /// fabricates a placeholder result.
38    fn compute_gradient(&self, op: &str, inputs: &[Tensor<f64>]) -> Result<Vec<Tensor<f64>>>;
39}
40
41static GRADIENT_EXECUTOR: OnceLock<Box<dyn GradientExecutor + Send + Sync>> = OnceLock::new();
42
43/// Register the global [`GradientExecutor`].
44///
45/// Intended to be called once, early in application or test startup (for example
46/// from `tenflowers_autograd::init()`). Idempotent: if an executor has already
47/// been registered, this call is a silent no-op and the original registration is
48/// kept (matching [`OnceLock`]'s first-write-wins semantics).
49pub fn register_gradient_executor(executor: Box<dyn GradientExecutor + Send + Sync>) {
50    let _ = GRADIENT_EXECUTOR.set(executor);
51}
52
53/// Get the globally registered [`GradientExecutor`], if any has been registered.
54///
55/// Returns `None` if no executor has been registered yet -- callers must not
56/// treat this as "the property holds"; see
57/// [`crate::gradient_validation_framework`] for how this is surfaced honestly.
58pub fn get_gradient_executor() -> Option<&'static (dyn GradientExecutor + Send + Sync)> {
59    GRADIENT_EXECUTOR.get().map(|b| b.as_ref())
60}
tenflowers_core/gradient_executor.rs

tenflowers_core/
gradient_executor.rs
