luma_tensor/grad/
global.rs

1//! Global autograd enable/disable switch, mirroring PyTorch's
2//! `torch.set_grad_enabled` / `torch.no_grad()`.
3//!
4//! Controlled via [`NoGradGuard`] (RAII) or the [`no_grad!`] macro.
5
6use std::cell::Cell;
7
8thread_local! {
9    static GRAD_ENABLED: Cell<bool> = Cell::new(true);
10}
11
12/// Set the global gradient enable flag.
13pub fn set_grad_enabled(enabled: bool) {
14    GRAD_ENABLED.with(|c| c.set(enabled));
15}
16
17/// Check whether gradient tracking is currently active.
18pub fn is_grad_enabled() -> bool {
19    GRAD_ENABLED.with(|c| c.get())
20}
21
22/// RAII guard that disables gradient tracking for the duration of its
23/// lifetime, restoring the previous setting on drop.
24///
25/// # Example
26/// ```ignore
27/// {
28///     let _guard = NoGradGuard::new();
29///     // ops here don't build the computation graph
30/// }
31/// // grad tracking is restored
32/// ```
33pub struct NoGradGuard {
34    prev: bool,
35}
36
37impl NoGradGuard {
38    pub fn new() -> Self {
39        let prev = is_grad_enabled();
40        set_grad_enabled(false);
41        Self { prev }
42    }
43}
44
45impl Default for NoGradGuard {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl Drop for NoGradGuard {
52    fn drop(&mut self) {
53        set_grad_enabled(self.prev);
54    }
55}
56
57/// Temporarily disable gradient tracking within the current scope.
58///
59/// Equivalent to `{ let _guard = NoGradGuard::new(); ... }`.
60///
61/// # Example
62/// ```ignore
63/// no_grad!();
64/// // .backward() will not traverse beyond this scope
65/// ```
66#[macro_export]
67macro_rules! no_grad {
68    () => {
69        let _guard = $crate::NoGradGuard::new();
70    };
71}
luma_tensor/grad/global.rs

luma_tensor/grad/
global.rs
