Struct NoGradTrack 

pub struct NoGradTrack { /* private fields */ }

RAII guard for temporarily disabling gradient tracking

NoGradTrack provides a scope-based mechanism for disabling gradient tracking, automatically restoring the previous state when the guard is dropped. It ensures proper gradient state management even in the presence of early returns or exceptions. A RAII guard that temporarily disables gradient tracking

Similar to PyTorch’s torch.no_grad(), this guard disables gradient computation within its scope and automatically restores the previous gradient tracking state when it goes out of scope.

Performance Benefits

• Prevents computation graph construction during inference
• Reduces memory usage by not storing intermediate values for backpropagation
• Improves computation speed by skipping gradient-related operations

Examples

use train_station::{NoGradTrack, Tensor};

let x = Tensor::ones(vec![3, 3]).with_requires_grad();
let y = Tensor::ones(vec![3, 3]).with_requires_grad();

// Normal computation with gradients
let z1 = x.add_tensor(&y);
assert!(z1.requires_grad());

// Computation without gradients
{
    let _guard = NoGradTrack::new();
    let z2 = x.add_tensor(&y);
    assert!(!z2.requires_grad()); // Gradients disabled
} // Guard drops here, gradients restored

// Gradients are automatically restored
let z3 = x.add_tensor(&y);
assert!(z3.requires_grad());

Nested Contexts

use train_station::{NoGradTrack, is_grad_enabled, Tensor};

assert!(is_grad_enabled());

{
    let _guard1 = NoGradTrack::new();
    assert!(!is_grad_enabled());

    {
        let _guard2 = NoGradTrack::new();
        assert!(!is_grad_enabled());
    } // guard2 drops

    assert!(!is_grad_enabled()); // Still disabled
} // guard1 drops

assert!(is_grad_enabled()); // Restored

Implementations

impl NoGradTrack

pub fn new() -> Self

Create a new NoGradTrack that disables gradient tracking

This function pushes the current gradient state onto the stack and disables gradient tracking. When the guard is dropped, the previous state is automatically restored.

Returns

A new NoGradTrack that will restore gradient state when dropped

Examples found in repository

examples/neural_networks/feedforward_network.rs (line 53)

    pub fn forward_no_grad(input: &Tensor) -> Tensor {
        let _guard = NoGradTrack::new();
        Self::forward(input)
    }
}

/// A basic linear layer implementation (reused from basic_linear_layer.rs)
#[derive(Debug)]
pub struct LinearLayer {
    pub weight: Tensor,
    pub bias: Tensor,
    pub input_size: usize,
    pub output_size: usize,
}

impl LinearLayer {
    pub fn new(input_size: usize, output_size: usize, seed: Option<u64>) -> Self {
        let scale = (1.0 / input_size as f32).sqrt();

        let weight = Tensor::randn(vec![input_size, output_size], seed)
            .mul_scalar(scale)
            .with_requires_grad();
        let bias = Tensor::zeros(vec![output_size]).with_requires_grad();

        Self {
            weight,
            bias,
            input_size,
            output_size,
        }
    }

    pub fn forward(&self, input: &Tensor) -> Tensor {
        let output = input.matmul(&self.weight);
        output.add_tensor(&self.bias)
    }

    pub fn forward_no_grad(&self, input: &Tensor) -> Tensor {
        let _guard = NoGradTrack::new();
        self.forward(input)
    }

    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
        vec![&mut self.weight, &mut self.bias]
    }
}

/// Configuration for feed-forward network
#[derive(Debug, Clone)]
pub struct FeedForwardConfig {
    pub input_size: usize,
    pub hidden_sizes: Vec<usize>,
    pub output_size: usize,
    pub use_bias: bool,
}

impl Default for FeedForwardConfig {
    fn default() -> Self {
        Self {
            input_size: 4,
            hidden_sizes: vec![8, 4],
            output_size: 2,
            use_bias: true,
        }
    }
}

/// A configurable feed-forward neural network
pub struct FeedForwardNetwork {
    layers: Vec<LinearLayer>,
    config: FeedForwardConfig,
}

impl FeedForwardNetwork {
    /// Create a new feed-forward network with the given configuration
    pub fn new(config: FeedForwardConfig, seed: Option<u64>) -> Self {
        let mut layers = Vec::new();
        let mut current_size = config.input_size;
        let mut current_seed = seed;

        // Create hidden layers
        for &hidden_size in &config.hidden_sizes {
            layers.push(LinearLayer::new(current_size, hidden_size, current_seed));
            current_size = hidden_size;
            current_seed = current_seed.map(|s| s + 1);
        }

        // Create output layer
        layers.push(LinearLayer::new(
            current_size,
            config.output_size,
            current_seed,
        ));

        Self { layers, config }
    }

    /// Forward pass through the entire network
    pub fn forward(&self, input: &Tensor) -> Tensor {
        let mut x = input.clone();

        // Pass through all layers except the last one with ReLU activation
        for layer in &self.layers[..self.layers.len() - 1] {
            x = layer.forward(&x);
            x = ReLU::forward(&x);
        }

        // Final layer without activation (raw logits)
        if let Some(final_layer) = self.layers.last() {
            x = final_layer.forward(&x);
        }

        x
    }

    /// Forward pass without gradients (for inference)
    pub fn forward_no_grad(&self, input: &Tensor) -> Tensor {
        let _guard = NoGradTrack::new();
        self.forward(input)
    }

examples/neural_networks/basic_linear_layer.rs (line 79)

    pub fn forward_no_grad(&self, input: &Tensor) -> Tensor {
        let _guard = NoGradTrack::new();
        self.forward(input)
    }

Trait Implementations

impl Default for NoGradTrack

fn default() -> Self

Returns the “default value” for a type.

impl Drop for NoGradTrack

fn drop(&mut self)

Automatically restore the previous gradient tracking state

This ensures that gradient tracking is properly restored even if the guard goes out of scope due to early returns or panics.