Struct Tensor 

pub struct Tensor { /* private fields */ }

High-performance multi-dimensional tensor with automatic differentiation support

The core data structure for machine learning operations, designed for maximum performance with zero-cost abstractions. Supports arbitrary dimensionality, SIMD optimization, gradient tracking, device placement, and natural mathematical expressions through operator overloading.

Key Features

• Raw Pointer Storage: Zero-overhead memory access for maximum performance
• SIMD Optimization: AVX2 alignment and vectorized operations
• Memory Efficiency: Optimized alignment strategies for different tensor sizes
• gradtrack Integration: Built-in gradient tracking and computation
• Device Support: CPU and future CUDA device placement
• View Tensors: Zero-copy tensor views with shared memory management
• Thread Safety: Send + Sync implementation for concurrent usage
• Operator Overloading: Natural mathematical expressions (+, -, *, /, +=, -=, *=, /=)

Memory Layout

Tensors use row-major memory layout with size-dependent alignment:

• Small tensors (≤8 elements): 16-byte SSE alignment
• Medium tensors (8-1024 elements): 32-byte AVX2 alignment
• Large tensors (>1024 elements): 64-byte cache-line alignment

Performance Characteristics

• Memory Overhead: ~64 bytes per tensor (excluding data)
• SIMD Ready: Properly aligned for vectorized operations
• Cache Friendly: Optimized memory layout for CPU cache hierarchies
• Zero-Cost Views: View tensors share memory without copying
• Thread Safe: Atomic ID generation and lock-free operations
• Operator Performance: Zero-cost operator overloading for mathematical expressions

Safety

This struct uses unsafe code for performance. The following invariants must be maintained:

• data must be valid for shape.size elements
• data must be properly aligned for f32
• data must not be aliased while the tensor exists
• shape.size must match the actual allocated memory
• allocation_owner must be valid if present

Examples

Basic Tensor Operations

use train_station::Tensor;

// Create tensors with different configurations
let tensor = Tensor::new(vec![2, 3]);
let tensor_with_grad = Tensor::ones(vec![10, 10]).with_requires_grad();

// Access tensor properties
assert_eq!(tensor.size(), 6);
assert_eq!(tensor.shape().dims(), vec![2, 3]);
assert!(tensor.is_contiguous());

Operator Overloading

use train_station::Tensor;

// Create tensors for operations
let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();

// Tensor operations with operators
let result = a.clone() + b.clone();                    // Tensor addition
let result = a.clone() * b.clone();                    // Element-wise multiplication
let result = a.clone() - b.clone();                    // Tensor subtraction
let result = a.clone() / b.clone();                    // Element-wise division

// Scalar operations
let result = a.clone() + 5.0;                          // Tensor + scalar
let result = 5.0 + a.clone();                          // Scalar + tensor
let result = a.clone() * 3.0;                          // Tensor * scalar
let result = 3.0 * a.clone();                          // Scalar * tensor

// Compound expressions
let result = (a.clone() + b.clone()) * 2.0 - 1.0;      // Complex mathematical expressions

// Assignment operators
let mut c = a.clone();
c += b.clone();                                        // In-place addition
c *= 2.0;                                              // In-place scalar multiplication

// Negation
let result = -a;                                       // Negate all elements

Thread Safety

This type is Send + Sync and can be safely shared between threads. All operations are thread-safe through atomic ID generation and thread-local gradtrack storage.

Implementations

impl Tensor

pub fn capacity_elems(&self) -> usize

Returns the allocated capacity in elements, which may be padded beyond logical size

pub fn new(shape_dims: Vec<usize>) -> Self

Creates a new tensor with the specified shape and optimized memory layout

Allocates memory with size-dependent alignment for optimal performance:

• Small tensors (≤8 elements): 16-byte SSE alignment
• Medium tensors (8-1024 elements): 32-byte AVX2 alignment
• Large tensors (>1024 elements): 64-byte cache-line alignment

Arguments

• shape_dims - Vector of dimension sizes defining the tensor shape

Returns

A new tensor with uninitialized data. The data must be initialized before use to avoid undefined behavior.

Performance

• Memory Allocation: Single allocation with optimized alignment
• SIMD Ready: Properly aligned for vectorized operations
• Cache Friendly: Optimized for CPU cache hierarchies
• Thread Safe: Atomic ID generation for gradtrack tracking

Safety

The returned tensor contains uninitialized memory. You must initialize the data before performing any operations that read from it.

Examples

use train_station::Tensor;

// Create tensors of different sizes
let small_tensor = Tensor::new(vec![2, 3]);      // 16-byte alignment
let medium_tensor = Tensor::new(vec![32, 32]);   // 32-byte alignment
let large_tensor = Tensor::new(vec![1000, 1000]); // 64-byte alignment

// Initialize data before use
let mut tensor = Tensor::new(vec![2, 3]);
tensor.fill(0.0); // Initialize with zeros

Examples found in repository

examples/getting_started/tensor_basics.rs (line 70)

fn demonstrate_tensor_creation() {
    println!("--- Tensor Creation ---");

    // Create tensors with different initializations
    let zeros = Tensor::zeros(vec![2, 3]);
    println!(
        "Zeros tensor: shape {:?}, data: {:?}",
        zeros.shape().dims(),
        zeros.data()
    );

    let ones = Tensor::ones(vec![3, 2]);
    println!(
        "Ones tensor: shape {:?}, data: {:?}",
        ones.shape().dims(),
        ones.data()
    );

    // Create tensor from slice
    let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let from_slice = Tensor::from_slice(&data, vec![2, 3]).unwrap();
    println!(
        "From slice: shape {:?}, data: {:?}",
        from_slice.shape().dims(),
        from_slice.data()
    );

    // Create tensor with specific value
    let mut filled = Tensor::new(vec![2, 2]);
    {
        let data = filled.data_mut();
        for value in data.iter_mut() {
            *value = 42.0;
        }
    }
    println!("Filled with 42: {:?}", filled.data());

    // Create tensor with random data
    let random = Tensor::randn(vec![2, 2], Some(42));
    println!(
        "Random tensor: shape {:?}, data: {:?}",
        random.shape().dims(),
        random.data()
    );
}

pub fn shape(&self) -> &Shape

Returns the shape and dimensional information of the tensor

Provides access to the tensor’s dimensions, size, strides, and memory layout information. This is used for shape validation, memory access calculations, and optimization decisions.

Returns

Reference to the tensor’s shape information containing dimensions, size, strides, and memory layout type.

Performance

• Time Complexity: O(1) - direct field access
• Memory: No allocation - returns reference to existing data

Examples

use train_station::Tensor;

let tensor = Tensor::new(vec![2, 3, 4]);
let shape = tensor.shape();
assert_eq!(shape.dims(), vec![2, 3, 4]);
assert_eq!(shape.size(), 24);
assert_eq!(shape.rank(), 3);

Examples found in repository

examples/neural_networks/basic_decoder.rs (line 78)

    fn triple(t: &Tensor) -> (usize, usize, usize) {
        let d = t.shape().dims();
        (d[0], d[1], d[2])
    }
}

#[allow(unused)]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Decoder Example ===");

    let batch = 2usize;
    let src = 7usize;
    let tgt = 5usize;
    let embed = 32usize;
    let heads = 4usize;

    let memory = Tensor::randn(vec![batch, src, embed], Some(21));
    let tgt_in = Tensor::randn(vec![batch, tgt, embed], Some(22));

    let mut dec = DecoderBlock::new(embed, heads, Some(456));
    let out = dec.forward(&tgt_in, &memory, None, None);
    println!("Output shape: {:?}", out.shape().dims());

    let mut opt = Adam::with_learning_rate(0.01);
    let mut params = dec.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/neural_networks/basic_encoder.rs (line 67)

    fn triple(t: &Tensor) -> (usize, usize, usize) {
        let d = t.shape().dims();
        (d[0], d[1], d[2])
    }
}

#[allow(unused)]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Encoder Example ===");

    let batch = 2usize;
    let seq = 6usize;
    let embed = 32usize;
    let heads = 4usize;

    let input = Tensor::randn(vec![batch, seq, embed], Some(11));
    let mut enc = EncoderBlock::new(embed, heads, Some(123));

    // Example: no mask (set Some(mask) to use masking)
    let out = enc.forward(&input, None);
    println!("Output shape: {:?}", out.shape().dims());

    // Verify gradients/optimization
    let mut opt = Adam::with_learning_rate(0.01);
    let mut params = enc.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/RL_training/ppo_continuous.rs (line 112)

    fn forward(&self, state: &Tensor) -> (Tensor, Tensor) {
        // Returns (mean [B, A], log_std [A])
        let mean = self.net.forward(state);
        (
            mean,
            self.log_std
                .view(vec![1, self.log_std.shape().dims()[0] as i32]),
        )
    }

examples/RL_training/ppo_discrete.rs (line 296)

fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

examples/RL_training/td3.rs (line 196)

    fn forward(&self, state: &Tensor, action: &Tensor) -> Tensor {
        // Concatenate along feature dim (dim=1) for batched inputs
        // IMPORTANT: use views to preserve gradient graph; cloning would detach autograd
        let s_view = state.view(state.shape().dims().iter().map(|&d| d as i32).collect());
        let a_view = action.view(action.shape().dims().iter().map(|&d| d as i32).collect());
        let sa = Tensor::cat(&[s_view, a_view], 1);
        self.net.forward(&sa, None)
    }

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 162)

fn demonstrate_layer_creation() {
    println!("--- Layer Creation ---");

    let layer = LinearLayer::new(3, 2, Some(42));

    println!("Created linear layer:");
    println!("  Input size: {}", layer.input_size);
    println!("  Output size: {}", layer.output_size);
    println!("  Parameter count: {}", layer.parameter_count());
    println!("  Weight shape: {:?}", layer.weight.shape().dims());
    println!("  Bias shape: {:?}", layer.bias.shape().dims());
    println!("  Weight requires grad: {}", layer.weight.requires_grad());
    println!("  Bias requires grad: {}", layer.bias.requires_grad());
}

/// Demonstrate forward pass with gradient tracking
fn demonstrate_forward_pass() {
    println!("\n--- Forward Pass (with gradients) ---");

    let layer = LinearLayer::new(3, 2, Some(43));

    // Single input
    let input = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3]).unwrap();
    let output = layer.forward(&input);

    println!("Single input:");
    println!("  Input: {:?}", input.data());
    println!("  Output: {:?}", output.data());
    println!("  Output requires grad: {}", output.requires_grad());

    // Batch input
    let batch_input = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
    let batch_output = layer.forward(&batch_input);

    println!("Batch input:");
    println!("  Input shape: {:?}", batch_input.shape().dims());
    println!("  Output shape: {:?}", batch_output.shape().dims());
    println!("  Output requires grad: {}", batch_output.requires_grad());
}

/// Demonstrate forward pass without gradient tracking
fn demonstrate_forward_pass_no_grad() {
    println!("\n--- Forward Pass (no gradients) ---");

    let layer = LinearLayer::new(3, 2, Some(44));

    // Single input
    let input = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3]).unwrap();
    let output = layer.forward_no_grad(&input);

    println!("Single input (no grad):");
    println!("  Input: {:?}", input.data());
    println!("  Output: {:?}", output.data());
    println!("  Output requires grad: {}", output.requires_grad());

    // Compare with grad version
    let output_with_grad = layer.forward(&input);
    println!("Comparison:");
    println!(
        "  Same values: {}",
        output.data() == output_with_grad.data()
    );
    println!("  No grad requires grad: {}", output.requires_grad());
    println!(
        "  With grad requires grad: {}",
        output_with_grad.requires_grad()
    );
}

/// Demonstrate complete training loop
fn demonstrate_training_loop() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Loop ---");

    // Create layer and training data
    let mut layer = LinearLayer::new(2, 1, Some(45));

    // Simple regression task: y = 2*x1 + 3*x2 + 1
    let x_data = Tensor::from_slice(
        &[
            1.0, 1.0, // x1=1, x2=1 -> y=6
            2.0, 1.0, // x1=2, x2=1 -> y=8
            1.0, 2.0, // x1=1, x2=2 -> y=9
            2.0, 2.0, // x1=2, x2=2 -> y=11
        ],
        vec![4, 2],
    )
    .unwrap();

    let y_true = Tensor::from_slice(&[6.0, 8.0, 9.0, 11.0], vec![4, 1]).unwrap();

    println!("Training data:");
    println!("  X shape: {:?}", x_data.shape().dims());
    println!("  Y shape: {:?}", y_true.shape().dims());
    println!("  Target function: y = 2*x1 + 3*x2 + 1");

    // Create optimizer
    let config = AdamConfig {
        learning_rate: 0.01,
        beta1: 0.9,
        beta2: 0.999,
        eps: 1e-8,
        weight_decay: 0.0,
        amsgrad: false,
    };

    let mut optimizer = Adam::with_config(config);
    let params = layer.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    println!("Optimizer setup complete. Starting training...");

    // Training loop
    let num_epochs = 100;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = layer.forward(&x_data);

        // Compute loss: MSE
        let diff = y_pred.sub_tensor(&y_true);
        let mut loss = diff.pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        let mut params = layer.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        losses.push(loss.value());

        // Print progress
        if epoch % 20 == 0 || epoch == num_epochs - 1 {
            println!("Epoch {:3}: Loss = {:.6}", epoch, loss.value());
        }
    }

    // Evaluate final model
    let final_predictions = layer.forward_no_grad(&x_data);

    println!("\nFinal model evaluation:");
    println!("  Learned weights: {:?}", layer.weight.data());
    println!("  Learned bias: {:?}", layer.bias.data());
    println!("  Target weights: [2.0, 3.0]");
    println!("  Target bias: [1.0]");

    println!("  Predictions vs True:");
    for i in 0..4 {
        let pred = final_predictions.data()[i];
        let true_val = y_true.data()[i];
        println!(
            "    Sample {}: pred={:.3}, true={:.1}, error={:.3}",
            i + 1,
            pred,
            true_val,
            (pred - true_val).abs()
        );
    }

    // Training analysis
    let initial_loss = losses[0];
    let final_loss = losses[losses.len() - 1];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);

    Ok(())
}

/// Demonstrate single vs batch inference
fn demonstrate_single_vs_batch_inference() {
    println!("\n--- Single vs Batch Inference ---");

    let layer = LinearLayer::new(4, 3, Some(46));

    // Single inference
    println!("Single inference:");
    let single_input = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![1, 4]).unwrap();
    let single_output = layer.forward_no_grad(&single_input);
    println!("  Input shape: {:?}", single_input.shape().dims());
    println!("  Output shape: {:?}", single_output.shape().dims());
    println!("  Output: {:?}", single_output.data());

    // Batch inference
    println!("Batch inference:");
    let batch_input = Tensor::from_slice(
        &[
            1.0, 2.0, 3.0, 4.0, // Sample 1
            5.0, 6.0, 7.0, 8.0, // Sample 2
            9.0, 10.0, 11.0, 12.0, // Sample 3
        ],
        vec![3, 4],
    )
    .unwrap();
    let batch_output = layer.forward_no_grad(&batch_input);
    println!("  Input shape: {:?}", batch_input.shape().dims());
    println!("  Output shape: {:?}", batch_output.shape().dims());

    // Verify batch consistency - first sample should match single inference
    let _first_batch_sample = batch_output.view(vec![3, 3]); // Reshape to access first sample
    let first_sample_data = &batch_output.data()[0..3]; // First 3 elements
    let single_sample_data = single_output.data();

    println!("Consistency check:");
    println!("  Single output: {:?}", single_sample_data);
    println!("  First batch sample: {:?}", first_sample_data);
    println!(
        "  Match: {}",
        single_sample_data
            .iter()
            .zip(first_sample_data.iter())
            .all(|(a, b)| (a - b).abs() < 1e-6)
    );
}

Additional examples can be found in:

• examples/iterators/element_iteration.rs
• examples/getting_started/tensor_basics.rs
• examples/getting_started/tensor_operators.rs
• examples/getting_started/optimizer_basics.rs
• examples/neural_networks/feedforward_network.rs
• examples/neural_networks/basic_transformer.rs
• examples/getting_started/serialization_basics.rs
• examples/neural_networks/multi_head_attention.rs
• examples/iterators/advanced_patterns.rs

pub fn size(&self) -> usize

Returns the total number of elements in the tensor

Provides the total count of elements across all dimensions. This is used for memory allocation, iteration bounds, and performance optimization.

Returns

Total number of elements as usize

Performance

• Time Complexity: O(1) - direct field access
• Memory: No allocation - returns stored value

Examples

use train_station::Tensor;

let tensor = Tensor::new(vec![2, 3, 4]);
assert_eq!(tensor.size(), 24); // 2 * 3 * 4

let scalar = Tensor::new(vec![1]);
assert_eq!(scalar.size(), 1);

let empty = Tensor::new(vec![0]);
assert_eq!(empty.size(), 0);

Examples found in repository

examples/getting_started/tensor_basics.rs (line 186)

fn demonstrate_utility_functions() {
    println!("\n--- Utility Functions ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

    // Basic properties
    println!("Shape: {:?}", tensor.shape().dims());
    println!("Size: {}", tensor.size());
    println!("Is contiguous: {}", tensor.is_contiguous());
    println!("Device: {:?}", tensor.device());

    // Mathematical operations
    let sum = tensor.sum();
    println!("Sum: {}", sum.value());

    let mean = tensor.mean();
    println!("Mean: {}", mean.value());

    let norm = tensor.norm();
    println!("Norm: {}", norm.value());

    // Device placement
    let cpu_tensor = Tensor::zeros_on_device(vec![3, 3], train_station::Device::cpu());
    println!(
        "CPU tensor: shape {:?}, device: {:?}",
        cpu_tensor.shape().dims(),
        cpu_tensor.device()
    );
}

examples/iterators/performance_optimization.rs (line 209)

fn demonstrate_memory_optimization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Memory Optimization ---");

    // Create a large tensor for memory testing
    let size = 10000;
    let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
    let tensor = Tensor::from_slice(&data, vec![size])?;

    println!("Processing tensor of size: {}", size);

    // Pattern 1: Streaming processing with iterator chunks (process in blocks, collect with shape)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();
    let flattened = tensor.view(vec![size as i32]);
    let _streamed_result: Tensor = flattened
        .chunks(chunk_size)
        .map(|c| c.pow_scalar(2.0).sqrt())
        .collect_shape(vec![size]);
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter_elements()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect_shape(vec![size]);
    let full_time = start.elapsed();

    println!("  Streaming time: {:?}", streamed_time);
    println!("  Full processing time: {:?}", full_time);
    println!(
        "  Memory efficiency ratio: {:.2}x",
        full_time.as_nanos() as f64 / streamed_time.as_nanos() as f64
    );

    // Pattern 3: Lazy evaluation with take
    println!("\nPattern 2: Lazy Evaluation");
    let start = Instant::now();
    let lazy_result: Tensor = tensor
        .iter_elements()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect_shape(vec![1000]);
    let lazy_time = start.elapsed();

    println!("  Lazy processing (1000 elements): {:?}", lazy_time);
    println!("  Lazy result size: {}", lazy_result.size());

    // Pattern 4: Memory-efficient filtering
    println!("\nPattern 3: Memory-Efficient Filtering");
    let start = Instant::now();
    let filtered_result: Tensor = tensor
        .iter_elements()
        .filter(|elem| elem.value() > size as f32 / 2.0) // Keep only large values
        .map(|elem| elem.mul_scalar(2.0))
        .collect();
    let filtered_time = start.elapsed();

    println!("  Filtered processing: {:?}", filtered_time);
    println!(
        "  Filtered result size: {} (reduced from {})",
        filtered_result.size(),
        size
    );

    Ok(())
}

/// Demonstrate large-scale processing techniques
///
/// Shows how to efficiently process very large datasets using
/// iterator patterns and optimization strategies.
fn demonstrate_large_scale_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Large-Scale Processing ---");

    // Simulate large dataset processing
    let sizes = vec![10000, 50000, 100000];

    for size in sizes {
        println!("\nProcessing dataset of size: {}", size);

        // Generate large dataset
        let data: Vec<f32> = (0..size)
            .map(|i| {
                let x = i as f32 / size as f32;
                x * x + 0.1 * (i % 10) as f32 // Quadratic with noise
            })
            .collect();

        let tensor = Tensor::from_slice(&data, vec![size])?;

        // Technique 1: Batch processing
        let batch_size = 1000;
        let start = Instant::now();

        let mut batch_results = Vec::new();
        for batch_start in (0..size).step_by(batch_size) {
            let batch_end = (batch_start + batch_size).min(size);
            let batch: Tensor = tensor
                .iter_range(batch_start, batch_end)
                .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
                .collect();
            batch_results.push(batch);
        }
        let batch_time = start.elapsed();

        // Technique 2: Parallel-like processing with stride
        let start = Instant::now();
        let stride = 4;
        let strided_result: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % stride == 0)
            .map(|(_, elem)| elem.pow_scalar(2.0).add_scalar(1.0))
            .collect();
        let strided_time = start.elapsed();

        // Technique 3: Hierarchical processing
        let start = Instant::now();
        let coarse: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % 10 == 0) // Every 10th element
            .map(|(_, elem)| elem.pow_scalar(2.0).add_scalar(1.0))
            .collect();
        let fine: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % 10 != 0) // Rest of elements
            .map(|(_, elem)| elem.pow_scalar(1.5).add_scalar(0.5))
            .collect();
        let hierarchical_time = start.elapsed();

        // Report performance
        println!("  Batch processing: {:?}", batch_time);
        println!("  Strided processing: {:?}", strided_time);
        println!("  Hierarchical processing: {:?}", hierarchical_time);

        // Memory usage analysis
        let total_batches = size.div_ceil(batch_size);
        println!("  Batch count: {}", total_batches);
        println!("  Strided result size: {}", strided_result.size());
        println!(
            "  Hierarchical: coarse={}, fine={}",
            coarse.size(),
            fine.size()
        );
    }

    Ok(())
}

/// Demonstrate advanced optimization techniques
///
/// Shows sophisticated optimization strategies and techniques
/// for maximizing performance in tensor iterator operations.
fn demonstrate_optimization_techniques() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Optimization Techniques ---");

    let size = 50000;
    let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
    let tensor = Tensor::from_slice(&data, vec![size])?;

    println!("Optimizing processing for size: {}", size);

    // Technique 1: Operation fusion
    println!("\nTechnique 1: Operation Fusion");
    let start = Instant::now();
    let fused_result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Fuse multiple operations into single chain
            elem.mul_scalar(2.0).add_scalar(1.0).pow_scalar(2.0).sqrt()
        })
        .collect();
    let fused_time = start.elapsed();

    // Technique 2: Conditional optimization
    println!("\nTechnique 2: Conditional Optimization");
    let start = Instant::now();
    let conditional_result: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val < size as f32 / 2.0 {
                elem.mul_scalar(2.0) // Simple operation for small values
            } else {
                elem.pow_scalar(2.0).sqrt() // Complex operation for large values
            }
        })
        .collect();
    let conditional_time = start.elapsed();

    // Technique 3: Cache-friendly processing
    println!("\nTechnique 3: Cache-Friendly Processing");
    let start = Instant::now();
    let cache_friendly_result: Tensor = tensor
        .iter()
        .take(1000) // Process in cache-friendly chunks
        .map(|elem| elem.mul_scalar(2.0))
        .collect();
    let cache_friendly_time = start.elapsed();

    // Technique 4: Memory pooling simulation
    println!("\nTechnique 4: Memory Pooling Simulation");
    let start = Instant::now();
    let pooled_result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 100 == 0) // Process every 100th element
        .map(|(_, elem)| elem.pow_scalar(2.0))
        .collect();
    let pooled_time = start.elapsed();

    // Report optimization results
    println!("  Fused operations: {:?}", fused_time);
    println!("  Conditional optimization: {:?}", conditional_time);
    println!("  Cache-friendly processing: {:?}", cache_friendly_time);
    println!("  Memory pooling simulation: {:?}", pooled_time);

    // Performance analysis
    let fastest = fused_time
        .min(conditional_time)
        .min(cache_friendly_time)
        .min(pooled_time);
    println!("  Fastest technique: {:?}", fastest);

    // Memory efficiency analysis
    println!("  Fused result size: {}", fused_result.size());
    println!("  Conditional result size: {}", conditional_result.size());
    println!(
        "  Cache-friendly result size: {}",
        cache_friendly_result.size()
    );
    println!("  Pooled result size: {}", pooled_result.size());

    // Technique 5: Gradient optimization
    println!("\nTechnique 5: Gradient Optimization");
    let grad_tensor = tensor.with_requires_grad();
    let start = Instant::now();

    let grad_result: Tensor = grad_tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
        .collect();

    let mut loss = grad_result.sum();
    loss.backward(None);
    let grad_time = start.elapsed();

    println!("  Gradient computation: {:?}", grad_time);
    println!(
        "  Gradient tracking enabled: {}",
        grad_result.requires_grad()
    );

    Ok(())
}

examples/iterators/advanced_patterns.rs (line 128)

fn demonstrate_data_pipeline() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Data Processing Pipeline ---");

    // Simulate raw sensor data with noise
    let raw_data: Vec<f32> = (0..20)
        .map(|i| {
            let base = i as f32 * 0.5;
            let noise = (i % 3) as f32 * 0.1;
            base + noise
        })
        .collect();

    let tensor = Tensor::from_slice(&raw_data, vec![20])?;
    println!("Raw sensor data: {:?}", tensor.data());

    // Multi-stage processing pipeline
    println!("\nProcessing pipeline:");
    println!("1. Normalize data (z-score)");
    println!("2. Apply smoothing filter");
    println!("3. Detect outliers");
    println!("4. Apply feature scaling");

    // Stage 1: Normalization
    let mean = tensor.mean().value();
    let std = tensor.std().value();
    let normalized: Tensor = tensor
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();
    println!(
        "  Normalized (mean={:.3}, std={:.3}): {:?}",
        mean,
        std,
        normalized.data()
    );

    // Stage 2: Smoothing (simple moving average)
    let smoothed: Tensor = normalized
        .iter()
        .enumerate()
        .map(|(i, elem)| {
            if i == 0 || i == normalized.size() - 1 {
                elem.clone()
            } else {
                // Simple 3-point average
                let prev = normalized.element_view(i - 1);
                let next = normalized.element_view(i + 1);
                elem.add_tensor(&prev).add_tensor(&next).div_scalar(3.0)
            }
        })
        .collect();
    println!("  Smoothed: {:?}", smoothed.data());

    // Stage 3: Outlier detection and removal
    let outlier_threshold = 2.0;
    let cleaned: Tensor = smoothed
        .iter()
        .filter(|elem| elem.value().abs() < outlier_threshold)
        .collect();
    println!(
        "  Outliers removed (threshold={}): {:?}",
        outlier_threshold,
        cleaned.data()
    );

    // Stage 4: Feature scaling to [0, 1] range
    let min_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);
    let scaled: Tensor = cleaned
        .iter()
        .map(|elem| elem.sub_scalar(min_val).div_scalar(max_val - min_val))
        .collect();
    println!("  Scaled to [0,1]: {:?}", scaled.data());

    Ok(())
}

/// Demonstrate conditional processing patterns
///
/// Shows how to implement dynamic filtering and transformation
/// based on data characteristics and conditions.
fn demonstrate_conditional_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Conditional Processing ---");

    // Create data with mixed characteristics
    let data = vec![1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0, 9.0, -10.0];
    let tensor = Tensor::from_slice(&data, vec![10])?;
    println!("Input data: {:?}", tensor.data());

    // Conditional transformation based on sign
    println!("\nConditional transformation (positive/negative handling):");
    let processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val > 0.0 {
                elem.pow_scalar(2.0) // Square positive values
            } else {
                elem.mul_scalar(-1.0).sqrt() // Square root of absolute negative values
            }
        })
        .collect();
    println!("  Processed: {:?}", processed.data());

    // Adaptive filtering based on local statistics
    println!("\nAdaptive filtering (remove values > 2 std from local mean):");
    let window_size = 3;
    let adaptive_filtered: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, elem)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(tensor.size());

            // Calculate local mean and std
            let local_values: Vec<f32> = (start..end)
                .map(|j| tensor.element_view(j).value())
                .collect();

            let local_mean = local_values.iter().sum::<f32>() / local_values.len() as f32;
            let local_variance = local_values
                .iter()
                .map(|v| (v - local_mean).powi(2))
                .sum::<f32>()
                / local_values.len() as f32;
            let local_std = local_variance.sqrt();

            let threshold = local_mean + 2.0 * local_std;
            elem.value() <= threshold
        })
        .map(|(_, elem)| elem)
        .collect();
    println!("  Adaptive filtered: {:?}", adaptive_filtered.data());

    // Multi-condition processing
    println!("\nMulti-condition processing:");
    let multi_processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            match () {
                _ if val > 5.0 => elem.mul_scalar(2.0), // Double large values
                _ if val < -5.0 => elem.div_scalar(2.0), // Halve small values
                _ if val.abs() < 2.0 => elem.add_scalar(1.0), // Add 1 to small values
                _ => elem.clone(),                      // Keep others unchanged
            }
        })
        .collect();
    println!("  Multi-condition: {:?}", multi_processed.data());

    Ok(())
}

/// Demonstrate batch processing operations
///
/// Shows efficient processing of large datasets using iterator
/// patterns and batch operations for performance optimization.
fn demonstrate_batch_operations() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Batch Operations ---");

    // Create a larger dataset for batch processing
    let size = 100;
    let data: Vec<f32> = (0..size)
        .map(|i| {
            let x = i as f32 / size as f32;
            x * x + 0.1 * (i % 7) as f32 // Quadratic with some noise
        })
        .collect();

    let tensor = Tensor::from_slice(&data, vec![size])?;
    println!("Dataset size: {}", tensor.size());

    // Batch processing with windowing (iterator views)
    println!("\nBatch processing with sliding windows:");
    let batch_size = 10;
    let batches: Vec<Tensor> = tensor
        .iter()
        .collect::<Vec<_>>()
        .chunks(batch_size)
        .map(|chunk| {
            // Process each batch independently
            chunk
                .iter()
                .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
                .collect()
        })
        .collect();

    println!(
        "  Processed {} batches of size {}",
        batches.len(),
        batch_size
    );
    for (i, batch) in batches.iter().enumerate() {
        println!(
            "    Batch {}: mean={:.3}, std={:.3}",
            i,
            batch.mean().value(),
            batch.std().value()
        );
    }

    // Parallel-like processing with stride
    println!("\nStrided processing (every nth element):");
    let stride = 5;
    let strided: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % stride == 0)
        .map(|(_, elem)| elem)
        .collect();
    println!("  Strided (every {}th): {:?}", stride, strided.data());

    // Hierarchical processing
    println!("\nHierarchical processing (coarse to fine):");
    let coarse: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 == 0) // Take every 4th element
        .map(|(_, elem)| elem)
        .collect();

    let fine: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 != 0) // Take the rest
        .map(|(_, elem)| elem)
        .collect();

    println!("  Coarse (every 4th): {:?}", coarse.data());
    println!("  Fine (rest): {:?}", fine.data());

    // Combine coarse and fine with different processing
    let combined: Tensor = coarse
        .iter()
        .map(|elem| elem.mul_scalar(2.0)) // Scale coarse
        .chain(fine.iter().map(|elem| elem.div_scalar(2.0))) // Scale fine
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

/// Demonstrate real-world processing scenarios
///
/// Shows practical applications of iterator patterns for
/// common data processing tasks in machine learning and analytics.
fn demonstrate_real_world_scenarios() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Real-world Scenarios ---");

    // Scenario 1: Time series analysis
    println!("\nScenario 1: Time Series Analysis");
    let time_series: Vec<f32> = (0..24)
        .map(|hour| {
            let base = 20.0 + 10.0 * (hour as f32 * std::f32::consts::PI / 12.0).sin();
            base + (hour % 3) as f32 * 2.0 // Add some noise
        })
        .collect();

    let series = Tensor::from_slice(&time_series, vec![24])?;
    println!("  Time series (24 hours): {:?}", series.data());

    // Calculate moving average with view-based iteration
    let window_size = 3;
    let moving_avg: Tensor = series
        .iter()
        .enumerate()
        .map(|(i, _)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(series.size());
            let window = series.iter_range(start, end);
            window.fold(0.0, |acc, elem| acc + elem.value()) / (end - start) as f32
        })
        .map(|val| Tensor::from_slice(&[val], vec![1]).unwrap())
        .collect();
    println!(
        "  Moving average (window={}): {:?}",
        window_size,
        moving_avg.data()
    );

    // Inference pipeline with NoGrad + streaming
    println!("\nInference pipeline (NoGrad + streaming)");
    let features = Tensor::from_slice(
        &(0..48).map(|i| i as f32 * 0.125).collect::<Vec<_>>(),
        vec![6, 8],
    )?;
    let fast = with_no_grad(|| {
        // Stream values directly, apply light affine, and collect back to same shape
        features
            .data()
            .iter()
            .copied()
            .map(|x| 0.75 * x + 0.1)
            .collect_shape(vec![6, 8])
    });
    println!(
        "  NoGrad streamed transform shape: {:?}",
        fast.shape().dims()
    );

    // Row-wise iteration with shape-preserving collection (GradTrack-friendly)
    let per_row: Tensor = features
        .iter()
        .map(|row| row.mul_scalar(0.5).add_scalar(2.0))
        .collect_shape(vec![6, 8]);
    println!("  Row-wise mapped shape: {:?}", per_row.shape().dims());

    // Scenario 2: Feature engineering
    println!("\nScenario 2: Feature Engineering");
    let features = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;
    println!("  Original features: {:?}", features.data());

    // Create polynomial features
    let poly_features: Tensor = features
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // x^1
                elem.pow_scalar(2.0), // x^2
                elem.pow_scalar(3.0), // x^3
            ]
        })
        .collect();
    println!(
        "  Polynomial features (x, x^2, x^3): {:?}",
        poly_features.data()
    );

    // Scenario 3: Data augmentation
    println!("\nScenario 3: Data Augmentation");
    let original = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?;
    println!("  Original data: {:?}", original.data());

    // Augment with noise and scaling
    let augmented: Tensor = original
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // Original
                elem.add_scalar(0.1), // Add noise
                elem.sub_scalar(0.1), // Subtract noise
                elem.mul_scalar(1.1), // Scale up
                elem.mul_scalar(0.9), // Scale down
            ]
        })
        .collect();
    println!("  Augmented data: {:?}", augmented.data());

    // Scenario 4: Statistical analysis
    println!("\nScenario 4: Statistical Analysis");
    let sample_data = Tensor::from_slice(&[1.1, 2.3, 1.8, 2.1, 1.9, 2.0, 1.7, 2.2], vec![8])?;
    println!("  Sample data: {:?}", sample_data.data());

    // Calculate various statistics
    let mean = sample_data.mean().value();
    let std = sample_data.std().value();
    let min = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);

    // Z-score normalization
    let z_scores: Tensor = sample_data
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();

    println!(
        "  Statistics: mean={:.3}, std={:.3}, min={:.3}, max={:.3}",
        mean, std, min, max
    );
    println!("  Z-scores: {:?}", z_scores.data());

    Ok(())
}

pub fn device(&self) -> Device

Returns the device where this tensor is located

Provides the physical location of the tensor data (CPU/GPU). This determines which operations can be performed on the tensor and where computations will be executed.

Returns

Device enum indicating the tensor’s physical location

Performance

• Time Complexity: O(1) - direct field access
• Memory: No allocation - returns stored value

Examples

use train_station::Tensor;

let tensor = Tensor::new(vec![2, 3]);
assert!(tensor.device().is_cpu());
assert!(!tensor.device().is_cuda());

Examples found in repository

examples/getting_started/tensor_basics.rs (line 188)

fn demonstrate_utility_functions() {
    println!("\n--- Utility Functions ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

    // Basic properties
    println!("Shape: {:?}", tensor.shape().dims());
    println!("Size: {}", tensor.size());
    println!("Is contiguous: {}", tensor.is_contiguous());
    println!("Device: {:?}", tensor.device());

    // Mathematical operations
    let sum = tensor.sum();
    println!("Sum: {}", sum.value());

    let mean = tensor.mean();
    println!("Mean: {}", mean.value());

    let norm = tensor.norm();
    println!("Norm: {}", norm.value());

    // Device placement
    let cpu_tensor = Tensor::zeros_on_device(vec![3, 3], train_station::Device::cpu());
    println!(
        "CPU tensor: shape {:?}, device: {:?}",
        cpu_tensor.shape().dims(),
        cpu_tensor.device()
    );
}

pub fn new_on_device(shape_dims: Vec<usize>, device: Device) -> Self

Creates a new tensor with the specified shape on a specific device

Allocates memory on the specified device with the same optimized alignment strategy as new(). Currently supports CPU device with future CUDA support.

Arguments

• shape_dims - Vector of dimension sizes defining the tensor shape
• device - The device where the tensor should be allocated

Returns

A new tensor with uninitialized data on the specified device

Performance

• Memory Allocation: Device-specific allocation with optimized alignment
• SIMD Ready: Properly aligned for vectorized operations on target device
• Thread Safe: Atomic ID generation for gradtrack tracking

Panics

Panics if the specified device is not supported (e.g., CUDA without feature flag)

Examples

use train_station::Tensor;

let tensor = Tensor::new_on_device(vec![2, 3], train_station::Device::cpu());
assert!(tensor.device().is_cpu());
assert_eq!(tensor.size(), 6);

Arguments

• shape_dims - Vector of dimension sizes defining the tensor shape
• device - Device where the tensor should be allocated

Returns

A new tensor with uninitialized data on the specified device

Panics

Panics if the device is not supported (currently only CPU is supported)

Performance

• Memory Allocation: Single allocation with optimized alignment
• SIMD Ready: Properly aligned for vectorized operations
• Cache Friendly: Optimized for CPU cache hierarchies
• Thread Safe: Atomic ID generation for gradtrack tracking

Examples

use train_station::{Tensor, Device};

// Create tensor on CPU device
let tensor = Tensor::new_on_device(vec![2, 3], Device::cpu());
assert_eq!(tensor.device(), Device::cpu());
assert_eq!(tensor.size(), 6);

pub fn with_requires_grad(self) -> Self

Enable gradient computation for this tensor

Builder method that enables automatic gradient tracking for this tensor. When enabled, all operations involving this tensor will be recorded in the computation graph for gradient computation during backward pass.

Returns

self with gradient tracking enabled

Performance

• Time Complexity: O(1) - simple field assignment
• Memory: No additional allocation
• Overhead: Minimal gradtrack tracking overhead when gradients computed

Examples

use train_station::Tensor;

let tensor = Tensor::ones(vec![2, 3]).with_requires_grad();
assert!(tensor.requires_grad());

Examples found in repository

examples/RL_training/ppo_continuous.rs (line 103)

    fn new(state_dim: usize, action_dim: usize, seed: Option<u64>) -> Self {
        let net = Mlp::new(&[state_dim, 64, 64, action_dim], seed);
        let log_std = Tensor::from_slice(&vec![0.0; action_dim], vec![action_dim])
            .unwrap()
            .with_requires_grad();
        Self { net, log_std }
    }

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 59)

    pub fn new(input_size: usize, output_size: usize, seed: Option<u64>) -> Self {
        // Xavier/Glorot initialization: scale by sqrt(1/input_size)
        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,
        }
    }

    /// Forward pass: output = input @ weight + bias
    pub fn forward(&self, input: &Tensor) -> Tensor {
        // Matrix multiplication: [batch_size, input_size] @ [input_size, output_size] = [batch_size, output_size]
        let output = input.matmul(&self.weight);
        // Add bias: [batch_size, output_size] + [output_size] = [batch_size, output_size]
        output.add_tensor(&self.bias)
    }

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

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

    /// Save layer parameters to JSON
    #[allow(unused)]
    pub fn save_json(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
        // Create directory if it doesn't exist
        if let Some(parent) = std::path::Path::new(path).parent() {
            fs::create_dir_all(parent)?;
        }

        let weight_path = format!("{}_weight.json", path);
        let bias_path = format!("{}_bias.json", path);

        self.weight.save_json(&weight_path)?;
        self.bias.save_json(&bias_path)?;

        println!("Saved linear layer to {} (weight and bias)", path);
        Ok(())
    }

    /// Load layer parameters from JSON
    #[allow(unused)]
    pub fn load_json(
        path: &str,
        input_size: usize,
        output_size: usize,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let weight_path = format!("{}_weight.json", path);
        let bias_path = format!("{}_bias.json", path);

        let weight = Tensor::load_json(&weight_path)?.with_requires_grad();
        let bias = Tensor::load_json(&bias_path)?.with_requires_grad();

        Ok(Self {
            weight,
            bias,
            input_size,
            output_size,
        })
    }

examples/iterators/element_iteration.rs (line 179)

fn demonstrate_gradient_tracking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Gradient Tracking ---");

    // Create a tensor with gradient tracking enabled
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?.with_requires_grad();
    println!("Input tensor (requires_grad): {:?}", tensor.data());

    // Perform element-wise operations through iteration
    let result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Apply a complex transformation: (x^2 + 1) * 2
            elem.pow_scalar(2.0).add_scalar(1.0).mul_scalar(2.0)
        })
        .collect();

    println!("Result tensor: {:?}", result.data());
    println!("Result requires_grad: {}", result.requires_grad());

    // Compute gradients
    let mut loss = result.sum();
    loss.backward(None);

    println!("Loss: {:.6}", loss.value());
    println!("Input gradients: {:?}", tensor.grad().map(|g| g.data()));

    Ok(())
}

examples/neural_networks/feedforward_network.rs (line 214)

    pub fn load_json(
        path: &str,
        config: FeedForwardConfig,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let mut layers = Vec::new();
        let mut current_size = config.input_size;
        let mut layer_idx = 0;

        // Load hidden layers
        for &hidden_size in &config.hidden_sizes {
            let layer_path = format!("{}_layer_{}", path, layer_idx);
            let weight_path = format!("{}_weight.json", layer_path);
            let bias_path = format!("{}_bias.json", layer_path);

            let weight = Tensor::load_json(&weight_path)?.with_requires_grad();
            let bias = Tensor::load_json(&bias_path)?.with_requires_grad();

            layers.push(LinearLayer {
                weight,
                bias,
                input_size: current_size,
                output_size: hidden_size,
            });

            current_size = hidden_size;
            layer_idx += 1;
        }

        // Load output layer
        let layer_path = format!("{}_layer_{}", path, layer_idx);
        let weight_path = format!("{}_weight.json", layer_path);
        let bias_path = format!("{}_bias.json", layer_path);

        let weight = Tensor::load_json(&weight_path)?.with_requires_grad();
        let bias = Tensor::load_json(&bias_path)?.with_requires_grad();

        layers.push(LinearLayer {
            weight,
            bias,
            input_size: current_size,
            output_size: config.output_size,
        });

        Ok(Self { layers, config })
    }

examples/getting_started/optimizer_basics.rs (line 51)

fn demonstrate_basic_optimizer_setup() {
    println!("--- Basic Optimizer Setup ---");

    // Create parameters that require gradients
    let weight = Tensor::randn(vec![3, 2], Some(42)).with_requires_grad();
    let bias = Tensor::zeros(vec![2]).with_requires_grad();

    println!("Created parameters:");
    println!(
        "  Weight: shape {:?}, requires_grad: {}",
        weight.shape().dims(),
        weight.requires_grad()
    );
    println!(
        "  Bias: shape {:?}, requires_grad: {}",
        bias.shape().dims(),
        bias.requires_grad()
    );

    // Create Adam optimizer with default configuration
    let mut optimizer = Adam::new();
    println!(
        "Created Adam optimizer with learning rate: {}",
        optimizer.learning_rate()
    );

    // Add parameters to optimizer
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);
    println!(
        "Added {} parameters to optimizer",
        optimizer.parameter_count()
    );

    // Create optimizer with custom configuration
    let config = AdamConfig {
        learning_rate: 0.01,
        beta1: 0.9,
        beta2: 0.999,
        eps: 1e-8,
        weight_decay: 0.0,
        amsgrad: false,
    };

    let mut custom_optimizer = Adam::with_config(config);
    custom_optimizer.add_parameter(&weight);
    custom_optimizer.add_parameter(&bias);

    println!(
        "Created custom optimizer with learning rate: {}",
        custom_optimizer.learning_rate()
    );

    // Demonstrate parameter linking
    println!("Parameter linking completed successfully");
}

/// Demonstrate simple linear regression training
fn demonstrate_linear_regression() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Linear Regression Training ---");

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(43)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Create simple training data: y = 2*x + 1
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    println!("Training data:");
    println!("  X: {:?}", x_data.data());
    println!("  Y: {:?}", y_true.data());
    println!("  Target: y = 2*x + 1");

    // Training loop
    let num_epochs = 100;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass: y_pred = x * weight + bias
        let y_pred = x_data.matmul(&weight) + &bias;

        // Compute loss: MSE
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        losses.push(loss.value());

        // Print progress every 20 epochs
        if epoch % 20 == 0 || epoch == num_epochs - 1 {
            println!("Epoch {:3}: Loss = {:.6}", epoch, loss.value());
        }
    }

    // Evaluate final model
    let final_predictions = x_data.matmul(&weight) + &bias;
    println!("\nFinal model evaluation:");
    println!("  Learned weight: {:.6}", weight.value());
    println!("  Learned bias: {:.6}", bias.value());
    println!("  Predictions vs True:");

    for i in 0..5 {
        let x1 = x_data.data()[i];
        let pred = final_predictions.data()[i];
        let true_val = y_true.data()[i];
        println!(
            "    x={:.1}: pred={:.3}, true={:.1}, error={:.3}",
            x1,
            pred,
            true_val,
            (pred - true_val).abs()
        );
    }

    Ok(())
}

/// Demonstrate advanced training patterns
fn demonstrate_advanced_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Training Patterns ---");

    // Create a more complex model
    let mut weight = Tensor::randn(vec![1, 2], Some(44)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![2]).with_requires_grad();

    // Create optimizer with different learning rate
    let mut optimizer = Adam::with_learning_rate(0.005);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Create training data: y = 2*x + [1, 3]
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(
        &[3.0, 5.0, 7.0, 9.0, 11.0, 6.0, 8.0, 10.0, 12.0, 14.0],
        vec![5, 2],
    )
    .unwrap();

    println!("Advanced training with monitoring:");
    println!("  Initial learning rate: {}", optimizer.learning_rate());

    // Training loop with monitoring
    let num_epochs = 50;
    let mut losses = Vec::new();
    let mut weight_norms = Vec::new();
    let mut gradient_norms = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Compute gradient norm before optimizer step
        let gradient_norm = weight.grad_owned().unwrap().norm();

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Learning rate scheduling: reduce every 10 epochs
        if epoch > 0 && epoch % 10 == 0 {
            let current_lr = optimizer.learning_rate();
            let new_lr = current_lr * 0.5;
            optimizer.set_learning_rate(new_lr);
            println!(
                "Epoch {:2}: Reduced learning rate from {:.3} to {:.3}",
                epoch, current_lr, new_lr
            );
        }

        // Record metrics
        losses.push(loss.value());
        weight_norms.push(weight.norm().value());
        gradient_norms.push(gradient_norm.value());

        // Print detailed progress
        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, Weight Norm = {:.6}, Gradient Norm = {:.6}",
                epoch,
                loss.value(),
                weight.norm().value(),
                gradient_norm.value()
            );
        }
    }

    println!("Final learning rate: {}", optimizer.learning_rate());

    // Analyze training progression
    let initial_loss = losses[0];
    let final_loss = losses[losses.len() - 1];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final weight norm: {:.6}", weight.norm().value());
    println!("  Final bias: {:?}", bias.data());

    Ok(())
}

/// Demonstrate learning rate scheduling
fn demonstrate_learning_rate_scheduling() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Learning Rate Scheduling ---");

    // Create simple model
    let mut weight = Tensor::randn(vec![1, 1], Some(45)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer with high initial learning rate
    let mut optimizer = Adam::with_learning_rate(0.1);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Simple data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3, 1]).unwrap();
    let y_true = Tensor::from_slice(&[2.0, 4.0, 6.0], vec![3, 1]).unwrap();

    println!("Initial learning rate: {}", optimizer.learning_rate());

    // Training loop with learning rate scheduling
    let num_epochs = 50;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Learning rate scheduling: reduce every 10 epochs
        if epoch > 0 && epoch % 10 == 0 {
            let current_lr = optimizer.learning_rate();
            let new_lr = current_lr * 0.5;
            optimizer.set_learning_rate(new_lr);
            println!(
                "Epoch {:2}: Reduced learning rate from {:.3} to {:.3}",
                epoch, current_lr, new_lr
            );
        }

        losses.push(loss.value());

        // Print progress
        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, LR = {:.3}",
                epoch,
                loss.value(),
                optimizer.learning_rate()
            );
        }
    }

    println!("Final learning rate: {}", optimizer.learning_rate());

    Ok(())
}

/// Demonstrate training monitoring and analysis
fn demonstrate_training_monitoring() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Monitoring ---");

    // Create model
    let mut weight = Tensor::randn(vec![1, 1], Some(46)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Training data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0], vec![4, 1]).unwrap();

    // Training loop with comprehensive monitoring
    let num_epochs = 30;
    let mut losses = Vec::new();
    let mut weight_history = Vec::new();
    let mut bias_history = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Record history
        losses.push(loss.value());
        weight_history.push(weight.value());
        bias_history.push(bias.value());

        // Print detailed monitoring
        if epoch % 5 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, Weight = {:.6}, Bias = {:.6}",
                epoch,
                loss.value(),
                weight.value(),
                bias.value()
            );
        }
    }

    // Analyze training progression
    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", losses[0]);
    println!("  Final loss: {:.6}", losses[losses.len() - 1]);
    println!(
        "  Loss reduction: {:.1}%",
        (losses[0] - losses[losses.len() - 1]) / losses[0] * 100.0
    );

    // Compute statistics
    let loss_mean = compute_mean(&losses);
    let loss_std = compute_std(&losses);
    let weight_change = (weight_history[weight_history.len() - 1] - weight_history[0]).abs();
    let bias_change = (bias_history[bias_history.len() - 1] - bias_history[0]).abs();

    println!("  Average loss: {:.6} ± {:.6}", loss_mean, loss_std);
    println!("  Weight change: {:.6}", weight_change);
    println!("  Bias change: {:.6}", bias_change);
    println!("  Final weight norm: {:.6}", weight.norm().value());
    println!("  Final bias: {:.6}", bias.value());

    Ok(())
}

examples/getting_started/serialization_basics.rs (line 113)

fn demonstrate_optimizer_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Optimizer Serialization ---");

    // Create an optimizer with some parameters
    let mut weight = Tensor::randn(vec![2, 2], Some(42)).with_requires_grad();
    let mut bias = Tensor::randn(vec![2], Some(43)).with_requires_grad();

    let config = AdamConfig {
        learning_rate: 0.001,
        beta1: 0.9,
        beta2: 0.999,
        eps: 1e-8,
        weight_decay: 0.0,
        amsgrad: false,
    };

    let mut optimizer = Adam::with_config(config);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    println!(
        "Created optimizer with {} parameters",
        optimizer.parameter_count()
    );
    println!("Learning rate: {}", optimizer.learning_rate());

    // Simulate some training steps
    for _ in 0..3 {
        let mut loss = weight.sum() + bias.sum();
        loss.backward(None);
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);
    }

    // Save optimizer state
    let optimizer_path = "temp_optimizer.json";
    optimizer.save_json(optimizer_path)?;
    println!("Saved optimizer to: {}", optimizer_path);

    // Load optimizer state
    let loaded_optimizer = Adam::load_json(optimizer_path)?;
    println!(
        "Loaded optimizer with {} parameters",
        loaded_optimizer.parameter_count()
    );
    println!("Learning rate: {}", loaded_optimizer.learning_rate());

    // Verify optimizer state
    assert_eq!(
        optimizer.parameter_count(),
        loaded_optimizer.parameter_count()
    );
    assert_eq!(optimizer.learning_rate(), loaded_optimizer.learning_rate());
    println!("Optimizer serialization verification: PASSED");

    Ok(())
}

/// Demonstrate format comparison and performance characteristics
fn demonstrate_format_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Format Comparison ---");

    // Create a larger tensor for comparison
    let tensor = Tensor::randn(vec![10, 10], Some(44));

    // Save in both formats
    tensor.save_json("temp_comparison.json")?;
    tensor.save_binary("temp_comparison.bin")?;

    // Compare file sizes
    let json_size = fs::metadata("temp_comparison.json")?.len();
    let binary_size = fs::metadata("temp_comparison.bin")?.len();

    println!("JSON file size: {} bytes", json_size);
    println!("Binary file size: {} bytes", binary_size);
    println!(
        "Compression ratio: {:.2}x",
        json_size as f64 / binary_size as f64
    );

    // Load and verify both formats
    let json_tensor = Tensor::load_json("temp_comparison.json")?;
    let binary_tensor = Tensor::load_binary("temp_comparison.bin")?;

    assert_eq!(tensor.shape().dims(), json_tensor.shape().dims());
    assert_eq!(tensor.shape().dims(), binary_tensor.shape().dims());
    assert_eq!(tensor.data(), json_tensor.data());
    assert_eq!(tensor.data(), binary_tensor.data());

    println!("Format comparison verification: PASSED");

    Ok(())
}

/// Demonstrate a basic model checkpointing workflow
fn demonstrate_model_checkpointing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Model Checkpointing ---");

    // Create a simple model (weights and bias)
    let mut weights = Tensor::randn(vec![2, 1], Some(45)).with_requires_grad();
    let mut bias = Tensor::randn(vec![1], Some(46)).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weights);
    optimizer.add_parameter(&bias);

    println!("Initial weights: {:?}", weights.data());
    println!("Initial bias: {:?}", bias.data());

    // Simulate training
    for epoch in 0..5 {
        let mut loss = weights.sum() + bias.sum();
        loss.backward(None);
        optimizer.step(&mut [&mut weights, &mut bias]);
        optimizer.zero_grad(&mut [&mut weights, &mut bias]);

        if epoch % 2 == 0 {
            // Save checkpoint
            let checkpoint_dir = format!("checkpoint_epoch_{}", epoch);
            fs::create_dir_all(&checkpoint_dir)?;

            weights.save_json(format!("{}/weights.json", checkpoint_dir))?;
            bias.save_json(format!("{}/bias.json", checkpoint_dir))?;
            optimizer.save_json(format!("{}/optimizer.json", checkpoint_dir))?;

            println!("Saved checkpoint for epoch {}", epoch);
        }
    }

    // Load from checkpoint
    let loaded_weights = Tensor::load_json("checkpoint_epoch_4/weights.json")?;
    let loaded_bias = Tensor::load_json("checkpoint_epoch_4/bias.json")?;
    let loaded_optimizer = Adam::load_json("checkpoint_epoch_4/optimizer.json")?;

    println!("Loaded weights: {:?}", loaded_weights.data());
    println!("Loaded bias: {:?}", loaded_bias.data());
    println!(
        "Loaded optimizer learning rate: {}",
        loaded_optimizer.learning_rate()
    );

    // Verify checkpoint integrity
    assert_eq!(weights.shape().dims(), loaded_weights.shape().dims());
    assert_eq!(bias.shape().dims(), loaded_bias.shape().dims());
    assert_eq!(optimizer.learning_rate(), loaded_optimizer.learning_rate());

    println!("Checkpointing verification: PASSED");

    Ok(())
}

Additional examples can be found in:

• examples/optimizers/adam_configurations.rs
• examples/optimizers/learning_rate_scheduling.rs
• examples/iterators/performance_optimization.rs

pub fn set_requires_grad(&mut self, requires_grad: bool)

Set gradient tracking for this tensor

Controls whether the gradtrack system tracks operations on this tensor and computes gradients during backward pass. When disabled, clears any existing gradients and gradient functions.

Arguments

• requires_grad - Whether to track gradients for this tensor

Performance

• Time Complexity: O(1) - simple field assignment
• Memory: May free gradient storage when disabled
• Overhead: Zero overhead when gradients disabled

Examples

use train_station::Tensor;

let mut tensor = Tensor::ones(vec![2, 3]);
tensor.set_requires_grad(true);
assert!(tensor.requires_grad());

// Disable gradient tracking
tensor.set_requires_grad(false);
assert!(!tensor.requires_grad());

Examples found in repository

examples/RL_training/dqn.rs (line 102)

    fn set_requires_grad_all(&mut self, enable: bool) {
        for l in &mut self.layers {
            l.weight.set_requires_grad(enable);
            l.bias.set_requires_grad(enable);
        }
    }

    // In-place copy (preserve tensor IDs and optimizer links on targets)
    fn copy_from(&mut self, other: &Self) {
        for (t, s) in self.layers.iter_mut().zip(other.layers.iter()) {
            {
                let src = s.weight.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = s.bias.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }

examples/RL_training/td3.rs (line 109)

    fn set_requires_grad_all(&mut self, enable: bool) {
        for l in &mut self.layers {
            l.weight.set_requires_grad(enable);
            l.bias.set_requires_grad(enable);
        }
    }

    fn copy_from(&mut self, other: &Self) {
        for (t, s) in self.layers.iter_mut().zip(other.layers.iter()) {
            {
                let src = s.weight.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = s.bias.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }

    fn soft_update_from(&mut self, source: &Self, tau: f32) {
        let _ng = NoGradTrack::new();
        for (t, s) in self.layers.iter_mut().zip(source.layers.iter()) {
            // In-place Polyak update to preserve tensor IDs (no optimizer relink needed)
            let new_w = t
                .weight
                .mul_scalar(1.0 - tau)
                .add_tensor(&s.weight.mul_scalar(tau));
            let new_b = t
                .bias
                .mul_scalar(1.0 - tau)
                .add_tensor(&s.bias.mul_scalar(tau));
            {
                let src = new_w.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = new_b.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }

pub fn retain_grad(self) -> Self

Mark this tensor to retain gradients after backward, even if it is non-leaf.

Builder-style API: returns self with retain_grad=true. Call materialize_grad() or grad_or_fetch() after backward to copy the accumulated gradient from the GradGraph into self.grad so grad() works.

pub fn retain_grad_(&mut self, enable: bool)

In-place variant to enable or disable gradient retention for non-leaf tensors

pub fn requires_grad(&self) -> bool

Check if this tensor requires gradients

Returns

true if gradient tracking is enabled for this tensor

Examples

use train_station::Tensor;

let tensor = Tensor::new(vec![2, 3]);
assert!(!tensor.requires_grad());

let grad_tensor = Tensor::ones(vec![2, 3]).with_requires_grad();
assert!(grad_tensor.requires_grad());

Examples found in repository

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 164)

fn demonstrate_layer_creation() {
    println!("--- Layer Creation ---");

    let layer = LinearLayer::new(3, 2, Some(42));

    println!("Created linear layer:");
    println!("  Input size: {}", layer.input_size);
    println!("  Output size: {}", layer.output_size);
    println!("  Parameter count: {}", layer.parameter_count());
    println!("  Weight shape: {:?}", layer.weight.shape().dims());
    println!("  Bias shape: {:?}", layer.bias.shape().dims());
    println!("  Weight requires grad: {}", layer.weight.requires_grad());
    println!("  Bias requires grad: {}", layer.bias.requires_grad());
}

/// Demonstrate forward pass with gradient tracking
fn demonstrate_forward_pass() {
    println!("\n--- Forward Pass (with gradients) ---");

    let layer = LinearLayer::new(3, 2, Some(43));

    // Single input
    let input = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3]).unwrap();
    let output = layer.forward(&input);

    println!("Single input:");
    println!("  Input: {:?}", input.data());
    println!("  Output: {:?}", output.data());
    println!("  Output requires grad: {}", output.requires_grad());

    // Batch input
    let batch_input = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
    let batch_output = layer.forward(&batch_input);

    println!("Batch input:");
    println!("  Input shape: {:?}", batch_input.shape().dims());
    println!("  Output shape: {:?}", batch_output.shape().dims());
    println!("  Output requires grad: {}", batch_output.requires_grad());
}

/// Demonstrate forward pass without gradient tracking
fn demonstrate_forward_pass_no_grad() {
    println!("\n--- Forward Pass (no gradients) ---");

    let layer = LinearLayer::new(3, 2, Some(44));

    // Single input
    let input = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3]).unwrap();
    let output = layer.forward_no_grad(&input);

    println!("Single input (no grad):");
    println!("  Input: {:?}", input.data());
    println!("  Output: {:?}", output.data());
    println!("  Output requires grad: {}", output.requires_grad());

    // Compare with grad version
    let output_with_grad = layer.forward(&input);
    println!("Comparison:");
    println!(
        "  Same values: {}",
        output.data() == output_with_grad.data()
    );
    println!("  No grad requires grad: {}", output.requires_grad());
    println!(
        "  With grad requires grad: {}",
        output_with_grad.requires_grad()
    );
}

examples/iterators/element_iteration.rs (line 192)

fn demonstrate_gradient_tracking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Gradient Tracking ---");

    // Create a tensor with gradient tracking enabled
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?.with_requires_grad();
    println!("Input tensor (requires_grad): {:?}", tensor.data());

    // Perform element-wise operations through iteration
    let result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Apply a complex transformation: (x^2 + 1) * 2
            elem.pow_scalar(2.0).add_scalar(1.0).mul_scalar(2.0)
        })
        .collect();

    println!("Result tensor: {:?}", result.data());
    println!("Result requires_grad: {}", result.requires_grad());

    // Compute gradients
    let mut loss = result.sum();
    loss.backward(None);

    println!("Loss: {:.6}", loss.value());
    println!("Input gradients: {:?}", tensor.grad().map(|g| g.data()));

    Ok(())
}

examples/getting_started/optimizer_basics.rs (line 58)

fn demonstrate_basic_optimizer_setup() {
    println!("--- Basic Optimizer Setup ---");

    // Create parameters that require gradients
    let weight = Tensor::randn(vec![3, 2], Some(42)).with_requires_grad();
    let bias = Tensor::zeros(vec![2]).with_requires_grad();

    println!("Created parameters:");
    println!(
        "  Weight: shape {:?}, requires_grad: {}",
        weight.shape().dims(),
        weight.requires_grad()
    );
    println!(
        "  Bias: shape {:?}, requires_grad: {}",
        bias.shape().dims(),
        bias.requires_grad()
    );

    // Create Adam optimizer with default configuration
    let mut optimizer = Adam::new();
    println!(
        "Created Adam optimizer with learning rate: {}",
        optimizer.learning_rate()
    );

    // Add parameters to optimizer
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);
    println!(
        "Added {} parameters to optimizer",
        optimizer.parameter_count()
    );

    // Create optimizer with custom configuration
    let config = AdamConfig {
        learning_rate: 0.01,
        beta1: 0.9,
        beta2: 0.999,
        eps: 1e-8,
        weight_decay: 0.0,
        amsgrad: false,
    };

    let mut custom_optimizer = Adam::with_config(config);
    custom_optimizer.add_parameter(&weight);
    custom_optimizer.add_parameter(&bias);

    println!(
        "Created custom optimizer with learning rate: {}",
        custom_optimizer.learning_rate()
    );

    // Demonstrate parameter linking
    println!("Parameter linking completed successfully");
}

examples/neural_networks/feedforward_network.rs (line 332)

fn demonstrate_forward_pass() {
    println!("\n--- Forward Pass ---");

    let config = FeedForwardConfig {
        input_size: 3,
        hidden_sizes: vec![5, 3],
        output_size: 2,
        use_bias: true,
    };
    let network = FeedForwardNetwork::new(config, Some(43));

    // Single input
    let input = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3]).unwrap();
    let output = network.forward(&input);

    println!("Single input forward pass:");
    println!("  Input shape: {:?}", input.shape().dims());
    println!("  Output shape: {:?}", output.shape().dims());
    println!("  Output: {:?}", output.data());
    println!("  Output requires grad: {}", output.requires_grad());

    // Batch input
    let batch_input = Tensor::from_slice(
        &[
            1.0, 2.0, 3.0, // Sample 1
            4.0, 5.0, 6.0, // Sample 2
            7.0, 8.0, 9.0, // Sample 3
        ],
        vec![3, 3],
    )
    .unwrap();
    let batch_output = network.forward(&batch_input);

    println!("Batch input forward pass:");
    println!("  Input shape: {:?}", batch_input.shape().dims());
    println!("  Output shape: {:?}", batch_output.shape().dims());
    println!("  Output requires grad: {}", batch_output.requires_grad());

    // Compare with no-grad version
    let output_no_grad = network.forward_no_grad(&input);
    println!("No-grad comparison:");
    println!("  Same values: {}", output.data() == output_no_grad.data());
    println!("  With grad requires grad: {}", output.requires_grad());
    println!(
        "  No grad requires grad: {}",
        output_no_grad.requires_grad()
    );
}

examples/iterators/performance_optimization.rs (line 416)

fn demonstrate_optimization_techniques() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Optimization Techniques ---");

    let size = 50000;
    let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
    let tensor = Tensor::from_slice(&data, vec![size])?;

    println!("Optimizing processing for size: {}", size);

    // Technique 1: Operation fusion
    println!("\nTechnique 1: Operation Fusion");
    let start = Instant::now();
    let fused_result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Fuse multiple operations into single chain
            elem.mul_scalar(2.0).add_scalar(1.0).pow_scalar(2.0).sqrt()
        })
        .collect();
    let fused_time = start.elapsed();

    // Technique 2: Conditional optimization
    println!("\nTechnique 2: Conditional Optimization");
    let start = Instant::now();
    let conditional_result: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val < size as f32 / 2.0 {
                elem.mul_scalar(2.0) // Simple operation for small values
            } else {
                elem.pow_scalar(2.0).sqrt() // Complex operation for large values
            }
        })
        .collect();
    let conditional_time = start.elapsed();

    // Technique 3: Cache-friendly processing
    println!("\nTechnique 3: Cache-Friendly Processing");
    let start = Instant::now();
    let cache_friendly_result: Tensor = tensor
        .iter()
        .take(1000) // Process in cache-friendly chunks
        .map(|elem| elem.mul_scalar(2.0))
        .collect();
    let cache_friendly_time = start.elapsed();

    // Technique 4: Memory pooling simulation
    println!("\nTechnique 4: Memory Pooling Simulation");
    let start = Instant::now();
    let pooled_result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 100 == 0) // Process every 100th element
        .map(|(_, elem)| elem.pow_scalar(2.0))
        .collect();
    let pooled_time = start.elapsed();

    // Report optimization results
    println!("  Fused operations: {:?}", fused_time);
    println!("  Conditional optimization: {:?}", conditional_time);
    println!("  Cache-friendly processing: {:?}", cache_friendly_time);
    println!("  Memory pooling simulation: {:?}", pooled_time);

    // Performance analysis
    let fastest = fused_time
        .min(conditional_time)
        .min(cache_friendly_time)
        .min(pooled_time);
    println!("  Fastest technique: {:?}", fastest);

    // Memory efficiency analysis
    println!("  Fused result size: {}", fused_result.size());
    println!("  Conditional result size: {}", conditional_result.size());
    println!(
        "  Cache-friendly result size: {}",
        cache_friendly_result.size()
    );
    println!("  Pooled result size: {}", pooled_result.size());

    // Technique 5: Gradient optimization
    println!("\nTechnique 5: Gradient Optimization");
    let grad_tensor = tensor.with_requires_grad();
    let start = Instant::now();

    let grad_result: Tensor = grad_tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
        .collect();

    let mut loss = grad_result.sum();
    loss.backward(None);
    let grad_time = start.elapsed();

    println!("  Gradient computation: {:?}", grad_time);
    println!(
        "  Gradient tracking enabled: {}",
        grad_result.requires_grad()
    );

    Ok(())
}

pub fn grad(&self) -> Option<&Tensor>

Get a reference to this tensor’s locally cached gradient (if any)

This accessor returns only the gradient cached on this tensor’s grad field. It does NOT query the global/autograd gradient storage. For leaf tensors, gradients are accumulated and tracked by the grad engine and are not automatically written back to the local grad field.

To make grad() return Some(&Tensor):

• Enable retention on this tensor (typically for non-leaf tensors) with retain_grad_(&mut tensor, true) or tensor.retain_grad()
• After backward(), call tensor.materialize_grad() or tensor.grad_or_fetch() to copy the accumulated gradient from the autograd engine into self.grad

If you want to read gradients without caching them locally, prefer grad_owned, which consults the global gradient storage.

Returns

Optional reference to the locally cached gradient tensor, or None if not materialized on this tensor.

Examples

use train_station::Tensor;

let mut x = Tensor::ones(vec![2, 3]).with_requires_grad();
let mut loss = x.sum();
loss.backward(None);
// Without materialization, grad() typically returns None for leaves
assert!(x.grad().is_none());
// Materialize to cache locally so grad() works
x.retain_grad_ (true);
x.materialize_grad();
assert!(x.grad().is_some());

Examples found in repository

examples/iterators/element_iteration.rs (line 199)

fn demonstrate_gradient_tracking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Gradient Tracking ---");

    // Create a tensor with gradient tracking enabled
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?.with_requires_grad();
    println!("Input tensor (requires_grad): {:?}", tensor.data());

    // Perform element-wise operations through iteration
    let result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Apply a complex transformation: (x^2 + 1) * 2
            elem.pow_scalar(2.0).add_scalar(1.0).mul_scalar(2.0)
        })
        .collect();

    println!("Result tensor: {:?}", result.data());
    println!("Result requires_grad: {}", result.requires_grad());

    // Compute gradients
    let mut loss = result.sum();
    loss.backward(None);

    println!("Loss: {:.6}", loss.value());
    println!("Input gradients: {:?}", tensor.grad().map(|g| g.data()));

    Ok(())
}

pub fn materialize_grad(&mut self) -> bool

Fetch the accumulated gradient after backward and cache it on this tensor if retain_grad is enabled.

Returns true if a gradient was found and cached. After a successful call, grad() will return Some(&Tensor) even for non-leaf tensors.

pub fn grad_or_fetch(&mut self) -> Option<&Tensor>

Convenience accessor: if retain_grad is enabled, fetch and cache the gradient on first access so callers can immediately get a reference.

pub fn grad_owned(&self) -> Option<Tensor>

Get the accumulated gradient from the autograd engine as an owned tensor

This accessor queries the global/shared gradient storage for this tensor’s ID and returns the current accumulated gradient by value. It complements grad, which only returns a locally cached reference.

• Works for leaf tensors without any prior materialization step
• Does not modify or clear internal gradient state
• Suitable for optimizers and logging that need the latest gradients

Returns

Some(Tensor) containing the current accumulated gradient when available, otherwise None.

Examples

use train_station::Tensor;

let mut x = Tensor::ones(vec![2, 3]).with_requires_grad();
let mut loss = x.sum();
loss.backward(None);

// Fetch directly from autograd storage (no materialization required)
let g = x.grad_owned().unwrap();
assert_eq!(g.shape().dims(), vec![2, 3]);

Examples found in repository

examples/RL_training/dqn.rs (line 280)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

fn params_l2_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let _ng = NoGradTrack::new();
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        for &v in p.data() {
            total_sq += v * v;
        }
    }
    total_sq.sqrt()
}

// Pseudo-Huber loss: sqrt(1 + diff^2) - 1 (smooth, robust)
fn pseudo_huber_mean(diff: &Tensor) -> Tensor {
    diff.pow_scalar(2.0)
        .add_scalar(1.0)
        .sqrt()
        .sub_scalar(1.0)
        .mean()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== DQN Example (YardEnv discrete) ===");

    // Dims
    let state_dim = 3usize;
    let action_dim = 3usize;

    // Hparams
    let gamma = 0.99f32;
    let batch_size = 64usize;
    let start_steps = 200usize;
    let target_update_interval = 200usize; // hard update cadence
    let max_grad_norm = 1.0f32;
    let mut epsilon = 1.0f32;
    let eps_min = 0.05f32;
    let eps_decay_steps = 2_000usize; // linear decay
    let total_steps = std::env::var("DQN_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3000usize);

    // Models
    let mut q_net = QNet::new(state_dim, action_dim, Some(7));
    let mut q_targ = QNet::new(state_dim, action_dim, Some(8));
    q_targ.net.copy_from(&q_net.net);
    q_targ.set_requires_grad_all(false);

    // Optimizer
    let mut q_opt = Adam::with_learning_rate(3e-4);
    for p in q_net.parameters() {
        q_opt.add_parameter(p);
    }

    // Replay + env
    let mut rb = ReplayBuffer::new(100_000, state_dim);
    let mut env = YardEnv::new(12345);
    let mut rng = SmallRng::new(999_111);

    // Metrics
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    for t in 0..total_steps {
        // Epsilon-greedy action
        let action_index = if t < start_steps || rng.next_f32() < epsilon {
            rng.sample_index(action_dim)
        } else {
            let _ng = NoGradTrack::new();
            let q_vals = q_net.forward(&state);
            let row = q_vals.data();
            let mut best_i = 0usize;
            let mut best_v = row[0];
            for (i, &r) in row.iter().enumerate().take(action_dim).skip(1) {
                if r > best_v {
                    best_v = r;
                    best_i = i;
                }
            }
            best_i
        };

        // Env step
        let (next_state, reward, done) = env.step(action_index);
        episode_return += reward;

        // Store
        let s_slice = state.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            action_index,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        // Reset on done
        state = if done {
            let st = env.reset();
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Epsilon linear decay
        if t < eps_decay_steps {
            epsilon = (1.0 - (t as f32) / (eps_decay_steps as f32)) * (1.0 - eps_min) + eps_min;
        }

        // Train
        if rb.can_sample(batch_size) {
            let (s, a_idx, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Double DQN target: a* = argmax_a Q_online(s2,a); y = r + (1-d)*gamma*Q_target(s2, a*)
            let target_q = {
                let _ng = NoGradTrack::new();
                let q_online_s2 = q_net.forward(&s2);
                // argmax per row (manual on CPU)
                let row_stride = action_dim;
                let qd = q_online_s2.data();
                let mut next_actions: Vec<usize> = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let base = i * row_stride;
                    let mut bi = 0usize;
                    let mut bv = qd[base];
                    for j in 1..action_dim {
                        let v = qd[base + j];
                        if v > bv {
                            bv = v;
                            bi = j;
                        }
                    }
                    next_actions.push(bi);
                }
                let q_targ_s2 = q_targ.forward(&s2);
                let q_targ_g = q_targ_s2.gather(1, &next_actions, &[batch_size, 1]);
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&q_targ_g))
            };

            // Q(s,a) for current actions
            // Zero grads first
            {
                let mut params = q_net.parameters();
                q_opt.zero_grad(&mut params);
            }

            let q_all = q_net.forward(&s);
            let q_sa = q_all.gather(1, &a_idx, &[batch_size, 1]);
            let diff = q_sa.sub_tensor(&target_q);
            let mut loss = pseudo_huber_mean(&diff);
            loss.backward(None);

            // Step (filter only params with grads)
            {
                let params = q_net.parameters();
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    let gn = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    q_opt.step(&mut with_grads);
                    q_opt.zero_grad(&mut with_grads);
                    if t % 100 == 0 {
                        let mut pn = q_net.parameters();
                        let pn_l2 = params_l2_norm(&mut pn);
                        let q_mean = q_all.mean().value();
                        println!(
                            "t={:5} | loss={:.4} | q_mean={:.3} | grad_norm={:.3} | param_norm={:.3} | eps={:.3}",
                            t, loss.value(), q_mean, gn, pn_l2, epsilon
                        );
                    }
                }
            }

            // Target hard update
            if t % target_update_interval == 0 {
                q_targ.net.copy_from(&q_net.net);
            }

            // Clear graphs
            clear_all_graphs_known();
        }
    }

    println!("=== DQN training finished ===");
    Ok(())
}

examples/RL_training/ppo_continuous.rs (line 296)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

examples/RL_training/ppo_discrete.rs (line 255)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

// log-softmax for selected actions: given logits [B,A] and actions Vec<usize> -> log_prob [B,1]
fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

// probability ratio = exp(new_logp - old_logp)
fn ratio_from_logps(new_logp: &Tensor, old_logp: &Tensor) -> Tensor {
    new_logp.sub_tensor(old_logp).exp()
}

// Clamp ratio to [1-clip, 1+clip] using ReLU-based clamp (no custom ops)
fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/supervised_training/supervised_bce.rs (line 26)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn accuracy(pred: &Tensor, targets: &Tensor) -> f32 {
    // pred: [B,1] with sigmoid; threshold at 0.5
    let p = pred.data();
    let t = targets.data();
    let mut correct = 0usize;
    for i in 0..p.len() {
        let yhat = if p[i] >= 0.5 { 1.0 } else { 0.0 };
        if (yhat - t[i]).abs() < 1e-6 {
            correct += 1;
        }
    }
    correct as f32 / (p.len() as f32)
}

// Numerically stable BCE with logits:
// L = mean( relu(z) - z*y + log(1 + exp(-|z|)) )
fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Supervised FFN Example (XOR) ===");

    // Dataset: XOR (repeat to form a small batch)
    let inputs: Vec<f32> = vec![
        0.0, 0.0, // -> 0
        0.0, 1.0, // -> 1
        1.0, 0.0, // -> 1
        1.0, 1.0, // -> 0
    ];
    let targets: Vec<f32> = vec![0.0, 1.0, 1.0, 0.0];

    // Repeat the base patterns to stabilize training
    let repeats = 64usize; // effective batch = 4 * repeats = 256
    let mut xs = Vec::with_capacity(repeats * inputs.len());
    let mut ys = Vec::with_capacity(repeats * targets.len());
    for _ in 0..repeats {
        xs.extend_from_slice(&inputs);
        ys.extend_from_slice(&targets);
    }

    let batch = xs.len() / 2; // two features
    let x_t = Tensor::from_slice(&xs, vec![batch, 2]).unwrap();
    let y_t = Tensor::from_slice(&ys, vec![batch, 1]).unwrap();

    // Model config: 2 -> 32 -> 32 -> 1, final sigmoid via loss path
    let cfg = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![32, 32],
        output_size: 1,
        use_bias: true,
    };
    let mut net = FeedForwardNetwork::new(cfg, Some(777));

    // Optimizer and parameter linking
    let mut opt = Adam::with_learning_rate(1e-3);
    for p in net.parameters() {
        opt.add_parameter(p);
    }

    let epochs = 1000usize;
    let max_grad_norm = 1.0f32;
    let mut best_loss = f32::INFINITY;
    let mut best_acc = 0.0f32;

    for e in 0..epochs {
        // Zero grads each iteration
        {
            let mut params = net.parameters();
            opt.zero_grad(&mut params);
        }

        // Forward -> logits; use numerically stable BCE-with-logits for loss
        let logits = net.forward(&x_t);
        let mut loss = bce_with_logits(&logits, &y_t);
        loss.backward(None);

        // Step only params with grads
        {
            let params = net.parameters();
            let mut with_grads: Vec<&mut Tensor> = Vec::new();
            for p in params {
                if p.grad_owned().is_some() {
                    with_grads.push(p);
                }
            }
            if !with_grads.is_empty() {
                clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                opt.step(&mut with_grads);
                opt.zero_grad(&mut with_grads);
            }
        }

        // Metrics (use sigmoid only for reporting accuracy)
        let preds = logits.sigmoid();
        let acc = accuracy(&preds, &y_t);
        if loss.value() < best_loss {
            best_loss = loss.value();
        }
        if acc > best_acc {
            best_acc = acc;
        }
        if e % 10 == 0 || e + 1 == epochs {
            println!(
                "epoch {:4} | loss={:.5} acc={:.3} | best_loss={:.5} best_acc={:.3}",
                e,
                loss.value(),
                acc,
                best_loss,
                best_acc
            );
        }

        // Clear graphs to avoid stale accumulation across epochs
        clear_all_graphs_known();
    }

    // Quick sanity check predictions
    let test = Tensor::from_slice(&inputs, vec![4, 2]).unwrap();
    let out = net.forward(&test).sigmoid();
    println!("predictions (approx): {:?}", out.data());

    println!("=== Supervised training finished ===");
    Ok(())
}

examples/supervised_training/supervised_classification.rs (line 26)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

// Cross-entropy over logits: CE = -mean(log_softmax(logits)[range, labels])
fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

fn accuracy_from_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    num_classes: usize,
) -> f32 {
    let row = logits.data();
    let mut correct = 0usize;
    for (i, &label) in labels.iter().enumerate().take(batch) {
        let base = i * num_classes;
        let mut best_j = 0usize;
        let mut best_v = row[base];
        for j in 1..num_classes {
            let v = row[base + j];
            if v > best_v {
                best_v = v;
                best_j = j;
            }
        }
        if best_j == label {
            correct += 1;
        }
    }
    correct as f32 / batch as f32
}

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Supervised Classification Example (Cross-Entropy) ===");

    // Synthetic 2D inputs, 3 classes with linear-ish separations
    let n = 1200usize;
    let classes = 3usize;
    let mut xs: Vec<f32> = Vec::with_capacity(n * 2);
    let mut ys: Vec<usize> = Vec::with_capacity(n);

    // Simple RNG
    let mut state: u64 = 424242;
    let mut rand_f32 = || {
        state = state.wrapping_mul(1664525).wrapping_add(1013904223);
        (state >> 16) as f32 / (u32::MAX as f32)
    };

    for _ in 0..n {
        let x1 = rand_f32() * 4.0 - 2.0;
        let x2 = rand_f32() * 4.0 - 2.0;
        // Class by quadrant-ish rule with noise
        let mut c = if x1 + 0.5 * x2 > 0.5 {
            0
        } else if x1 - x2 < -0.5 {
            1
        } else {
            2
        };
        if rand_f32() < 0.05 {
            c = (c + 1) % classes;
        }
        xs.push(x1);
        xs.push(x2);
        ys.push(c);
    }

    // Normalize inputs per-feature to [-1, 1]
    let mut min1 = f32::INFINITY;
    let mut max1 = f32::NEG_INFINITY;
    let mut min2 = f32::INFINITY;
    let mut max2 = f32::NEG_INFINITY;
    for i in (0..xs.len()).step_by(2) {
        let a = xs[i];
        let b = xs[i + 1];
        if a < min1 {
            min1 = a;
        }
        if a > max1 {
            max1 = a;
        }
        if b < min2 {
            min2 = b;
        }
        if b > max2 {
            max2 = b;
        }
    }
    let rng1 = (max1 - min1).max(1e-8);
    let rng2 = (max2 - min2).max(1e-8);
    for i in (0..xs.len()).step_by(2) {
        let a = xs[i];
        let b = xs[i + 1];
        xs[i] = 2.0 * (a - min1) / rng1 - 1.0;
        xs[i + 1] = 2.0 * (b - min2) / rng2 - 1.0;
    }

    // Train/Val split (80/20)
    let n_train = (n as f32 * 0.8) as usize;
    let x_train = Tensor::from_slice(&xs[..n_train * 2], vec![n_train, 2]).unwrap();
    let y_train = ys[..n_train].to_vec();
    let x_val = Tensor::from_slice(&xs[n_train * 2..], vec![n - n_train, 2]).unwrap();
    let y_val = ys[n_train..].to_vec();

    // Model: 2 -> 64 -> 64 -> 3 (logits)
    let cfg = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![64, 64],
        output_size: classes,
        use_bias: true,
    };
    let mut net = FeedForwardNetwork::new(cfg, Some(303));

    // Optimizer
    let mut opt = Adam::with_learning_rate(1e-3);
    for p in net.parameters() {
        opt.add_parameter(p);
    }

    let epochs = 300usize;
    let max_grad_norm = 1.0f32;
    let mut best_val_acc = 0.0f32;
    let mut best_val_loss = f32::INFINITY;

    for e in 0..epochs {
        // Zero grads
        {
            let mut params = net.parameters();
            opt.zero_grad(&mut params);
        }

        // Forward logits
        let logits = net.forward(&x_train);
        let mut loss = cross_entropy_logits(&logits, &y_train, n_train, classes);
        loss.backward(None);

        // Step clipped
        {
            let params = net.parameters();
            let mut with_grads: Vec<&mut Tensor> = Vec::new();
            for p in params {
                if p.grad_owned().is_some() {
                    with_grads.push(p);
                }
            }
            if !with_grads.is_empty() {
                clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                opt.step(&mut with_grads);
                opt.zero_grad(&mut with_grads);
            }
        }

        // Metrics
        let train_acc = accuracy_from_logits(&logits, &y_train, n_train, classes);
        let val_logits = net.forward(&x_val);
        let val_loss = cross_entropy_logits(&val_logits, &y_val, n - n_train, classes).value();
        let val_acc = accuracy_from_logits(&val_logits, &y_val, n - n_train, classes);
        if val_acc > best_val_acc {
            best_val_acc = val_acc;
        }
        if val_loss < best_val_loss {
            best_val_loss = val_loss;
        }

        if e % 10 == 0 || e + 1 == epochs {
            println!(
                "epoch {:4} | loss={:.4} acc={:.3} | val_loss={:.4} val_acc={:.3} | best_val_acc={:.3}",
                e, loss.value(), train_acc, val_loss, val_acc, best_val_acc
            );
        }

        clear_all_graphs_known();
    }

    // Quick sample preds via softmax
    let samples = Tensor::from_slice(&[-1.0, -1.0, 0.0, 0.0, 1.0, 1.0], vec![3, 2]).unwrap();
    let sm = net.forward(&samples).softmax(1);
    println!("sample class probs: {:?}", sm.data());

    println!("=== Supervised classification finished ===");
    Ok(())
}

examples/supervised_training/supervised_regression.rs (line 25)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn mse(pred: &Tensor, target: &Tensor) -> Tensor {
    pred.sub_tensor(target).pow_scalar(2.0).mean()
}

fn rmse(pred: &Tensor, target: &Tensor) -> f32 {
    mse(pred, target).sqrt().value()
}

fn r2_score(pred: &Tensor, target: &Tensor) -> f32 {
    // R^2 = 1 - SS_res / SS_tot
    let y = target;
    let y_mean = y.mean();
    let ss_res = pred.sub_tensor(y).pow_scalar(2.0).sum();
    let ss_tot = y.sub_tensor(&y_mean).pow_scalar(2.0).sum();
    let ss_res_v = ss_res.value();
    let ss_tot_v = ss_tot.value().max(1e-12); // avoid divide by zero
    1.0 - (ss_res_v / ss_tot_v)
}

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Supervised Regression Example (MSE) ===");

    // Generate simple synthetic data: y = 2*x1 - 3*x2 + 0.5 + noise
    let n = 1024usize;
    let mut xs: Vec<f32> = Vec::with_capacity(n * 2);
    let mut ys: Vec<f32> = Vec::with_capacity(n);
    // Simple LCG RNG for reproducibility
    let mut state: u64 = 123456789;
    let mut rand_f32 = || {
        state = state.wrapping_mul(1664525).wrapping_add(1013904223);
        (state >> 16) as f32 / (u32::MAX as f32)
    };
    for _ in 0..n {
        let x1 = rand_f32() * 2.0 - 1.0;
        let x2 = rand_f32() * 2.0 - 1.0;
        let noise = (rand_f32() * 2.0 - 1.0) * 0.05;
        let y = 2.0 * x1 - 3.0 * x2 + 0.5 + noise;
        xs.push(x1);
        xs.push(x2);
        ys.push(y);
    }

    // Normalize targets to [-1, 1] (max-abs scaling) for reasonable loss magnitudes
    let mut max_abs = 0.0f32;
    for &v in &ys {
        let a = v.abs();
        if a > max_abs {
            max_abs = a;
        }
    }
    if max_abs < 1e-8 {
        max_abs = 1.0;
    }
    for v in ys.iter_mut() {
        *v /= max_abs;
    }

    // Normalize inputs per-feature to [-1, 1] (min-max scaling)
    let mut min1 = f32::INFINITY;
    let mut max1 = f32::NEG_INFINITY;
    let mut min2 = f32::INFINITY;
    let mut max2 = f32::NEG_INFINITY;
    for i in (0..xs.len()).step_by(2) {
        let a = xs[i];
        let b = xs[i + 1];
        if a < min1 {
            min1 = a;
        }
        if a > max1 {
            max1 = a;
        }
        if b < min2 {
            min2 = b;
        }
        if b > max2 {
            max2 = b;
        }
    }
    let rng1 = (max1 - min1).max(1e-8);
    let rng2 = (max2 - min2).max(1e-8);
    for i in (0..xs.len()).step_by(2) {
        let a = xs[i];
        let b = xs[i + 1];
        xs[i] = 2.0 * (a - min1) / rng1 - 1.0;
        xs[i + 1] = 2.0 * (b - min2) / rng2 - 1.0;
    }

    // Train/Val split (80/20)
    let n_train = (n as f32 * 0.8) as usize;
    let x_train = Tensor::from_slice(&xs[..n_train * 2], vec![n_train, 2]).unwrap();
    let y_train = Tensor::from_slice(&ys[..n_train], vec![n_train, 1]).unwrap();
    let x_val = Tensor::from_slice(&xs[n_train * 2..], vec![n - n_train, 2]).unwrap();
    let y_val = Tensor::from_slice(&ys[n_train..], vec![n - n_train, 1]).unwrap();

    // Model config: 2 -> 64 -> 64 -> 1
    let cfg = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![64, 64],
        output_size: 1,
        use_bias: true,
    };
    let mut net = FeedForwardNetwork::new(cfg, Some(2025));

    // Optimizer and parameter linking
    let mut opt = Adam::with_learning_rate(1e-3);
    for p in net.parameters() {
        opt.add_parameter(p);
    }

    let epochs = 400usize;
    let max_grad_norm = 1.0f32;
    let mut best_val_rmse = f32::INFINITY;
    let mut best_val_r2 = -f32::INFINITY;

    for e in 0..epochs {
        // Zero grads
        {
            let mut params = net.parameters();
            opt.zero_grad(&mut params);
        }

        // Forward
        let pred = net.forward(&x_train);
        let mut loss = mse(&pred, &y_train);
        loss.backward(None);

        // Step
        {
            let params = net.parameters();
            let mut with_grads: Vec<&mut Tensor> = Vec::new();
            for p in params {
                if p.grad_owned().is_some() {
                    with_grads.push(p);
                }
            }
            if !with_grads.is_empty() {
                clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                opt.step(&mut with_grads);
                opt.zero_grad(&mut with_grads);
            }
        }

        // Metrics
        let train_rmse = rmse(&pred, &y_train);
        let train_r2 = r2_score(&pred, &y_train);
        let val_pred = net.forward(&x_val);
        let val_rmse = rmse(&val_pred, &y_val);
        let val_r2 = r2_score(&val_pred, &y_val);
        if val_rmse < best_val_rmse {
            best_val_rmse = val_rmse;
        }
        if val_r2 > best_val_r2 {
            best_val_r2 = val_r2;
        }

        if e % 20 == 0 || e + 1 == epochs {
            // Clamp displayed R^2 to avoid huge negative prints on early epochs
            let train_r2_disp = train_r2.max(-10.0);
            let val_r2_disp = val_r2.max(-10.0);
            println!(
                "epoch {:4} | train_rmse={:.4} r2={:.3} | val_rmse={:.4} r2={:.3} | best_val_rmse={:.4} best_val_r2={:.3}",
                e, train_rmse, train_r2_disp, val_rmse, val_r2_disp, best_val_rmse, best_val_r2
            );
        }

        clear_all_graphs_known();
    }

    // Quick sanity predictions on small samples
    let sample = Tensor::from_slice(&[0.5, -0.25, -0.8, 0.3], vec![2, 2]).unwrap();
    let sample_pred = net.forward(&sample);
    println!("samples pred: {:?}", sample_pred.data());

    println!("=== Supervised regression finished ===");
    Ok(())
}

Additional examples can be found in:

• examples/RL_training/td3.rs
• examples/getting_started/optimizer_basics.rs

pub fn id(&self) -> usize

Get the unique ID of this tensor

Returns the unique identifier assigned to this tensor during creation. This ID is used for gradtrack tracking and tensor identification.

Returns

Unique tensor ID as usize

Examples

use train_station::Tensor;

let tensor1 = Tensor::new(vec![2, 3]);
let tensor2 = Tensor::new(vec![2, 3]);
assert_ne!(tensor1.id(), tensor2.id()); // Each tensor has unique ID

pub fn detach(&self) -> Self

Detach this tensor from the computation graph

Returns a new tensor with the same data but no gradient tracking. This is useful when you want to use a tensor in inference without affecting the computation graph.

Returns

A new tensor with the same data but gradient tracking disabled

Examples

use train_station::Tensor;

let tensor = Tensor::ones(vec![2, 3]).with_requires_grad();
let detached = tensor.detach();
assert!(!detached.requires_grad());
assert_eq!(tensor.size(), detached.size());

pub fn detach_(&mut self)

Create a new tensor that doesn’t track gradients from this one

Similar to detach() but modifies this tensor in place. This is useful when you want to disable gradient tracking for the current tensor without creating a copy.

Examples

use train_station::Tensor;

let mut tensor = Tensor::ones(vec![2, 3]).with_requires_grad();
assert!(tensor.requires_grad());
tensor.detach_();
assert!(!tensor.requires_grad());

pub fn backward(&mut self, grad_output: Option<Tensor>)

Entry point for backward pass on this tensor

Computes gradients for all tensors in the computation graph that have requires_grad set to true. This is the main entry point for automatic differentiation.

Arguments

• grad_output - Optional gradient tensor for the output. If None, assumes the tensor is a scalar (e.g., loss value) and uses a tensor of ones.

Examples

use train_station::Tensor;

let mut tensor = Tensor::ones(vec![2, 3]).with_requires_grad();
let mut result = tensor.add_scalar(5.0);
result.backward(None);
// Note: Gradient computation depends on the gradtrack system implementation

Examples found in repository

examples/neural_networks/basic_encoder.rs (line 95)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Encoder Example ===");

    let batch = 2usize;
    let seq = 6usize;
    let embed = 32usize;
    let heads = 4usize;

    let input = Tensor::randn(vec![batch, seq, embed], Some(11));
    let mut enc = EncoderBlock::new(embed, heads, Some(123));

    // Example: no mask (set Some(mask) to use masking)
    let out = enc.forward(&input, None);
    println!("Output shape: {:?}", out.shape().dims());

    // Verify gradients/optimization
    let mut opt = Adam::with_learning_rate(0.01);
    let mut params = enc.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/neural_networks/basic_decoder.rs (line 106)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Decoder Example ===");

    let batch = 2usize;
    let src = 7usize;
    let tgt = 5usize;
    let embed = 32usize;
    let heads = 4usize;

    let memory = Tensor::randn(vec![batch, src, embed], Some(21));
    let tgt_in = Tensor::randn(vec![batch, tgt, embed], Some(22));

    let mut dec = DecoderBlock::new(embed, heads, Some(456));
    let out = dec.forward(&tgt_in, &memory, None, None);
    println!("Output shape: {:?}", out.shape().dims());

    let mut opt = Adam::with_learning_rate(0.01);
    let mut params = dec.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/iterators/element_iteration.rs (line 196)

fn demonstrate_gradient_tracking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Gradient Tracking ---");

    // Create a tensor with gradient tracking enabled
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?.with_requires_grad();
    println!("Input tensor (requires_grad): {:?}", tensor.data());

    // Perform element-wise operations through iteration
    let result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Apply a complex transformation: (x^2 + 1) * 2
            elem.pow_scalar(2.0).add_scalar(1.0).mul_scalar(2.0)
        })
        .collect();

    println!("Result tensor: {:?}", result.data());
    println!("Result requires_grad: {}", result.requires_grad());

    // Compute gradients
    let mut loss = result.sum();
    loss.backward(None);

    println!("Loss: {:.6}", loss.value());
    println!("Input gradients: {:?}", tensor.grad().map(|g| g.data()));

    Ok(())
}

examples/neural_networks/basic_transformer.rs (line 171)

    pub fn train_non_autoregressive_steps(
        &mut self,
        src: &Tensor,
        tgt: &Tensor,
        steps: usize,
        lr: f32,
    ) {
        let mut opt = Adam::with_learning_rate(lr);
        {
            let params_once = self.parameters();
            for p in &params_once {
                opt.add_parameter(p);
            }
        }
        for step in 0..steps {
            // forward + backward scope (immutable borrow)
            {
                let pred = self.forward(src, tgt);
                let diff = pred.sub_tensor(tgt);
                let mut loss = diff.pow_scalar(2.0).mean();
                if step == 0 || step + 1 == steps {
                    println!("NAR train step {}: loss={:.6}", step, loss.value());
                }
                loss.backward(None);
            }
            // step + zero_grad scope (mutable borrow)
            let mut params_step = self.parameters();
            opt.step(&mut params_step);
            opt.zero_grad(&mut params_step);
        }
    }

    /// Auto-regressive training (teacher forcing): predict next token with causal mask
    pub fn train_autoregressive_steps(
        &mut self,
        src: &Tensor,
        tgt: &Tensor,
        steps: usize,
        lr: f32,
    ) {
        let mut opt = Adam::with_learning_rate(lr);
        {
            let params_once = self.parameters();
            for p in &params_once {
                opt.add_parameter(p);
            }
        }

        // Build encoder memory once (static dataset demo)
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }

        let (b, t, _e) = Self::triple(tgt);
        // Predict y[t] from y[:t] using causal mask; here we simply predict full seq with mask
        let causal = Self::build_causal_mask_static(b, self.num_heads, t);
        for step in 0..steps {
            // forward + backward scope
            {
                let mut out = tgt.clone();
                for dec in &self.decoders {
                    out = dec.forward(&out, &memory, Some(&causal), None);
                }
                let diff = out.sub_tensor(tgt);
                let mut loss = diff.pow_scalar(2.0).mean();
                if step == 0 || step + 1 == steps {
                    println!("AR  train step {}: loss={:.6}", step, loss.value());
                }
                loss.backward(None);
            }
            let mut params_step = self.parameters();
            opt.step(&mut params_step);
            opt.zero_grad(&mut params_step);
        }
    }

    fn triple(t: &Tensor) -> (usize, usize, usize) {
        let d = t.shape().dims();
        (d[0], d[1], d[2])
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Transformer Example ===");

    let batch = 2usize;
    let src_len = 8usize;
    let tgt_len = 6usize;
    let embed = 32usize;
    let heads = 4usize;
    let layers = 2usize;

    let src = Tensor::randn(vec![batch, src_len, embed], Some(1001));
    let tgt = Tensor::randn(vec![batch, tgt_len, embed], Some(1002));

    let mut trf = BasicTransformer::new(embed, heads, layers, Some(999));
    let out = trf.forward(&src, &tgt);
    println!("Output shape: {:?}", out.shape().dims());

    // Quick optimization step
    let mut opt = Adam::with_learning_rate(0.005);
    let mut params = trf.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());

    // Demo: non auto-regressive inference (single pass)
    let nar = trf.infer_non_autoregressive(&src, tgt_len);
    println!("NAR output shape: {:?}", nar.shape().dims());

    // Demo: auto-regressive inference (toy)
    let ar = trf.infer_autoregressive(&src, 3);
    println!("AR output shape: {:?}", ar.shape().dims());

    // NAR training demo
    let nar_tgt = tgt.clone();
    trf.train_non_autoregressive_steps(&src, &nar_tgt, 3, 0.01);

    // AR training demo (teacher-forced)
    let ar_tgt = tgt.clone();
    trf.train_autoregressive_steps(&src, &ar_tgt, 3, 0.01);
    println!("=== Done ===");
    Ok(())
}

examples/getting_started/serialization_basics.rs (line 138)

fn demonstrate_optimizer_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Optimizer Serialization ---");

    // Create an optimizer with some parameters
    let mut weight = Tensor::randn(vec![2, 2], Some(42)).with_requires_grad();
    let mut bias = Tensor::randn(vec![2], Some(43)).with_requires_grad();

    let config = AdamConfig {
        learning_rate: 0.001,
        beta1: 0.9,
        beta2: 0.999,
        eps: 1e-8,
        weight_decay: 0.0,
        amsgrad: false,
    };

    let mut optimizer = Adam::with_config(config);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    println!(
        "Created optimizer with {} parameters",
        optimizer.parameter_count()
    );
    println!("Learning rate: {}", optimizer.learning_rate());

    // Simulate some training steps
    for _ in 0..3 {
        let mut loss = weight.sum() + bias.sum();
        loss.backward(None);
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);
    }

    // Save optimizer state
    let optimizer_path = "temp_optimizer.json";
    optimizer.save_json(optimizer_path)?;
    println!("Saved optimizer to: {}", optimizer_path);

    // Load optimizer state
    let loaded_optimizer = Adam::load_json(optimizer_path)?;
    println!(
        "Loaded optimizer with {} parameters",
        loaded_optimizer.parameter_count()
    );
    println!("Learning rate: {}", loaded_optimizer.learning_rate());

    // Verify optimizer state
    assert_eq!(
        optimizer.parameter_count(),
        loaded_optimizer.parameter_count()
    );
    assert_eq!(optimizer.learning_rate(), loaded_optimizer.learning_rate());
    println!("Optimizer serialization verification: PASSED");

    Ok(())
}

/// Demonstrate format comparison and performance characteristics
fn demonstrate_format_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Format Comparison ---");

    // Create a larger tensor for comparison
    let tensor = Tensor::randn(vec![10, 10], Some(44));

    // Save in both formats
    tensor.save_json("temp_comparison.json")?;
    tensor.save_binary("temp_comparison.bin")?;

    // Compare file sizes
    let json_size = fs::metadata("temp_comparison.json")?.len();
    let binary_size = fs::metadata("temp_comparison.bin")?.len();

    println!("JSON file size: {} bytes", json_size);
    println!("Binary file size: {} bytes", binary_size);
    println!(
        "Compression ratio: {:.2}x",
        json_size as f64 / binary_size as f64
    );

    // Load and verify both formats
    let json_tensor = Tensor::load_json("temp_comparison.json")?;
    let binary_tensor = Tensor::load_binary("temp_comparison.bin")?;

    assert_eq!(tensor.shape().dims(), json_tensor.shape().dims());
    assert_eq!(tensor.shape().dims(), binary_tensor.shape().dims());
    assert_eq!(tensor.data(), json_tensor.data());
    assert_eq!(tensor.data(), binary_tensor.data());

    println!("Format comparison verification: PASSED");

    Ok(())
}

/// Demonstrate a basic model checkpointing workflow
fn demonstrate_model_checkpointing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Model Checkpointing ---");

    // Create a simple model (weights and bias)
    let mut weights = Tensor::randn(vec![2, 1], Some(45)).with_requires_grad();
    let mut bias = Tensor::randn(vec![1], Some(46)).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weights);
    optimizer.add_parameter(&bias);

    println!("Initial weights: {:?}", weights.data());
    println!("Initial bias: {:?}", bias.data());

    // Simulate training
    for epoch in 0..5 {
        let mut loss = weights.sum() + bias.sum();
        loss.backward(None);
        optimizer.step(&mut [&mut weights, &mut bias]);
        optimizer.zero_grad(&mut [&mut weights, &mut bias]);

        if epoch % 2 == 0 {
            // Save checkpoint
            let checkpoint_dir = format!("checkpoint_epoch_{}", epoch);
            fs::create_dir_all(&checkpoint_dir)?;

            weights.save_json(format!("{}/weights.json", checkpoint_dir))?;
            bias.save_json(format!("{}/bias.json", checkpoint_dir))?;
            optimizer.save_json(format!("{}/optimizer.json", checkpoint_dir))?;

            println!("Saved checkpoint for epoch {}", epoch);
        }
    }

    // Load from checkpoint
    let loaded_weights = Tensor::load_json("checkpoint_epoch_4/weights.json")?;
    let loaded_bias = Tensor::load_json("checkpoint_epoch_4/bias.json")?;
    let loaded_optimizer = Adam::load_json("checkpoint_epoch_4/optimizer.json")?;

    println!("Loaded weights: {:?}", loaded_weights.data());
    println!("Loaded bias: {:?}", loaded_bias.data());
    println!(
        "Loaded optimizer learning rate: {}",
        loaded_optimizer.learning_rate()
    );

    // Verify checkpoint integrity
    assert_eq!(weights.shape().dims(), loaded_weights.shape().dims());
    assert_eq!(bias.shape().dims(), loaded_bias.shape().dims());
    assert_eq!(optimizer.learning_rate(), loaded_optimizer.learning_rate());

    println!("Checkpointing verification: PASSED");

    Ok(())
}

examples/neural_networks/multi_head_attention.rs (line 209)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Multi-Head Attention Example ===");

    let batch = 2usize;
    let src_len = 5usize;
    let tgt_len = 4usize;
    let embed = 16usize;
    let heads = 4usize;

    let query = Tensor::randn(vec![batch, tgt_len, embed], Some(7));
    let key = Tensor::randn(vec![batch, src_len, embed], Some(8));
    let value = Tensor::randn(vec![batch, src_len, embed], Some(9));

    let mut mha = MultiHeadAttention::new(embed, heads, Some(42));

    // Simple causal mask for target self-attention shape [b, h, tq, tk]
    let mut mask = Tensor::zeros(vec![batch, heads, tgt_len, src_len]);
    // Disallow attending to future positions when tgt_len <= src_len by adding -1e9
    // Here, just demonstrate mask broadcast/add mechanics with a light mask on last head
    if src_len >= tgt_len {
        // set upper triangle to a large negative value for head 0
        for i in 0..tgt_len {
            for j in (i + 1)..src_len {
                let idx = [0usize, 0usize, i, j];
                // Quick set via data_mut using a slice view
                let offset = mask.memory_offset(&idx);
                let data = mask.data_mut();
                data[offset] = -1e9;
            }
        }
    }

    let out = mha.forward(&query, &key, &value, Some(&mask));
    println!("Output shape: {:?}", out.shape().dims());

    // Tiny training step to confirm gradients are wired
    let mut optimizer = Adam::with_learning_rate(0.01);
    let mut params = mha.parameters();
    for p in &params {
        optimizer.add_parameter(p);
    }

    // Dummy loss = mean of output
    let mut loss = out.mean();
    loss.backward(None);
    optimizer.step(&mut params);
    optimizer.zero_grad(&mut params);

    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

Additional examples can be found in:

• examples/optimizers/adam_configurations.rs
• examples/optimizers/learning_rate_scheduling.rs
• examples/getting_started/optimizer_basics.rs
• examples/neural_networks/feedforward_network.rs
• examples/RL_training/../neural_networks/basic_linear_layer.rs
• examples/supervised_training/supervised_bce.rs
• examples/iterators/performance_optimization.rs
• examples/supervised_training/supervised_classification.rs
• examples/supervised_training/supervised_regression.rs
• examples/RL_training/dqn.rs
• examples/RL_training/ppo_discrete.rs
• examples/RL_training/ppo_continuous.rs
• examples/RL_training/td3.rs

pub unsafe fn as_ptr(&self) -> *const f32

Returns a raw pointer to the tensor data for unsafe operations

Safety

This is unsafe because it provides direct access to the underlying memory. The caller must ensure:

• The tensor is not dropped while the pointer is used
• No concurrent mutable access occurs
• Bounds are respected

pub unsafe fn as_mut_ptr(&mut self) -> *mut f32

Returns a mutable raw pointer to the tensor data for unsafe operations

Safety

This is unsafe because it provides direct mutable access to the underlying memory. The caller must ensure:

• The tensor is not dropped while the pointer is used
• No concurrent access occurs
• Bounds are respected

pub fn grad_fn(&self) -> &GradFn

Get a reference to the gradient function (for gradtrack)

Returns a reference to the gradient function associated with this tensor. This is used internally by the gradtrack system to compute gradients.

Returns

Reference to the gradient function

Implementation Details

This method is used by the gradtrack engine to access the gradient computation function during backward pass.

pub fn set_grad(&mut self, grad: Tensor)

Set gradient from external source

Sets the gradient tensor for this tensor. This is used internally by the gradtrack system to set gradients during backward pass.

Arguments

• grad - The gradient tensor to set

Implementation Details

This method is used internally by the gradtrack engine to set gradients during backward pass. It only sets the gradient if gradient tracking is enabled for this tensor.

Examples found in repository

examples/RL_training/dqn.rs (line 291)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

examples/RL_training/ppo_continuous.rs (line 307)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

examples/RL_training/ppo_discrete.rs (line 266)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

examples/supervised_training/supervised_bce.rs (line 37)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

examples/supervised_training/supervised_classification.rs (line 37)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

examples/supervised_training/supervised_regression.rs (line 36)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

Additional examples can be found in:

• examples/RL_training/td3.rs

pub fn zero_grad(&mut self)

Clear accumulated gradients for this tensor

This method is used by optimizers to zero gradients before each backward pass. It clears any accumulated gradients, allowing for fresh gradient computation.

Examples

use train_station::Tensor;

let mut tensor = Tensor::ones(vec![2, 3]).with_requires_grad();
tensor.set_grad(Tensor::ones(vec![2, 3]));
assert!(tensor.grad().is_some());
tensor.zero_grad();
assert!(tensor.grad().is_none());

pub fn is_contiguous(&self) -> bool

Checks if the tensor data is stored contiguously in memory

Returns

true if the tensor data is contiguous, enabling optimized SIMD operations

Examples

use train_station::Tensor;

let tensor = Tensor::new(vec![2, 3, 4]);
assert!(tensor.is_contiguous());

Examples found in repository

examples/getting_started/tensor_basics.rs (line 187)

fn demonstrate_utility_functions() {
    println!("\n--- Utility Functions ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

    // Basic properties
    println!("Shape: {:?}", tensor.shape().dims());
    println!("Size: {}", tensor.size());
    println!("Is contiguous: {}", tensor.is_contiguous());
    println!("Device: {:?}", tensor.device());

    // Mathematical operations
    let sum = tensor.sum();
    println!("Sum: {}", sum.value());

    let mean = tensor.mean();
    println!("Mean: {}", mean.value());

    let norm = tensor.norm();
    println!("Norm: {}", norm.value());

    // Device placement
    let cpu_tensor = Tensor::zeros_on_device(vec![3, 3], train_station::Device::cpu());
    println!(
        "CPU tensor: shape {:?}, device: {:?}",
        cpu_tensor.shape().dims(),
        cpu_tensor.device()
    );
}

pub fn strides(&self) -> &[usize]

Gets the memory strides for all dimensions

Returns

Reference to the stride vector for efficient memory access calculations

Examples

use train_station::Tensor;

let tensor = Tensor::new(vec![2, 3, 4]);
assert_eq!(tensor.strides(), &[12, 4, 1]);

pub fn stride(&self, dim: usize) -> usize

Gets the memory stride for a specific dimension

Arguments

• dim - The dimension index

Returns

The memory stride for the given dimension

Panics

Panics if dim is out of bounds

pub fn memory_offset(&self, indices: &[usize]) -> usize

Calculates the linear memory offset for given multi-dimensional indices

Arguments

• indices - Vector of indices for each dimension

Returns

Linear memory offset for direct memory access

Examples

use train_station::Tensor;

let tensor = Tensor::new(vec![2, 3, 4]);
let offset = tensor.memory_offset(&[1, 2, 3]);
// offset = 1*12 + 2*4 + 3*1 = 23

Examples found in repository

examples/neural_networks/basic_transformer.rs (line 95)

    pub fn infer_autoregressive(&self, src: &Tensor, max_steps: usize) -> Tensor {
        let (b, _s, e) = Self::triple(src);
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }

        let mut out_seq: Vec<Tensor> = Vec::new();
        // Start token: zeros
        let mut current = Tensor::zeros(vec![b, 1, e]);
        for _step in 0..max_steps {
            // Build causal mask for length t
            let t = current.shape().dims()[1];
            let mut causal = Tensor::ones(vec![b, self.num_heads, t, t]);
            // Upper triangle as false -> masked for all batches and heads
            for bb in 0..b {
                for hh in 0..self.num_heads {
                    for i in 0..t {
                        for j in (i + 1)..t {
                            let offset = causal.memory_offset(&[bb, hh, i, j]);
                            let data = causal.data_mut();
                            data[offset] = 0.0;
                        }
                    }
                }
            }
            let mut step_out = current.clone();
            for dec in &self.decoders {
                step_out = dec.forward(&step_out, &memory, Some(&causal), None);
            }
            // (Toy) append placeholder token; real models would project last token
            out_seq.push(step_out.clone());
            // Append a zero token to grow sequence by 1 for next causal computation
            current = Tensor::zeros(vec![b, t + 1, e]);
        }
        // Simple return of final sequence placeholder
        current
    }

    /// Non auto-regressive inference: single forward pass
    pub fn infer_non_autoregressive(&self, src: &Tensor, tgt_len: usize) -> Tensor {
        let (b, _s, e) = Self::triple(src);
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }
        let tgt = Tensor::zeros(vec![b, tgt_len, e]);
        let mut out = tgt.clone();
        for dec in &self.decoders {
            out = dec.forward(&out, &memory, None, None);
        }
        out
    }

    /// Helper: build boolean-like causal mask [b, heads, t, t] with 1.0 keep, 0.0 masked
    fn build_causal_mask_static(batch: usize, heads: usize, t: usize) -> Tensor {
        let mut mask = Tensor::ones(vec![batch, heads, t, t]);
        for bb in 0..batch {
            for hh in 0..heads {
                for i in 0..t {
                    for j in (i + 1)..t {
                        let offset = mask.memory_offset(&[bb, hh, i, j]);
                        let data = mask.data_mut();
                        data[offset] = 0.0;
                    }
                }
            }
        }
        mask
    }

examples/neural_networks/multi_head_attention.rs (line 190)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Multi-Head Attention Example ===");

    let batch = 2usize;
    let src_len = 5usize;
    let tgt_len = 4usize;
    let embed = 16usize;
    let heads = 4usize;

    let query = Tensor::randn(vec![batch, tgt_len, embed], Some(7));
    let key = Tensor::randn(vec![batch, src_len, embed], Some(8));
    let value = Tensor::randn(vec![batch, src_len, embed], Some(9));

    let mut mha = MultiHeadAttention::new(embed, heads, Some(42));

    // Simple causal mask for target self-attention shape [b, h, tq, tk]
    let mut mask = Tensor::zeros(vec![batch, heads, tgt_len, src_len]);
    // Disallow attending to future positions when tgt_len <= src_len by adding -1e9
    // Here, just demonstrate mask broadcast/add mechanics with a light mask on last head
    if src_len >= tgt_len {
        // set upper triangle to a large negative value for head 0
        for i in 0..tgt_len {
            for j in (i + 1)..src_len {
                let idx = [0usize, 0usize, i, j];
                // Quick set via data_mut using a slice view
                let offset = mask.memory_offset(&idx);
                let data = mask.data_mut();
                data[offset] = -1e9;
            }
        }
    }

    let out = mha.forward(&query, &key, &value, Some(&mask));
    println!("Output shape: {:?}", out.shape().dims());

    // Tiny training step to confirm gradients are wired
    let mut optimizer = Adam::with_learning_rate(0.01);
    let mut params = mha.parameters();
    for p in &params {
        optimizer.add_parameter(p);
    }

    // Dummy loss = mean of output
    let mut loss = out.mean();
    loss.backward(None);
    optimizer.step(&mut params);
    optimizer.zero_grad(&mut params);

    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

pub fn memory_alignment(&self) -> usize

Gets the memory alignment of the tensor data

Returns

The memory alignment in bytes (typically 32 for SIMD optimization)

pub fn is_broadcastable_with(&self, other: &Tensor) -> bool

Checks if this tensor is broadcastable with another tensor

Arguments

• other - The other tensor to check broadcasting compatibility

Returns

true if the tensors are broadcastable according to NumPy broadcasting rules

Examples

use train_station::Tensor;

let a = Tensor::new(vec![2, 3, 4]);
let b = Tensor::new(vec![1, 3, 4]);
assert!(a.is_broadcastable_with(&b));

pub fn memory_footprint(&self) -> usize

Gets the total number of bytes allocated for this tensor

Returns

Total memory footprint in bytes

pub fn get(&self, indices: &[usize]) -> f32

Get a single element from the tensor at the specified indices

Arguments

• indices - Multi-dimensional indices to access the element

Returns

The value at the specified position

Panics

Panics if indices are out of bounds or indices length doesn’t match tensor rank

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let value = tensor.get(&[0, 1]);
assert_eq!(value, 2.0);

Examples found in repository

examples/getting_started/tensor_basics.rs (line 162)

fn demonstrate_data_access() {
    println!("\n--- Data Access ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

    // Access individual elements
    println!("Element [0, 0]: {}", tensor.get(&[0, 0]));
    println!("Element [0, 1]: {}", tensor.get(&[0, 1]));
    println!("Element [1, 0]: {}", tensor.get(&[1, 0]));
    println!("Element [1, 1]: {}", tensor.get(&[1, 1]));

    // Access data as slice
    let data = tensor.data();
    println!("Data as slice: {:?}", data);

    // Iterate over elements
    println!("Elements:");
    for (i, &value) in data.iter().enumerate() {
        println!("  [{}]: {}", i, value);
    }
}

pub fn set(&mut self, indices: &[usize], value: f32)

Set a single element in the tensor at the specified indices

Arguments

• indices - Multi-dimensional indices to set the element
• value - The value to set

Panics

Panics if indices are out of bounds or indices length doesn’t match tensor rank

Examples

use train_station::Tensor;

let mut tensor = Tensor::new(vec![2, 2]);
tensor.set(&[0, 1], 42.0);
assert_eq!(tensor.get(&[0, 1]), 42.0);

pub fn data(&self) -> &[f32]

Returns a safe slice of the tensor’s underlying data

Provides safe access to the tensor’s data without requiring unsafe pointer operations. This is the preferred way to access tensor data for reading values, comparisons, and other operations that don’t require direct pointer manipulation.

Returns

A slice containing all tensor elements in row-major order

Performance

• Zero-Cost: Direct slice creation with no copying
• Cache-Friendly: Sequential memory access pattern
• Safe: No unsafe code required for basic data access

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let data = tensor.data();

// Safe indexing and comparisons
assert_eq!(data[0], 1.0);
assert_eq!(data.len(), tensor.size());

Examples found in repository

examples/RL_training/dqn.rs (line 111)

    fn copy_from(&mut self, other: &Self) {
        for (t, s) in self.layers.iter_mut().zip(other.layers.iter()) {
            {
                let src = s.weight.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = s.bias.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }
}

// -------------------------------
// Q-Network (state -> Q-values over actions)
// -------------------------------

struct QNet {
    net: Mlp,
}

impl QNet {
    fn new(state_dim: usize, action_dim: usize, seed: Option<u64>) -> Self {
        let net = Mlp::new(&[state_dim, 64, 64, action_dim], seed);
        Self { net }
    }
    fn forward(&self, state: &Tensor) -> Tensor {
        self.net.forward(state, None)
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.net.parameters()
    }
    fn set_requires_grad_all(&mut self, enable: bool) {
        self.net.set_requires_grad_all(enable);
    }
}

// -------------------------------
// Discrete YardEnv (3 actions: -1, 0, +1)
// -------------------------------

struct YardEnv {
    pos: f32,
    vel: f32,
    steps: usize,
    max_steps: usize,
    rng: SmallRng,
}

impl YardEnv {
    const ACTIONS: [f32; 3] = [-1.0, 0.0, 1.0];

    fn new(seed: u64) -> Self {
        let mut env = Self {
            pos: 0.0,
            vel: 0.0,
            steps: 0,
            max_steps: 200,
            rng: SmallRng::new(seed),
        };
        env.reset();
        env
    }

    fn reset(&mut self) -> Tensor {
        self.pos = self.rng.uniform(-0.5, 0.5);
        self.vel = self.rng.uniform(-0.1, 0.1);
        self.steps = 0;
        self.state_tensor()
    }

    fn state_tensor(&self) -> Tensor {
        Tensor::from_slice(&[self.pos, self.vel, 0.0], vec![1, 3]).unwrap()
    }

    fn step(&mut self, action_index: usize) -> (Tensor, f32, bool) {
        let a = Self::ACTIONS[action_index.min(2)];
        self.vel += 0.1 * a - 0.01 * self.pos;
        self.pos += self.vel;
        self.steps += 1;
        let reward = -(self.pos * self.pos) - 0.05 * (a * a);
        let done = self.pos.abs() > 3.0 || self.steps >= self.max_steps;
        (self.state_tensor(), reward, done)
    }
}

// -------------------------------
// Replay Buffer
// -------------------------------

struct ReplayBuffer {
    capacity: usize,
    size: usize,
    pos: usize,
    state_dim: usize,
    states: Vec<f32>,
    actions: Vec<usize>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    next_states: Vec<f32>,
}

impl ReplayBuffer {
    fn new(capacity: usize, state_dim: usize) -> Self {
        Self {
            capacity,
            size: 0,
            pos: 0,
            state_dim,
            states: vec![0.0; capacity * state_dim],
            actions: vec![0usize; capacity],
            rewards: vec![0.0; capacity],
            dones: vec![0.0; capacity],
            next_states: vec![0.0; capacity * state_dim],
        }
    }

    fn push(&mut self, s: &[f32], a_idx: usize, r: f32, d: f32, s2: &[f32]) {
        let i = self.pos;
        let so = i * self.state_dim;
        self.states[so..so + self.state_dim].copy_from_slice(s);
        self.actions[i] = a_idx;
        self.rewards[i] = r;
        self.dones[i] = d;
        self.next_states[so..so + self.state_dim].copy_from_slice(s2);
        self.pos = (self.pos + 1) % self.capacity;
        self.size = self.size.saturating_add(1).min(self.capacity);
    }

    fn can_sample(&self, batch_size: usize) -> bool {
        self.size >= batch_size
    }

    fn sample(
        &self,
        batch_size: usize,
        rng: &mut SmallRng,
    ) -> (Tensor, Vec<usize>, Tensor, Tensor, Tensor) {
        let mut s_vec = Vec::with_capacity(batch_size * self.state_dim);
        let mut a_idx = Vec::with_capacity(batch_size);
        let mut r_vec = Vec::with_capacity(batch_size);
        let mut d_vec = Vec::with_capacity(batch_size);
        let mut s2_vec = Vec::with_capacity(batch_size * self.state_dim);
        for _ in 0..batch_size {
            let idx = rng.sample_index(self.size);
            let so = idx * self.state_dim;
            s_vec.extend_from_slice(&self.states[so..so + self.state_dim]);
            a_idx.push(self.actions[idx]);
            r_vec.push(self.rewards[idx]);
            d_vec.push(self.dones[idx]);
            s2_vec.extend_from_slice(&self.next_states[so..so + self.state_dim]);
        }
        let s = Tensor::from_slice(&s_vec, vec![batch_size, self.state_dim]).unwrap();
        let r = Tensor::from_slice(&r_vec, vec![batch_size, 1]).unwrap();
        let d = Tensor::from_slice(&d_vec, vec![batch_size, 1]).unwrap();
        let s2 = Tensor::from_slice(&s2_vec, vec![batch_size, self.state_dim]).unwrap();
        (s, a_idx, r, d, s2)
    }
}

// -------------------------------
// Helpers
// -------------------------------

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

fn params_l2_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let _ng = NoGradTrack::new();
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        for &v in p.data() {
            total_sq += v * v;
        }
    }
    total_sq.sqrt()
}

// Pseudo-Huber loss: sqrt(1 + diff^2) - 1 (smooth, robust)
fn pseudo_huber_mean(diff: &Tensor) -> Tensor {
    diff.pow_scalar(2.0)
        .add_scalar(1.0)
        .sqrt()
        .sub_scalar(1.0)
        .mean()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== DQN Example (YardEnv discrete) ===");

    // Dims
    let state_dim = 3usize;
    let action_dim = 3usize;

    // Hparams
    let gamma = 0.99f32;
    let batch_size = 64usize;
    let start_steps = 200usize;
    let target_update_interval = 200usize; // hard update cadence
    let max_grad_norm = 1.0f32;
    let mut epsilon = 1.0f32;
    let eps_min = 0.05f32;
    let eps_decay_steps = 2_000usize; // linear decay
    let total_steps = std::env::var("DQN_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3000usize);

    // Models
    let mut q_net = QNet::new(state_dim, action_dim, Some(7));
    let mut q_targ = QNet::new(state_dim, action_dim, Some(8));
    q_targ.net.copy_from(&q_net.net);
    q_targ.set_requires_grad_all(false);

    // Optimizer
    let mut q_opt = Adam::with_learning_rate(3e-4);
    for p in q_net.parameters() {
        q_opt.add_parameter(p);
    }

    // Replay + env
    let mut rb = ReplayBuffer::new(100_000, state_dim);
    let mut env = YardEnv::new(12345);
    let mut rng = SmallRng::new(999_111);

    // Metrics
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    for t in 0..total_steps {
        // Epsilon-greedy action
        let action_index = if t < start_steps || rng.next_f32() < epsilon {
            rng.sample_index(action_dim)
        } else {
            let _ng = NoGradTrack::new();
            let q_vals = q_net.forward(&state);
            let row = q_vals.data();
            let mut best_i = 0usize;
            let mut best_v = row[0];
            for (i, &r) in row.iter().enumerate().take(action_dim).skip(1) {
                if r > best_v {
                    best_v = r;
                    best_i = i;
                }
            }
            best_i
        };

        // Env step
        let (next_state, reward, done) = env.step(action_index);
        episode_return += reward;

        // Store
        let s_slice = state.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            action_index,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        // Reset on done
        state = if done {
            let st = env.reset();
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Epsilon linear decay
        if t < eps_decay_steps {
            epsilon = (1.0 - (t as f32) / (eps_decay_steps as f32)) * (1.0 - eps_min) + eps_min;
        }

        // Train
        if rb.can_sample(batch_size) {
            let (s, a_idx, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Double DQN target: a* = argmax_a Q_online(s2,a); y = r + (1-d)*gamma*Q_target(s2, a*)
            let target_q = {
                let _ng = NoGradTrack::new();
                let q_online_s2 = q_net.forward(&s2);
                // argmax per row (manual on CPU)
                let row_stride = action_dim;
                let qd = q_online_s2.data();
                let mut next_actions: Vec<usize> = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let base = i * row_stride;
                    let mut bi = 0usize;
                    let mut bv = qd[base];
                    for j in 1..action_dim {
                        let v = qd[base + j];
                        if v > bv {
                            bv = v;
                            bi = j;
                        }
                    }
                    next_actions.push(bi);
                }
                let q_targ_s2 = q_targ.forward(&s2);
                let q_targ_g = q_targ_s2.gather(1, &next_actions, &[batch_size, 1]);
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&q_targ_g))
            };

            // Q(s,a) for current actions
            // Zero grads first
            {
                let mut params = q_net.parameters();
                q_opt.zero_grad(&mut params);
            }

            let q_all = q_net.forward(&s);
            let q_sa = q_all.gather(1, &a_idx, &[batch_size, 1]);
            let diff = q_sa.sub_tensor(&target_q);
            let mut loss = pseudo_huber_mean(&diff);
            loss.backward(None);

            // Step (filter only params with grads)
            {
                let params = q_net.parameters();
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    let gn = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    q_opt.step(&mut with_grads);
                    q_opt.zero_grad(&mut with_grads);
                    if t % 100 == 0 {
                        let mut pn = q_net.parameters();
                        let pn_l2 = params_l2_norm(&mut pn);
                        let q_mean = q_all.mean().value();
                        println!(
                            "t={:5} | loss={:.4} | q_mean={:.3} | grad_norm={:.3} | param_norm={:.3} | eps={:.3}",
                            t, loss.value(), q_mean, gn, pn_l2, epsilon
                        );
                    }
                }
            }

            // Target hard update
            if t % target_update_interval == 0 {
                q_targ.net.copy_from(&q_net.net);
            }

            // Clear graphs
            clear_all_graphs_known();
        }
    }

    println!("=== DQN training finished ===");
    Ok(())
}

examples/RL_training/td3.rs (line 117)

    fn copy_from(&mut self, other: &Self) {
        for (t, s) in self.layers.iter_mut().zip(other.layers.iter()) {
            {
                let src = s.weight.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = s.bias.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }

    fn soft_update_from(&mut self, source: &Self, tau: f32) {
        let _ng = NoGradTrack::new();
        for (t, s) in self.layers.iter_mut().zip(source.layers.iter()) {
            // In-place Polyak update to preserve tensor IDs (no optimizer relink needed)
            let new_w = t
                .weight
                .mul_scalar(1.0 - tau)
                .add_tensor(&s.weight.mul_scalar(tau));
            let new_b = t
                .bias
                .mul_scalar(1.0 - tau)
                .add_tensor(&s.bias.mul_scalar(tau));
            {
                let src = new_w.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = new_b.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }
}

// -------------------------------
// Actor and Critic
// -------------------------------

struct Actor {
    net: Mlp,
}

impl Actor {
    fn new(state_dim: usize, action_dim: usize, seed: Option<u64>) -> Self {
        // Smaller net for faster demo: sd -> 64 -> 64 -> ad, tanh output
        let net = Mlp::new(&[state_dim, 64, 64, action_dim], seed);
        Self { net }
    }
    fn forward(&self, state: &Tensor) -> Tensor {
        self.net.forward(state, Some(tanh_bounded))
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.net.parameters()
    }
    fn set_requires_grad_all(&mut self, enable: bool) {
        self.net.set_requires_grad_all(enable);
    }
}

struct Critic {
    net: Mlp,
}

impl Critic {
    fn new(state_dim: usize, action_dim: usize, seed: Option<u64>) -> Self {
        let net = Mlp::new(&[state_dim + action_dim, 64, 64, 1], seed);
        Self { net }
    }
    fn forward(&self, state: &Tensor, action: &Tensor) -> Tensor {
        // Concatenate along feature dim (dim=1) for batched inputs
        // IMPORTANT: use views to preserve gradient graph; cloning would detach autograd
        let s_view = state.view(state.shape().dims().iter().map(|&d| d as i32).collect());
        let a_view = action.view(action.shape().dims().iter().map(|&d| d as i32).collect());
        let sa = Tensor::cat(&[s_view, a_view], 1);
        self.net.forward(&sa, None)
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.net.parameters()
    }
    fn set_requires_grad_all(&mut self, enable: bool) {
        self.net.set_requires_grad_all(enable);
    }
}

// -------------------------------
// Simple continuous control environment: YardEnv
// State: normalized features [pos/3, clamp(vel/1, -1..1), bias(=0)] ; Action: scalar in [-1, 1]
// Dynamics: vel += 0.1*act - 0.01*pos; pos += vel
// Reward: -(pos^2) - 0.1*act^2 ; Episode ends if |pos| > 3 or step >= max_steps
// -------------------------------

struct YardEnv {
    pos: f32,
    vel: f32,
    steps: usize,
    max_steps: usize,
    rng: SmallRng,
}

impl YardEnv {
    fn new(seed: u64) -> Self {
        let mut env = Self {
            pos: 0.0,
            vel: 0.0,
            steps: 0,
            max_steps: 200,
            rng: SmallRng::new(seed),
        };
        env.reset();
        env
    }

    fn reset(&mut self) -> Tensor {
        self.pos = self.rng.uniform(-0.5, 0.5);
        self.vel = self.rng.uniform(-0.1, 0.1);
        self.steps = 0;
        self.state_tensor()
    }

    fn state_tensor(&self) -> Tensor {
        // Normalize to keep critic inputs bounded:
        // - Position is bounded by termination at |pos|>3 → scale by 3 to [-1,1]
        // - Velocity scaled by 1.0 and clamped to [-1,1]
        let pos_n = self.pos / 3.0;
        let vel_n = self.vel.clamp(-1.0, 1.0);
        Tensor::from_slice(&[pos_n, vel_n, 0.0], vec![1, 3]).unwrap()
    }

    fn step(&mut self, action_value: f32) -> (Tensor, f32, bool) {
        let a = action_value.clamp(-1.0, 1.0);
        self.vel += 0.1 * a - 0.01 * self.pos;
        self.pos += self.vel;
        self.steps += 1;

        let reward = -(self.pos * self.pos) - 0.1 * (a * a);
        let done = self.pos.abs() > 3.0 || self.steps >= self.max_steps;
        (self.state_tensor(), reward, done)
    }
}

// -------------------------------
// Replay Buffer
// -------------------------------

struct ReplayBuffer {
    capacity: usize,
    size: usize,
    pos: usize,
    state_dim: usize,
    action_dim: usize,
    states: Vec<f32>,
    actions: Vec<f32>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    next_states: Vec<f32>,
}

impl ReplayBuffer {
    fn new(capacity: usize, state_dim: usize, action_dim: usize) -> Self {
        Self {
            capacity,
            size: 0,
            pos: 0,
            state_dim,
            action_dim,
            states: vec![0.0; capacity * state_dim],
            actions: vec![0.0; capacity * action_dim],
            rewards: vec![0.0; capacity],
            dones: vec![0.0; capacity],
            next_states: vec![0.0; capacity * state_dim],
        }
    }

    fn push(&mut self, s: &[f32], a: &[f32], r: f32, d: f32, s2: &[f32]) {
        let i = self.pos;
        let so = i * self.state_dim;
        let ao = i * self.action_dim;
        self.states[so..so + self.state_dim].copy_from_slice(s);
        self.actions[ao..ao + self.action_dim].copy_from_slice(a);
        self.rewards[i] = r;
        self.dones[i] = d;
        self.next_states[so..so + self.state_dim].copy_from_slice(s2);

        self.pos = (self.pos + 1) % self.capacity;
        self.size = self.size.saturating_add(1).min(self.capacity);
    }

    fn can_sample(&self, batch_size: usize) -> bool {
        self.size >= batch_size
    }

    fn sample(
        &self,
        batch_size: usize,
        rng: &mut SmallRng,
    ) -> (Tensor, Tensor, Tensor, Tensor, Tensor) {
        let mut s_vec = Vec::with_capacity(batch_size * self.state_dim);
        let mut a_vec = Vec::with_capacity(batch_size * self.action_dim);
        let mut r_vec = Vec::with_capacity(batch_size);
        let mut d_vec = Vec::with_capacity(batch_size);
        let mut s2_vec = Vec::with_capacity(batch_size * self.state_dim);

        for _ in 0..batch_size {
            let idx = rng.sample_index(self.size);
            let so = idx * self.state_dim;
            let ao = idx * self.action_dim;
            s_vec.extend_from_slice(&self.states[so..so + self.state_dim]);
            a_vec.extend_from_slice(&self.actions[ao..ao + self.action_dim]);
            r_vec.push(self.rewards[idx]);
            d_vec.push(self.dones[idx]);
            s2_vec.extend_from_slice(&self.next_states[so..so + self.state_dim]);
        }

        let s = Tensor::from_slice(&s_vec, vec![batch_size, self.state_dim]).unwrap();
        let a = Tensor::from_slice(&a_vec, vec![batch_size, self.action_dim]).unwrap();
        let r = Tensor::from_slice(&r_vec, vec![batch_size, 1]).unwrap();
        let d = Tensor::from_slice(&d_vec, vec![batch_size, 1]).unwrap();
        let s2 = Tensor::from_slice(&s2_vec, vec![batch_size, self.state_dim]).unwrap();
        (s, a, r, d, s2)
    }
}

// -------------------------------
// Helper: gradient clipping by global norm
// -------------------------------

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    // Compute global L2 norm of all grads
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                let scaled = g.mul_scalar(scale);
                p.set_grad(scaled);
            }
        }
    }
}

// Compute global L2 norm of gradients across a parameter list (read-only)
fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// Compute L2 norm of parameters (weights/biases) across a parameter list
fn params_l2_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let _ng = NoGradTrack::new();
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        for &v in p.data() {
            total_sq += v * v;
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main: TD3 training on YardEnv
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== TD3 Example (YardEnv) ===");

    // Environment / problem dims
    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hyperparameters (small for demo)
    let gamma = 0.99f32;
    let tau = 0.005f32; // Polyak
    let policy_noise = 0.2f32; // target smoothing noise stddev
    let exploration_noise = 0.1f32; // behavior policy noise stddev
    let policy_delay = 2usize;
    let batch_size = 64usize;
    let start_steps = 500usize; // random exploration steps
    let total_steps = 1500usize;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(11));
    let mut actor_targ = Actor::new(state_dim, action_dim, Some(12));
    actor_targ.net.copy_from(&actor.net);
    actor_targ.set_requires_grad_all(false);

    let mut critic1 = Critic::new(state_dim, action_dim, Some(21));
    let mut critic2 = Critic::new(state_dim, action_dim, Some(22));
    let mut critic1_targ = Critic::new(state_dim, action_dim, Some(23));
    let mut critic2_targ = Critic::new(state_dim, action_dim, Some(24));
    critic1_targ.net.copy_from(&critic1.net);
    critic2_targ.net.copy_from(&critic2.net);
    critic1_targ.set_requires_grad_all(false);
    critic2_targ.set_requires_grad_all(false);

    // Optimizers
    let mut actor_opt = Adam::with_learning_rate(1e-3);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }

    let mut critic_opt = Adam::with_learning_rate(1e-4);
    for p in critic1.parameters() {
        critic_opt.add_parameter(p);
    }
    for p in critic2.parameters() {
        critic_opt.add_parameter(p);
    }

    // Replay buffer and env
    let mut rb = ReplayBuffer::new(100_000, state_dim, action_dim);
    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(987654321);

    // Reset & metric trackers
    let mut state = env.reset(); // [1, state_dim]
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32; // smooth short-term
    let mut best_return = f32::NEG_INFINITY;
    let mut policy_updates: usize = 0;

    for t in 0..total_steps {
        // Select action
        let action_tensor = if t < start_steps {
            let a = rng.uniform(-1.0, 1.0);
            Tensor::from_slice(&[a], vec![1, action_dim]).unwrap()
        } else {
            // Behavior policy with exploration noise
            let _ng = NoGradTrack::new();
            let det = actor.forward(&state);
            let noise = Tensor::randn(vec![1, action_dim], None).mul_scalar(exploration_noise);
            tanh_bounded(&det.add_tensor(&noise))
        };
        let action_value = action_tensor.data()[0];

        // Environment step
        let (next_state, reward, done) = env.step(action_value);
        episode_return += reward;

        // Store transition
        let s_slice = state.data().to_vec();
        let a_slice = action_tensor.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            &a_slice,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        state = if done {
            let st = env.reset();
            // Metrics: update EMA and best
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={} | policy_updates={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size,
                policy_updates
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Training
        if rb.can_sample(batch_size) {
            // Sample batch
            let (s, a, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Compute target values y = r + (1-d)*gamma*min(Q1', Q2') using target networks (no grad)
            let target_q = {
                let _ng = NoGradTrack::new();
                // Target actions with smoothing noise (tanh bounds)
                let noise =
                    Tensor::randn(vec![batch_size, action_dim], None).mul_scalar(policy_noise);
                let a_targ = tanh_bounded(&actor_targ.forward(&s2).add_tensor(&noise));
                let q1_t = critic1_targ.forward(&s2, &a_targ);
                let q2_t = critic2_targ.forward(&s2, &a_targ);

                // Elementwise min via data() since this path is no-grad
                let q1d = q1_t.data();
                let q2d = q2_t.data();
                let mut min_vec = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let v1 = q1d[i];
                    let v2 = q2d[i];
                    min_vec.push(v1.min(v2));
                }
                let min_q = Tensor::from_slice(&min_vec, vec![batch_size, 1]).unwrap();
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&min_q))
            };

            // Critic update (both critics)
            // Zero grads in a short scope, then drop borrows before forward
            {
                let mut params = {
                    let c_params = critic1.parameters();
                    let c2_params = critic2.parameters();
                    let mut tmp: Vec<&mut Tensor> = Vec::new();
                    tmp.extend(c_params);
                    tmp.extend(c2_params);
                    tmp
                };
                critic_opt.zero_grad(&mut params);
            }

            // Forward current Q estimates
            let q1 = critic1.forward(&s, &a);
            let q2 = critic2.forward(&s, &a);
            let diff1 = q1.sub_tensor(&target_q);
            let diff2 = q2.sub_tensor(&target_q);
            let mut critic_loss = diff1
                .pow_scalar(2.0)
                .mean()
                .add_tensor(&diff2.pow_scalar(2.0).mean());

            // Backward
            critic_loss.backward(None);

            // Optional gradient clipping + step (only for params that received grads)
            {
                let params = {
                    let c_params = critic1.parameters();
                    let c2_params = critic2.parameters();
                    let mut tmp: Vec<&mut Tensor> = Vec::new();
                    tmp.extend(c_params);
                    tmp.extend(c2_params);
                    tmp
                };
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    // Pre-step metrics
                    let grad_norm_before = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    critic_opt.step(&mut with_grads);
                    critic_opt.zero_grad(&mut with_grads);

                    // Post-step metrics (param norm)
                    let mut for_norm_params = {
                        let c_params = critic1.parameters();
                        let c2_params = critic2.parameters();
                        let mut tmp: Vec<&mut Tensor> = Vec::new();
                        tmp.extend(c_params);
                        tmp.extend(c2_params);
                        tmp
                    };
                    let param_norm = params_l2_norm(&mut for_norm_params);

                    // Print compact critic metrics occasionally
                    if t % 100 == 0 {
                        let q1_mean = q1.mean().value();
                        let q2_mean = q2.mean().value();
                        let tq_mean = target_q.mean().value();
                        println!(
                            "t={:5} | critic_loss={:.4} | q1_mean={:.3} q2_mean={:.3} tq_mean={:.3} | grad_norm={:.3} | crit_param_norm={:.3}",
                            t,
                            critic_loss.value(),
                            q1_mean,
                            q2_mean,
                            tq_mean,
                            grad_norm_before,
                            param_norm
                        );
                    }
                }
            }

            // Delayed policy update
            if t % policy_delay == 0 {
                // Actor update: maximize Q1(s, actor(s)) -> minimize -Q1
                // Zero actor grads before backward
                {
                    let mut a_params: Vec<&mut Tensor> = actor.parameters();
                    actor_opt.zero_grad(&mut a_params);
                }

                let a_pred = actor.forward(&s);
                let q_for_actor = critic1.forward(&s, &a_pred);
                let mut actor_loss = q_for_actor.mul_scalar(-1.0).mean();
                actor_loss.backward(None);

                {
                    let a_params: Vec<&mut Tensor> = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in a_params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let grad_norm_before = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);

                        // Post-step param norm
                        let mut for_norm_params = actor.parameters();
                        let param_norm = params_l2_norm(&mut for_norm_params);

                        policy_updates += 1;
                        if t % 200 == 0 {
                            println!(
                                "t={:5} | actor_loss={:.4} | act_grad_norm={:.3} | act_param_norm={:.3} | lr_a={:.4e} lr_c={:.4e} | policy_updates={}",
                                t,
                                actor_loss.value(),
                                grad_norm_before,
                                param_norm,
                                actor_opt.learning_rate(),
                                critic_opt.learning_rate(),
                                policy_updates
                            );
                        }
                    }
                }

                // Target updates (Polyak averaging, no grad)
                actor_targ.net.soft_update_from(&actor.net, tau);
                critic1_targ.net.soft_update_from(&critic1.net, tau);
                critic2_targ.net.soft_update_from(&critic2.net, tau);
            }

            // Clear entire graphs to avoid stale accumulation across iterations
            clear_all_graphs_known();
        }
    }

    println!("=== TD3 training finished ===");
    Ok(())
}

examples/RL_training/ppo_continuous.rs (line 297)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

examples/RL_training/ppo_discrete.rs (line 256)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

// log-softmax for selected actions: given logits [B,A] and actions Vec<usize> -> log_prob [B,1]
fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

// probability ratio = exp(new_logp - old_logp)
fn ratio_from_logps(new_logp: &Tensor, old_logp: &Tensor) -> Tensor {
    new_logp.sub_tensor(old_logp).exp()
}

// Clamp ratio to [1-clip, 1+clip] using ReLU-based clamp (no custom ops)
fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/supervised_training/supervised_bce.rs (line 27)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn accuracy(pred: &Tensor, targets: &Tensor) -> f32 {
    // pred: [B,1] with sigmoid; threshold at 0.5
    let p = pred.data();
    let t = targets.data();
    let mut correct = 0usize;
    for i in 0..p.len() {
        let yhat = if p[i] >= 0.5 { 1.0 } else { 0.0 };
        if (yhat - t[i]).abs() < 1e-6 {
            correct += 1;
        }
    }
    correct as f32 / (p.len() as f32)
}

// Numerically stable BCE with logits:
// L = mean( relu(z) - z*y + log(1 + exp(-|z|)) )
fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Supervised FFN Example (XOR) ===");

    // Dataset: XOR (repeat to form a small batch)
    let inputs: Vec<f32> = vec![
        0.0, 0.0, // -> 0
        0.0, 1.0, // -> 1
        1.0, 0.0, // -> 1
        1.0, 1.0, // -> 0
    ];
    let targets: Vec<f32> = vec![0.0, 1.0, 1.0, 0.0];

    // Repeat the base patterns to stabilize training
    let repeats = 64usize; // effective batch = 4 * repeats = 256
    let mut xs = Vec::with_capacity(repeats * inputs.len());
    let mut ys = Vec::with_capacity(repeats * targets.len());
    for _ in 0..repeats {
        xs.extend_from_slice(&inputs);
        ys.extend_from_slice(&targets);
    }

    let batch = xs.len() / 2; // two features
    let x_t = Tensor::from_slice(&xs, vec![batch, 2]).unwrap();
    let y_t = Tensor::from_slice(&ys, vec![batch, 1]).unwrap();

    // Model config: 2 -> 32 -> 32 -> 1, final sigmoid via loss path
    let cfg = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![32, 32],
        output_size: 1,
        use_bias: true,
    };
    let mut net = FeedForwardNetwork::new(cfg, Some(777));

    // Optimizer and parameter linking
    let mut opt = Adam::with_learning_rate(1e-3);
    for p in net.parameters() {
        opt.add_parameter(p);
    }

    let epochs = 1000usize;
    let max_grad_norm = 1.0f32;
    let mut best_loss = f32::INFINITY;
    let mut best_acc = 0.0f32;

    for e in 0..epochs {
        // Zero grads each iteration
        {
            let mut params = net.parameters();
            opt.zero_grad(&mut params);
        }

        // Forward -> logits; use numerically stable BCE-with-logits for loss
        let logits = net.forward(&x_t);
        let mut loss = bce_with_logits(&logits, &y_t);
        loss.backward(None);

        // Step only params with grads
        {
            let params = net.parameters();
            let mut with_grads: Vec<&mut Tensor> = Vec::new();
            for p in params {
                if p.grad_owned().is_some() {
                    with_grads.push(p);
                }
            }
            if !with_grads.is_empty() {
                clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                opt.step(&mut with_grads);
                opt.zero_grad(&mut with_grads);
            }
        }

        // Metrics (use sigmoid only for reporting accuracy)
        let preds = logits.sigmoid();
        let acc = accuracy(&preds, &y_t);
        if loss.value() < best_loss {
            best_loss = loss.value();
        }
        if acc > best_acc {
            best_acc = acc;
        }
        if e % 10 == 0 || e + 1 == epochs {
            println!(
                "epoch {:4} | loss={:.5} acc={:.3} | best_loss={:.5} best_acc={:.3}",
                e,
                loss.value(),
                acc,
                best_loss,
                best_acc
            );
        }

        // Clear graphs to avoid stale accumulation across epochs
        clear_all_graphs_known();
    }

    // Quick sanity check predictions
    let test = Tensor::from_slice(&inputs, vec![4, 2]).unwrap();
    let out = net.forward(&test).sigmoid();
    println!("predictions (approx): {:?}", out.data());

    println!("=== Supervised training finished ===");
    Ok(())
}

examples/supervised_training/supervised_classification.rs (line 27)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

// Cross-entropy over logits: CE = -mean(log_softmax(logits)[range, labels])
fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

fn accuracy_from_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    num_classes: usize,
) -> f32 {
    let row = logits.data();
    let mut correct = 0usize;
    for (i, &label) in labels.iter().enumerate().take(batch) {
        let base = i * num_classes;
        let mut best_j = 0usize;
        let mut best_v = row[base];
        for j in 1..num_classes {
            let v = row[base + j];
            if v > best_v {
                best_v = v;
                best_j = j;
            }
        }
        if best_j == label {
            correct += 1;
        }
    }
    correct as f32 / batch as f32
}

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Supervised Classification Example (Cross-Entropy) ===");

    // Synthetic 2D inputs, 3 classes with linear-ish separations
    let n = 1200usize;
    let classes = 3usize;
    let mut xs: Vec<f32> = Vec::with_capacity(n * 2);
    let mut ys: Vec<usize> = Vec::with_capacity(n);

    // Simple RNG
    let mut state: u64 = 424242;
    let mut rand_f32 = || {
        state = state.wrapping_mul(1664525).wrapping_add(1013904223);
        (state >> 16) as f32 / (u32::MAX as f32)
    };

    for _ in 0..n {
        let x1 = rand_f32() * 4.0 - 2.0;
        let x2 = rand_f32() * 4.0 - 2.0;
        // Class by quadrant-ish rule with noise
        let mut c = if x1 + 0.5 * x2 > 0.5 {
            0
        } else if x1 - x2 < -0.5 {
            1
        } else {
            2
        };
        if rand_f32() < 0.05 {
            c = (c + 1) % classes;
        }
        xs.push(x1);
        xs.push(x2);
        ys.push(c);
    }

    // Normalize inputs per-feature to [-1, 1]
    let mut min1 = f32::INFINITY;
    let mut max1 = f32::NEG_INFINITY;
    let mut min2 = f32::INFINITY;
    let mut max2 = f32::NEG_INFINITY;
    for i in (0..xs.len()).step_by(2) {
        let a = xs[i];
        let b = xs[i + 1];
        if a < min1 {
            min1 = a;
        }
        if a > max1 {
            max1 = a;
        }
        if b < min2 {
            min2 = b;
        }
        if b > max2 {
            max2 = b;
        }
    }
    let rng1 = (max1 - min1).max(1e-8);
    let rng2 = (max2 - min2).max(1e-8);
    for i in (0..xs.len()).step_by(2) {
        let a = xs[i];
        let b = xs[i + 1];
        xs[i] = 2.0 * (a - min1) / rng1 - 1.0;
        xs[i + 1] = 2.0 * (b - min2) / rng2 - 1.0;
    }

    // Train/Val split (80/20)
    let n_train = (n as f32 * 0.8) as usize;
    let x_train = Tensor::from_slice(&xs[..n_train * 2], vec![n_train, 2]).unwrap();
    let y_train = ys[..n_train].to_vec();
    let x_val = Tensor::from_slice(&xs[n_train * 2..], vec![n - n_train, 2]).unwrap();
    let y_val = ys[n_train..].to_vec();

    // Model: 2 -> 64 -> 64 -> 3 (logits)
    let cfg = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![64, 64],
        output_size: classes,
        use_bias: true,
    };
    let mut net = FeedForwardNetwork::new(cfg, Some(303));

    // Optimizer
    let mut opt = Adam::with_learning_rate(1e-3);
    for p in net.parameters() {
        opt.add_parameter(p);
    }

    let epochs = 300usize;
    let max_grad_norm = 1.0f32;
    let mut best_val_acc = 0.0f32;
    let mut best_val_loss = f32::INFINITY;

    for e in 0..epochs {
        // Zero grads
        {
            let mut params = net.parameters();
            opt.zero_grad(&mut params);
        }

        // Forward logits
        let logits = net.forward(&x_train);
        let mut loss = cross_entropy_logits(&logits, &y_train, n_train, classes);
        loss.backward(None);

        // Step clipped
        {
            let params = net.parameters();
            let mut with_grads: Vec<&mut Tensor> = Vec::new();
            for p in params {
                if p.grad_owned().is_some() {
                    with_grads.push(p);
                }
            }
            if !with_grads.is_empty() {
                clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                opt.step(&mut with_grads);
                opt.zero_grad(&mut with_grads);
            }
        }

        // Metrics
        let train_acc = accuracy_from_logits(&logits, &y_train, n_train, classes);
        let val_logits = net.forward(&x_val);
        let val_loss = cross_entropy_logits(&val_logits, &y_val, n - n_train, classes).value();
        let val_acc = accuracy_from_logits(&val_logits, &y_val, n - n_train, classes);
        if val_acc > best_val_acc {
            best_val_acc = val_acc;
        }
        if val_loss < best_val_loss {
            best_val_loss = val_loss;
        }

        if e % 10 == 0 || e + 1 == epochs {
            println!(
                "epoch {:4} | loss={:.4} acc={:.3} | val_loss={:.4} val_acc={:.3} | best_val_acc={:.3}",
                e, loss.value(), train_acc, val_loss, val_acc, best_val_acc
            );
        }

        clear_all_graphs_known();
    }

    // Quick sample preds via softmax
    let samples = Tensor::from_slice(&[-1.0, -1.0, 0.0, 0.0, 1.0, 1.0], vec![3, 2]).unwrap();
    let sm = net.forward(&samples).softmax(1);
    println!("sample class probs: {:?}", sm.data());

    println!("=== Supervised classification finished ===");
    Ok(())
}

Additional examples can be found in:

• examples/supervised_training/supervised_regression.rs
• examples/getting_started/tensor_operators.rs
• examples/RL_training/../neural_networks/basic_linear_layer.rs
• examples/iterators/element_iteration.rs
• examples/getting_started/tensor_basics.rs
• examples/neural_networks/feedforward_network.rs
• examples/getting_started/serialization_basics.rs
• examples/getting_started/optimizer_basics.rs
• examples/neural_networks/multi_head_attention.rs
• examples/iterators/advanced_patterns.rs
• examples/iterators/performance_optimization.rs

pub fn data_mut(&mut self) -> &mut [f32]

Returns a mutable slice of the tensor’s underlying data

Provides safe mutable access to the tensor’s data without requiring unsafe pointer operations. Use this for in-place modifications of tensor values.

Returns

A mutable slice containing all tensor elements in row-major order

Performance

• Zero-Cost: Direct slice creation with no copying
• Cache-Friendly: Sequential memory access pattern
• Safe: No unsafe code required for basic data modification

Examples

use train_station::Tensor;

let mut tensor = Tensor::new(vec![2, 2]);
let data = tensor.data_mut();

// Safe indexing for modification
data[0] = 1.0;
data[1] = 2.0;

assert_eq!(tensor.get(&[0, 0]), 1.0);
assert_eq!(tensor.get(&[0, 1]), 2.0);

Examples found in repository

examples/RL_training/dqn.rs (line 112)

    fn copy_from(&mut self, other: &Self) {
        for (t, s) in self.layers.iter_mut().zip(other.layers.iter()) {
            {
                let src = s.weight.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = s.bias.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }

examples/RL_training/td3.rs (line 118)

    fn copy_from(&mut self, other: &Self) {
        for (t, s) in self.layers.iter_mut().zip(other.layers.iter()) {
            {
                let src = s.weight.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = s.bias.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }

    fn soft_update_from(&mut self, source: &Self, tau: f32) {
        let _ng = NoGradTrack::new();
        for (t, s) in self.layers.iter_mut().zip(source.layers.iter()) {
            // In-place Polyak update to preserve tensor IDs (no optimizer relink needed)
            let new_w = t
                .weight
                .mul_scalar(1.0 - tau)
                .add_tensor(&s.weight.mul_scalar(tau));
            let new_b = t
                .bias
                .mul_scalar(1.0 - tau)
                .add_tensor(&s.bias.mul_scalar(tau));
            {
                let src = new_w.data();
                let dst = t.weight.data_mut();
                dst.copy_from_slice(src);
            }
            {
                let src = new_b.data();
                let dst = t.bias.data_mut();
                dst.copy_from_slice(src);
            }
            t.weight.set_requires_grad(false);
            t.bias.set_requires_grad(false);
        }
    }

examples/getting_started/tensor_basics.rs (line 72)

fn demonstrate_tensor_creation() {
    println!("--- Tensor Creation ---");

    // Create tensors with different initializations
    let zeros = Tensor::zeros(vec![2, 3]);
    println!(
        "Zeros tensor: shape {:?}, data: {:?}",
        zeros.shape().dims(),
        zeros.data()
    );

    let ones = Tensor::ones(vec![3, 2]);
    println!(
        "Ones tensor: shape {:?}, data: {:?}",
        ones.shape().dims(),
        ones.data()
    );

    // Create tensor from slice
    let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let from_slice = Tensor::from_slice(&data, vec![2, 3]).unwrap();
    println!(
        "From slice: shape {:?}, data: {:?}",
        from_slice.shape().dims(),
        from_slice.data()
    );

    // Create tensor with specific value
    let mut filled = Tensor::new(vec![2, 2]);
    {
        let data = filled.data_mut();
        for value in data.iter_mut() {
            *value = 42.0;
        }
    }
    println!("Filled with 42: {:?}", filled.data());

    // Create tensor with random data
    let random = Tensor::randn(vec![2, 2], Some(42));
    println!(
        "Random tensor: shape {:?}, data: {:?}",
        random.shape().dims(),
        random.data()
    );
}

examples/neural_networks/basic_transformer.rs (line 96)

    pub fn infer_autoregressive(&self, src: &Tensor, max_steps: usize) -> Tensor {
        let (b, _s, e) = Self::triple(src);
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }

        let mut out_seq: Vec<Tensor> = Vec::new();
        // Start token: zeros
        let mut current = Tensor::zeros(vec![b, 1, e]);
        for _step in 0..max_steps {
            // Build causal mask for length t
            let t = current.shape().dims()[1];
            let mut causal = Tensor::ones(vec![b, self.num_heads, t, t]);
            // Upper triangle as false -> masked for all batches and heads
            for bb in 0..b {
                for hh in 0..self.num_heads {
                    for i in 0..t {
                        for j in (i + 1)..t {
                            let offset = causal.memory_offset(&[bb, hh, i, j]);
                            let data = causal.data_mut();
                            data[offset] = 0.0;
                        }
                    }
                }
            }
            let mut step_out = current.clone();
            for dec in &self.decoders {
                step_out = dec.forward(&step_out, &memory, Some(&causal), None);
            }
            // (Toy) append placeholder token; real models would project last token
            out_seq.push(step_out.clone());
            // Append a zero token to grow sequence by 1 for next causal computation
            current = Tensor::zeros(vec![b, t + 1, e]);
        }
        // Simple return of final sequence placeholder
        current
    }

    /// Non auto-regressive inference: single forward pass
    pub fn infer_non_autoregressive(&self, src: &Tensor, tgt_len: usize) -> Tensor {
        let (b, _s, e) = Self::triple(src);
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }
        let tgt = Tensor::zeros(vec![b, tgt_len, e]);
        let mut out = tgt.clone();
        for dec in &self.decoders {
            out = dec.forward(&out, &memory, None, None);
        }
        out
    }

    /// Helper: build boolean-like causal mask [b, heads, t, t] with 1.0 keep, 0.0 masked
    fn build_causal_mask_static(batch: usize, heads: usize, t: usize) -> Tensor {
        let mut mask = Tensor::ones(vec![batch, heads, t, t]);
        for bb in 0..batch {
            for hh in 0..heads {
                for i in 0..t {
                    for j in (i + 1)..t {
                        let offset = mask.memory_offset(&[bb, hh, i, j]);
                        let data = mask.data_mut();
                        data[offset] = 0.0;
                    }
                }
            }
        }
        mask
    }

examples/neural_networks/multi_head_attention.rs (line 191)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Multi-Head Attention Example ===");

    let batch = 2usize;
    let src_len = 5usize;
    let tgt_len = 4usize;
    let embed = 16usize;
    let heads = 4usize;

    let query = Tensor::randn(vec![batch, tgt_len, embed], Some(7));
    let key = Tensor::randn(vec![batch, src_len, embed], Some(8));
    let value = Tensor::randn(vec![batch, src_len, embed], Some(9));

    let mut mha = MultiHeadAttention::new(embed, heads, Some(42));

    // Simple causal mask for target self-attention shape [b, h, tq, tk]
    let mut mask = Tensor::zeros(vec![batch, heads, tgt_len, src_len]);
    // Disallow attending to future positions when tgt_len <= src_len by adding -1e9
    // Here, just demonstrate mask broadcast/add mechanics with a light mask on last head
    if src_len >= tgt_len {
        // set upper triangle to a large negative value for head 0
        for i in 0..tgt_len {
            for j in (i + 1)..src_len {
                let idx = [0usize, 0usize, i, j];
                // Quick set via data_mut using a slice view
                let offset = mask.memory_offset(&idx);
                let data = mask.data_mut();
                data[offset] = -1e9;
            }
        }
    }

    let out = mha.forward(&query, &key, &value, Some(&mask));
    println!("Output shape: {:?}", out.shape().dims());

    // Tiny training step to confirm gradients are wired
    let mut optimizer = Adam::with_learning_rate(0.01);
    let mut params = mha.parameters();
    for p in &params {
        optimizer.add_parameter(p);
    }

    // Dummy loss = mean of output
    let mut loss = out.mean();
    loss.backward(None);
    optimizer.step(&mut params);
    optimizer.zero_grad(&mut params);

    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

pub fn value(&self) -> f32

Extract scalar value from single-element tensor

This method provides a convenient way to extract the scalar value from tensors that contain exactly one element. This is commonly used with element iterator results and scalar tensor operations.

Returns

The scalar value contained in this tensor

Panics

Panics if the tensor does not contain exactly one element

Examples

use train_station::Tensor;

// Single-element tensor
let scalar = Tensor::from_slice(&[42.0], vec![1]).unwrap();
assert_eq!(scalar.value(), 42.0);

Examples found in repository

examples/supervised_training/supervised_regression.rs (line 47)

fn rmse(pred: &Tensor, target: &Tensor) -> f32 {
    mse(pred, target).sqrt().value()
}

fn r2_score(pred: &Tensor, target: &Tensor) -> f32 {
    // R^2 = 1 - SS_res / SS_tot
    let y = target;
    let y_mean = y.mean();
    let ss_res = pred.sub_tensor(y).pow_scalar(2.0).sum();
    let ss_tot = y.sub_tensor(&y_mean).pow_scalar(2.0).sum();
    let ss_res_v = ss_res.value();
    let ss_tot_v = ss_tot.value().max(1e-12); // avoid divide by zero
    1.0 - (ss_res_v / ss_tot_v)
}

examples/neural_networks/basic_encoder.rs (line 98)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Encoder Example ===");

    let batch = 2usize;
    let seq = 6usize;
    let embed = 32usize;
    let heads = 4usize;

    let input = Tensor::randn(vec![batch, seq, embed], Some(11));
    let mut enc = EncoderBlock::new(embed, heads, Some(123));

    // Example: no mask (set Some(mask) to use masking)
    let out = enc.forward(&input, None);
    println!("Output shape: {:?}", out.shape().dims());

    // Verify gradients/optimization
    let mut opt = Adam::with_learning_rate(0.01);
    let mut params = enc.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/getting_started/tensor_basics.rs (line 192)

fn demonstrate_utility_functions() {
    println!("\n--- Utility Functions ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

    // Basic properties
    println!("Shape: {:?}", tensor.shape().dims());
    println!("Size: {}", tensor.size());
    println!("Is contiguous: {}", tensor.is_contiguous());
    println!("Device: {:?}", tensor.device());

    // Mathematical operations
    let sum = tensor.sum();
    println!("Sum: {}", sum.value());

    let mean = tensor.mean();
    println!("Mean: {}", mean.value());

    let norm = tensor.norm();
    println!("Norm: {}", norm.value());

    // Device placement
    let cpu_tensor = Tensor::zeros_on_device(vec![3, 3], train_station::Device::cpu());
    println!(
        "CPU tensor: shape {:?}, device: {:?}",
        cpu_tensor.shape().dims(),
        cpu_tensor.device()
    );
}

examples/neural_networks/basic_decoder.rs (line 109)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Decoder Example ===");

    let batch = 2usize;
    let src = 7usize;
    let tgt = 5usize;
    let embed = 32usize;
    let heads = 4usize;

    let memory = Tensor::randn(vec![batch, src, embed], Some(21));
    let tgt_in = Tensor::randn(vec![batch, tgt, embed], Some(22));

    let mut dec = DecoderBlock::new(embed, heads, Some(456));
    let out = dec.forward(&tgt_in, &memory, None, None);
    println!("Output shape: {:?}", out.shape().dims());

    let mut opt = Adam::with_learning_rate(0.01);
    let mut params = dec.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/neural_networks/basic_transformer.rs (line 169)

    pub fn train_non_autoregressive_steps(
        &mut self,
        src: &Tensor,
        tgt: &Tensor,
        steps: usize,
        lr: f32,
    ) {
        let mut opt = Adam::with_learning_rate(lr);
        {
            let params_once = self.parameters();
            for p in &params_once {
                opt.add_parameter(p);
            }
        }
        for step in 0..steps {
            // forward + backward scope (immutable borrow)
            {
                let pred = self.forward(src, tgt);
                let diff = pred.sub_tensor(tgt);
                let mut loss = diff.pow_scalar(2.0).mean();
                if step == 0 || step + 1 == steps {
                    println!("NAR train step {}: loss={:.6}", step, loss.value());
                }
                loss.backward(None);
            }
            // step + zero_grad scope (mutable borrow)
            let mut params_step = self.parameters();
            opt.step(&mut params_step);
            opt.zero_grad(&mut params_step);
        }
    }

    /// Auto-regressive training (teacher forcing): predict next token with causal mask
    pub fn train_autoregressive_steps(
        &mut self,
        src: &Tensor,
        tgt: &Tensor,
        steps: usize,
        lr: f32,
    ) {
        let mut opt = Adam::with_learning_rate(lr);
        {
            let params_once = self.parameters();
            for p in &params_once {
                opt.add_parameter(p);
            }
        }

        // Build encoder memory once (static dataset demo)
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }

        let (b, t, _e) = Self::triple(tgt);
        // Predict y[t] from y[:t] using causal mask; here we simply predict full seq with mask
        let causal = Self::build_causal_mask_static(b, self.num_heads, t);
        for step in 0..steps {
            // forward + backward scope
            {
                let mut out = tgt.clone();
                for dec in &self.decoders {
                    out = dec.forward(&out, &memory, Some(&causal), None);
                }
                let diff = out.sub_tensor(tgt);
                let mut loss = diff.pow_scalar(2.0).mean();
                if step == 0 || step + 1 == steps {
                    println!("AR  train step {}: loss={:.6}", step, loss.value());
                }
                loss.backward(None);
            }
            let mut params_step = self.parameters();
            opt.step(&mut params_step);
            opt.zero_grad(&mut params_step);
        }
    }

    fn triple(t: &Tensor) -> (usize, usize, usize) {
        let d = t.shape().dims();
        (d[0], d[1], d[2])
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Transformer Example ===");

    let batch = 2usize;
    let src_len = 8usize;
    let tgt_len = 6usize;
    let embed = 32usize;
    let heads = 4usize;
    let layers = 2usize;

    let src = Tensor::randn(vec![batch, src_len, embed], Some(1001));
    let tgt = Tensor::randn(vec![batch, tgt_len, embed], Some(1002));

    let mut trf = BasicTransformer::new(embed, heads, layers, Some(999));
    let out = trf.forward(&src, &tgt);
    println!("Output shape: {:?}", out.shape().dims());

    // Quick optimization step
    let mut opt = Adam::with_learning_rate(0.005);
    let mut params = trf.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());

    // Demo: non auto-regressive inference (single pass)
    let nar = trf.infer_non_autoregressive(&src, tgt_len);
    println!("NAR output shape: {:?}", nar.shape().dims());

    // Demo: auto-regressive inference (toy)
    let ar = trf.infer_autoregressive(&src, 3);
    println!("AR output shape: {:?}", ar.shape().dims());

    // NAR training demo
    let nar_tgt = tgt.clone();
    trf.train_non_autoregressive_steps(&src, &nar_tgt, 3, 0.01);

    // AR training demo (teacher-forced)
    let ar_tgt = tgt.clone();
    trf.train_autoregressive_steps(&src, &ar_tgt, 3, 0.01);
    println!("=== Done ===");
    Ok(())
}

examples/iterators/element_iteration.rs (line 106)

fn demonstrate_basic_iteration() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Basic Element Iteration ---");

    // Create a simple tensor for demonstration
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;
    println!("Original tensor: {:?}", tensor.data());

    // Basic iteration with for loop
    println!("\nBasic iteration with for loop:");
    for (i, element) in tensor.iter().enumerate() {
        println!(
            "  Element {}: value = {:.1}, shape = {:?}",
            i,
            element.value(),
            element.shape().dims()
        );
    }

    // Element-wise transformation
    println!("\nElement-wise transformation (2x + 1):");
    let transformed: Tensor = tensor
        .iter()
        .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0))
        .collect();
    println!("  Result: {:?}", transformed.data());

    // Filtering elements
    println!("\nFiltering elements (values > 3.0):");
    let filtered: Tensor = tensor.iter().filter(|elem| elem.value() > 3.0).collect();
    println!("  Filtered: {:?}", filtered.data());

    Ok(())
}

/// Demonstrate standard iterator trait methods
///
/// Shows compatibility with Rust's standard library iterator methods
/// and demonstrates various functional programming patterns.
fn demonstrate_standard_methods() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Standard Iterator Methods ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;

    // Using map for transformations
    println!("\nMap transformation (square each element):");
    let squared: Tensor = tensor.iter().map(|elem| elem.pow_scalar(2.0)).collect();
    println!("  Squared: {:?}", squared.data());

    // Using enumerate for indexed operations
    println!("\nEnumerate with indexed operations:");
    let indexed: Tensor = tensor
        .iter()
        .enumerate()
        .map(|(i, elem)| elem.add_scalar(i as f32))
        .collect();
    println!("  Indexed: {:?}", indexed.data());

    // Using fold for reduction
    println!("\nFold for sum calculation:");
    let sum: f32 = tensor.iter().fold(0.0, |acc, elem| acc + elem.value());
    println!("  Sum: {:.1}", sum);

    // Using find for element search
    println!("\nFind specific element:");
    if let Some(found) = tensor.iter().find(|elem| elem.value() == 3.0) {
        println!("  Found element with value 3.0: {:.1}", found.value());
    }

    // Using any/all for condition checking
    println!("\nCondition checking:");
    let all_positive = tensor.iter().all(|elem| elem.value() > 0.0);
    let any_large = tensor.iter().any(|elem| elem.value() > 4.0);
    println!("  All positive: {}", all_positive);
    println!("  Any > 4.0: {}", any_large);

    Ok(())
}

/// Demonstrate gradient tracking through element operations
///
/// Shows how gradient tracking works seamlessly through iterator
/// operations, maintaining the computational graph for backpropagation.
fn demonstrate_gradient_tracking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Gradient Tracking ---");

    // Create a tensor with gradient tracking enabled
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?.with_requires_grad();
    println!("Input tensor (requires_grad): {:?}", tensor.data());

    // Perform element-wise operations through iteration
    let result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Apply a complex transformation: (x^2 + 1) * 2
            elem.pow_scalar(2.0).add_scalar(1.0).mul_scalar(2.0)
        })
        .collect();

    println!("Result tensor: {:?}", result.data());
    println!("Result requires_grad: {}", result.requires_grad());

    // Compute gradients
    let mut loss = result.sum();
    loss.backward(None);

    println!("Loss: {:.6}", loss.value());
    println!("Input gradients: {:?}", tensor.grad().map(|g| g.data()));

    Ok(())
}

/// Demonstrate advanced iterator patterns
///
/// Shows complex iterator chains and advanced functional programming
/// patterns for sophisticated data processing workflows.
fn demonstrate_advanced_patterns() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Iterator Patterns ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6])?;
    println!("Input tensor: {:?}", tensor.data());

    // Complex chain: enumerate -> filter -> map -> collect
    println!("\nComplex chain (even indices only, add index to value):");
    let result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 2 == 0) // Take even indices
        .map(|(i, elem)| elem.add_scalar(i as f32)) // Add index to value
        .collect();
    println!("  Result: {:?}", result.data());

    // Using take and skip for windowing
    println!("\nWindowing with take and skip:");
    let window1: Tensor = tensor.iter().take(3).collect();
    let window2: Tensor = tensor.iter().skip(2).take(3).collect();
    println!("  Window 1 (first 3): {:?}", window1.data());
    println!("  Window 2 (middle 3): {:?}", window2.data());

    // Using rev() for reverse iteration
    println!("\nReverse iteration:");
    let reversed: Tensor = tensor.iter().rev().collect();
    println!("  Reversed: {:?}", reversed.data());

    // Chaining with mathematical operations
    println!("\nMathematical operation chain:");
    let math_result: Tensor = tensor
        .iter()
        .map(|elem| elem.exp()) // e^x
        .filter(|elem| elem.value() < 50.0) // Filter large values
        .map(|elem| elem.log()) // ln(x)
        .collect();
    println!("  Math chain result: {:?}", math_result.data());

    // Using zip for element-wise combinations
    println!("\nElement-wise combination with zip:");
    let tensor2 = Tensor::from_slice(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0], vec![6])?;
    let combined: Tensor = tensor
        .iter()
        .zip(tensor2.iter())
        .map(|(a, b)| a.mul_tensor(&b)) // Element-wise multiplication
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

Additional examples can be found in:

• examples/neural_networks/multi_head_attention.rs
• examples/optimizers/adam_configurations.rs
• examples/optimizers/learning_rate_scheduling.rs
• examples/getting_started/optimizer_basics.rs
• examples/neural_networks/feedforward_network.rs
• examples/iterators/advanced_patterns.rs
• examples/iterators/performance_optimization.rs
• examples/RL_training/../neural_networks/basic_linear_layer.rs
• examples/supervised_training/supervised_bce.rs
• examples/supervised_training/supervised_classification.rs
• examples/RL_training/dqn.rs
• examples/RL_training/ppo_discrete.rs
• examples/RL_training/ppo_continuous.rs
• examples/RL_training/td3.rs

pub fn view(&self, new_shape: Vec<i32>) -> Tensor

Create a view with a new shape (requires contiguous memory)

Behaves like PyTorch view: tensor must be contiguous and the total number of elements must remain the same. Supports -1 inference for one dimension.

Arguments

• new_shape - New shape for the tensor (can contain -1 for inference)

Returns

A tensor viewing the same data with a new shape

Examples

use train_station::Tensor;

let x = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
let y = x.view(vec![2, 2]);
assert_eq!(y.shape().dims(), vec![2, 2]);

Examples found in repository

examples/RL_training/ppo_continuous.rs (line 112)

    fn forward(&self, state: &Tensor) -> (Tensor, Tensor) {
        // Returns (mean [B, A], log_std [A])
        let mean = self.net.forward(state);
        (
            mean,
            self.log_std
                .view(vec![1, self.log_std.shape().dims()[0] as i32]),
        )
    }

examples/RL_training/td3.rs (line 196)

    fn forward(&self, state: &Tensor, action: &Tensor) -> Tensor {
        // Concatenate along feature dim (dim=1) for batched inputs
        // IMPORTANT: use views to preserve gradient graph; cloning would detach autograd
        let s_view = state.view(state.shape().dims().iter().map(|&d| d as i32).collect());
        let a_view = action.view(action.shape().dims().iter().map(|&d| d as i32).collect());
        let sa = Tensor::cat(&[s_view, a_view], 1);
        self.net.forward(&sa, None)
    }

examples/neural_networks/basic_encoder.rs (line 59)

    pub fn forward(&self, input: &Tensor, attn_mask: Option<&Tensor>) -> Tensor {
        let attn = self.mha.forward(input, input, input, attn_mask);
        let res1 = attn.add_tensor(input);

        // Feed-forward network with ReLU and residual
        let (b, t, e) = Self::triple(input);
        let x2d = res1.contiguous().view(vec![(b * t) as i32, e as i32]);
        let hidden = self.ffn_in.forward(&x2d).relu();
        let out2d = self.ffn_out.forward(&hidden);
        let out = out2d.view(vec![b as i32, t as i32, e as i32]);
        out.add_tensor(&res1)
    }

examples/neural_networks/basic_decoder.rs (line 70)

    pub fn forward(
        &self,
        tgt: &Tensor,
        memory: &Tensor,
        causal_mask: Option<&Tensor>,
        cross_mask: Option<&Tensor>,
    ) -> Tensor {
        let self_attn = self.self_attn.forward(tgt, tgt, tgt, causal_mask);
        let res1 = self_attn.add_tensor(tgt);

        let cross = self.cross_attn.forward(&res1, memory, memory, cross_mask);
        let res2 = cross.add_tensor(&res1);

        let (b, t, e) = Self::triple(tgt);
        let x2d = res2.contiguous().view(vec![(b * t) as i32, e as i32]);
        let hidden = self.ffn_in.forward(&x2d).relu();
        let out2d = self.ffn_out.forward(&hidden);
        let out = out2d.view(vec![b as i32, t as i32, e as i32]);
        out.add_tensor(&res2)
    }

examples/getting_started/tensor_basics.rs (line 131)

fn demonstrate_shape_operations() {
    println!("\n--- Shape Operations ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
    println!(
        "Original: shape {:?}, data: {:?}",
        tensor.shape().dims(),
        tensor.data()
    );

    // Reshape (view)
    let reshaped = tensor.view(vec![3, 2]);
    println!(
        "Reshaped to [3, 2]: shape {:?}, data: {:?}",
        reshaped.shape().dims(),
        reshaped.data()
    );

    // Create a different shaped tensor for demonstration
    let tensor_2d = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    println!(
        "2D tensor: shape {:?}, data: {:?}",
        tensor_2d.shape().dims(),
        tensor_2d.data()
    );

    // Create a 1D tensor
    let tensor_1d = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
    println!(
        "1D tensor: shape {:?}, data: {:?}",
        tensor_1d.shape().dims(),
        tensor_1d.data()
    );
}

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 359)

fn demonstrate_single_vs_batch_inference() {
    println!("\n--- Single vs Batch Inference ---");

    let layer = LinearLayer::new(4, 3, Some(46));

    // Single inference
    println!("Single inference:");
    let single_input = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![1, 4]).unwrap();
    let single_output = layer.forward_no_grad(&single_input);
    println!("  Input shape: {:?}", single_input.shape().dims());
    println!("  Output shape: {:?}", single_output.shape().dims());
    println!("  Output: {:?}", single_output.data());

    // Batch inference
    println!("Batch inference:");
    let batch_input = Tensor::from_slice(
        &[
            1.0, 2.0, 3.0, 4.0, // Sample 1
            5.0, 6.0, 7.0, 8.0, // Sample 2
            9.0, 10.0, 11.0, 12.0, // Sample 3
        ],
        vec![3, 4],
    )
    .unwrap();
    let batch_output = layer.forward_no_grad(&batch_input);
    println!("  Input shape: {:?}", batch_input.shape().dims());
    println!("  Output shape: {:?}", batch_output.shape().dims());

    // Verify batch consistency - first sample should match single inference
    let _first_batch_sample = batch_output.view(vec![3, 3]); // Reshape to access first sample
    let first_sample_data = &batch_output.data()[0..3]; // First 3 elements
    let single_sample_data = single_output.data();

    println!("Consistency check:");
    println!("  Single output: {:?}", single_sample_data);
    println!("  First batch sample: {:?}", first_sample_data);
    println!(
        "  Match: {}",
        single_sample_data
            .iter()
            .zip(first_sample_data.iter())
            .all(|(a, b)| (a - b).abs() < 1e-6)
    );
}

Additional examples can be found in:

• examples/neural_networks/multi_head_attention.rs
• examples/iterators/performance_optimization.rs

pub fn element_view(&self, index: usize) -> Tensor

Create an element view for the specified index

Returns a scalar tensor (shape [1]) that views a single element of the source tensor. Maintains gradient tracking.

Arguments

• index - Linear index of the element to view

Returns

A scalar tensor viewing the specified element

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let element = tensor.element_view(1);
assert_eq!(element.value(), 2.0);

Examples found in repository

examples/iterators/advanced_patterns.rs (line 132)

fn demonstrate_data_pipeline() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Data Processing Pipeline ---");

    // Simulate raw sensor data with noise
    let raw_data: Vec<f32> = (0..20)
        .map(|i| {
            let base = i as f32 * 0.5;
            let noise = (i % 3) as f32 * 0.1;
            base + noise
        })
        .collect();

    let tensor = Tensor::from_slice(&raw_data, vec![20])?;
    println!("Raw sensor data: {:?}", tensor.data());

    // Multi-stage processing pipeline
    println!("\nProcessing pipeline:");
    println!("1. Normalize data (z-score)");
    println!("2. Apply smoothing filter");
    println!("3. Detect outliers");
    println!("4. Apply feature scaling");

    // Stage 1: Normalization
    let mean = tensor.mean().value();
    let std = tensor.std().value();
    let normalized: Tensor = tensor
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();
    println!(
        "  Normalized (mean={:.3}, std={:.3}): {:?}",
        mean,
        std,
        normalized.data()
    );

    // Stage 2: Smoothing (simple moving average)
    let smoothed: Tensor = normalized
        .iter()
        .enumerate()
        .map(|(i, elem)| {
            if i == 0 || i == normalized.size() - 1 {
                elem.clone()
            } else {
                // Simple 3-point average
                let prev = normalized.element_view(i - 1);
                let next = normalized.element_view(i + 1);
                elem.add_tensor(&prev).add_tensor(&next).div_scalar(3.0)
            }
        })
        .collect();
    println!("  Smoothed: {:?}", smoothed.data());

    // Stage 3: Outlier detection and removal
    let outlier_threshold = 2.0;
    let cleaned: Tensor = smoothed
        .iter()
        .filter(|elem| elem.value().abs() < outlier_threshold)
        .collect();
    println!(
        "  Outliers removed (threshold={}): {:?}",
        outlier_threshold,
        cleaned.data()
    );

    // Stage 4: Feature scaling to [0, 1] range
    let min_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);
    let scaled: Tensor = cleaned
        .iter()
        .map(|elem| elem.sub_scalar(min_val).div_scalar(max_val - min_val))
        .collect();
    println!("  Scaled to [0,1]: {:?}", scaled.data());

    Ok(())
}

/// Demonstrate conditional processing patterns
///
/// Shows how to implement dynamic filtering and transformation
/// based on data characteristics and conditions.
fn demonstrate_conditional_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Conditional Processing ---");

    // Create data with mixed characteristics
    let data = vec![1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0, 9.0, -10.0];
    let tensor = Tensor::from_slice(&data, vec![10])?;
    println!("Input data: {:?}", tensor.data());

    // Conditional transformation based on sign
    println!("\nConditional transformation (positive/negative handling):");
    let processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val > 0.0 {
                elem.pow_scalar(2.0) // Square positive values
            } else {
                elem.mul_scalar(-1.0).sqrt() // Square root of absolute negative values
            }
        })
        .collect();
    println!("  Processed: {:?}", processed.data());

    // Adaptive filtering based on local statistics
    println!("\nAdaptive filtering (remove values > 2 std from local mean):");
    let window_size = 3;
    let adaptive_filtered: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, elem)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(tensor.size());

            // Calculate local mean and std
            let local_values: Vec<f32> = (start..end)
                .map(|j| tensor.element_view(j).value())
                .collect();

            let local_mean = local_values.iter().sum::<f32>() / local_values.len() as f32;
            let local_variance = local_values
                .iter()
                .map(|v| (v - local_mean).powi(2))
                .sum::<f32>()
                / local_values.len() as f32;
            let local_std = local_variance.sqrt();

            let threshold = local_mean + 2.0 * local_std;
            elem.value() <= threshold
        })
        .map(|(_, elem)| elem)
        .collect();
    println!("  Adaptive filtered: {:?}", adaptive_filtered.data());

    // Multi-condition processing
    println!("\nMulti-condition processing:");
    let multi_processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            match () {
                _ if val > 5.0 => elem.mul_scalar(2.0), // Double large values
                _ if val < -5.0 => elem.div_scalar(2.0), // Halve small values
                _ if val.abs() < 2.0 => elem.add_scalar(1.0), // Add 1 to small values
                _ => elem.clone(),                      // Keep others unchanged
            }
        })
        .collect();
    println!("  Multi-condition: {:?}", multi_processed.data());

    Ok(())
}

pub fn slice_view(&self, start: usize, step: usize, length: usize) -> Tensor

Create a slice view of the tensor

Returns a view of a contiguous or strided slice of the source tensor.

Arguments

• start - Starting index
• step - Step size (1 for contiguous)
• length - Number of elements

Returns

A tensor viewing the specified slice

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5]).unwrap();
let slice = tensor.slice_view(1, 2, 2); // [2.0, 4.0]
assert_eq!(slice.get(&[0]), 2.0);
assert_eq!(slice.get(&[1]), 4.0);

Examples found in repository

examples/RL_training/ppo_discrete.rs (line 475)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/RL_training/ppo_continuous.rs (line 489)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

pub fn allocation_owner(&self) -> Option<&Arc<Allocation>>

Get the allocation owner for this tensor

Returns the shared allocation owner if this tensor is a view, or None if this tensor owns its memory directly.

Returns

Optional reference to the allocation owner

Implementation Details

This method is used internally to manage memory lifecycle for tensor views. It helps determine whether a tensor shares memory with another tensor.

pub fn new_uninitialized(shape_dims: Vec<usize>) -> Self

Create a new tensor with uninitialized memory

This method allocates memory for a tensor without initializing it to any value. This is useful for performance-critical operations where the memory will be immediately overwritten, such as matrix multiplication results.

Safety

The caller must ensure that all memory is written before reading from the tensor. Reading from uninitialized memory is undefined behavior.

Arguments

• shape_dims - The dimensions of the tensor

Returns

A tensor with uninitialized memory

Performance

• Zero Initialization: Skips memory initialization for maximum performance
• SIMD Ready: Properly aligned for vectorized operations
• Memory Efficient: Uses optimized alignment strategies

Example

use train_station::Tensor;

// Create uninitialized tensor for matmul result
let mut result = Tensor::new_uninitialized(vec![100, 100]);
// Initialize the memory before use
for value in result.data_mut() {
    *value = 0.0;
}

pub fn new_uninitialized_aligned( shape_dims: Vec<usize>, alignment_bytes: usize, ) -> Self

Create a new uninitialized tensor with an explicit alignment request (in bytes)

This is intended for internal high-performance paths (e.g., packed GEMM panels) where stronger alignment such as 64 bytes is desired even on AVX2 systems.

impl Tensor

pub fn gather( &self, dim: usize, indices: &[usize], index_shape: &[usize], ) -> Tensor

Gather values along a dimension using a tensor of indices

This operation extracts elements from the input tensor based on indices provided along a specified dimension. The output tensor has the same shape as the index tensor, with each element taken from the input tensor at the corresponding position with the index value substituted for the specified dimension.

The gather operation is commonly used in machine learning for operations like embedding lookups, attention mechanisms, and advanced indexing patterns.

Arguments

• dim - The dimension along which to gather values (must be < tensor rank)
• indices - Flattened indices buffer containing the positions to gather from
• index_shape - Shape of the indices tensor and output tensor

Returns

A new tensor with shape index_shape containing the gathered values

Examples

Basic Gather Operation

use train_station::Tensor;

// Create a 2x3 tensor: [[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]]
let tensor = Tensor::from_slice(&[0.0, 0.1, 0.2, 0.3, 0.4, 0.5], vec![2, 3]).unwrap();

// Gather along dimension 1 (columns) with indices [2, 0, 1, 1]
let indices = [2, 0, 1, 1];
let index_shape = [2, 2];
let result = tensor.gather(1, &indices, &index_shape);

// Result shape is [2, 2]
assert_eq!(result.shape().dims(), vec![2, 2]);

// Row 0: indices [2, 0] -> [0.2, 0.0]
assert!((result.get(&[0, 0]) - 0.2).abs() < 1e-6);
assert!((result.get(&[0, 1]) - 0.0).abs() < 1e-6);

// Row 1: indices [1, 1] -> [0.4, 0.4]
assert!((result.get(&[1, 0]) - 0.4).abs() < 1e-6);
assert!((result.get(&[1, 1]) - 0.4).abs() < 1e-6);

Gather with Gradient Tracking

use train_station::Tensor;

let tensor = Tensor::from_slice(&[0.0, 0.1, 0.2, 0.3, 0.4, 0.5], vec![2, 3]).unwrap()
    .with_requires_grad();

let indices = [1, 1, 0, 2];
let index_shape = [2, 2];
let mut result = tensor.gather(1, &indices, &index_shape);

// Compute gradients
result.backward(None);
let grad = tensor.grad_owned().expect("gradient missing");

// Verify gradient accumulation for repeated indices
assert!((grad.get(&[0, 1]) - 2.0).abs() < 1e-6); // Index 1 used twice in row 0

Performance Characteristics

• Time Complexity: O(n) where n is the number of elements in the output
• Memory Usage: Creates a new tensor with the same size as the index tensor
• Optimization: Uses precomputed strides for efficient memory access
• GradTrack Overhead: Minimal overhead when gradient tracking is enabled

Implementation Details

The gather operation works by:

1. Validating input dimensions and index bounds
2. Creating an output tensor with the specified index shape
3. Iterating through all positions in the output tensor
4. Computing source offsets using the input tensor’s strides
5. Copying values from the input tensor to the output tensor
6. Registering the operation for gradient computation if needed

Safety

This function performs bounds checking to ensure:

• The specified dimension is within the tensor’s rank
• All indices are within bounds for the specified dimension
• The index shape is compatible with the input tensor shape
• The indices buffer length matches the product of index shape dimensions

Panics

This function will panic if:

• dim is greater than or equal to the tensor’s rank
• Any index in indices is out of bounds for the specified dimension
• The index_shape rank doesn’t match the input tensor’s rank
• The index_shape dimensions don’t match the input tensor (except along dim)
• The indices length doesn’t equal the product of index_shape dimensions

Examples found in repository

examples/supervised_training/supervised_classification.rs (line 57)

fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

examples/RL_training/ppo_discrete.rs (line 286)

fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

examples/RL_training/dqn.rs (line 470)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== DQN Example (YardEnv discrete) ===");

    // Dims
    let state_dim = 3usize;
    let action_dim = 3usize;

    // Hparams
    let gamma = 0.99f32;
    let batch_size = 64usize;
    let start_steps = 200usize;
    let target_update_interval = 200usize; // hard update cadence
    let max_grad_norm = 1.0f32;
    let mut epsilon = 1.0f32;
    let eps_min = 0.05f32;
    let eps_decay_steps = 2_000usize; // linear decay
    let total_steps = std::env::var("DQN_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3000usize);

    // Models
    let mut q_net = QNet::new(state_dim, action_dim, Some(7));
    let mut q_targ = QNet::new(state_dim, action_dim, Some(8));
    q_targ.net.copy_from(&q_net.net);
    q_targ.set_requires_grad_all(false);

    // Optimizer
    let mut q_opt = Adam::with_learning_rate(3e-4);
    for p in q_net.parameters() {
        q_opt.add_parameter(p);
    }

    // Replay + env
    let mut rb = ReplayBuffer::new(100_000, state_dim);
    let mut env = YardEnv::new(12345);
    let mut rng = SmallRng::new(999_111);

    // Metrics
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    for t in 0..total_steps {
        // Epsilon-greedy action
        let action_index = if t < start_steps || rng.next_f32() < epsilon {
            rng.sample_index(action_dim)
        } else {
            let _ng = NoGradTrack::new();
            let q_vals = q_net.forward(&state);
            let row = q_vals.data();
            let mut best_i = 0usize;
            let mut best_v = row[0];
            for (i, &r) in row.iter().enumerate().take(action_dim).skip(1) {
                if r > best_v {
                    best_v = r;
                    best_i = i;
                }
            }
            best_i
        };

        // Env step
        let (next_state, reward, done) = env.step(action_index);
        episode_return += reward;

        // Store
        let s_slice = state.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            action_index,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        // Reset on done
        state = if done {
            let st = env.reset();
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Epsilon linear decay
        if t < eps_decay_steps {
            epsilon = (1.0 - (t as f32) / (eps_decay_steps as f32)) * (1.0 - eps_min) + eps_min;
        }

        // Train
        if rb.can_sample(batch_size) {
            let (s, a_idx, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Double DQN target: a* = argmax_a Q_online(s2,a); y = r + (1-d)*gamma*Q_target(s2, a*)
            let target_q = {
                let _ng = NoGradTrack::new();
                let q_online_s2 = q_net.forward(&s2);
                // argmax per row (manual on CPU)
                let row_stride = action_dim;
                let qd = q_online_s2.data();
                let mut next_actions: Vec<usize> = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let base = i * row_stride;
                    let mut bi = 0usize;
                    let mut bv = qd[base];
                    for j in 1..action_dim {
                        let v = qd[base + j];
                        if v > bv {
                            bv = v;
                            bi = j;
                        }
                    }
                    next_actions.push(bi);
                }
                let q_targ_s2 = q_targ.forward(&s2);
                let q_targ_g = q_targ_s2.gather(1, &next_actions, &[batch_size, 1]);
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&q_targ_g))
            };

            // Q(s,a) for current actions
            // Zero grads first
            {
                let mut params = q_net.parameters();
                q_opt.zero_grad(&mut params);
            }

            let q_all = q_net.forward(&s);
            let q_sa = q_all.gather(1, &a_idx, &[batch_size, 1]);
            let diff = q_sa.sub_tensor(&target_q);
            let mut loss = pseudo_huber_mean(&diff);
            loss.backward(None);

            // Step (filter only params with grads)
            {
                let params = q_net.parameters();
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    let gn = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    q_opt.step(&mut with_grads);
                    q_opt.zero_grad(&mut with_grads);
                    if t % 100 == 0 {
                        let mut pn = q_net.parameters();
                        let pn_l2 = params_l2_norm(&mut pn);
                        let q_mean = q_all.mean().value();
                        println!(
                            "t={:5} | loss={:.4} | q_mean={:.3} | grad_norm={:.3} | param_norm={:.3} | eps={:.3}",
                            t, loss.value(), q_mean, gn, pn_l2, epsilon
                        );
                    }
                }
            }

            // Target hard update
            if t % target_update_interval == 0 {
                q_targ.net.copy_from(&q_net.net);
            }

            // Clear graphs
            clear_all_graphs_known();
        }
    }

    println!("=== DQN training finished ===");
    Ok(())
}

impl Tensor

pub fn index_select(&self, dim: usize, indices: &[usize]) -> Tensor

Select elements along a dimension using a list of indices

This operation extracts elements from the input tensor along a specified dimension using the provided indices. The output tensor has the same shape as the input except along the specified dimension, where the size becomes the length of the indices array.

The index_select operation is commonly used for extracting specific rows, columns, or slices from tensors, and is particularly useful in machine learning for operations like embedding lookups and attention mechanisms.

Arguments

• dim - The dimension along which to select elements (must be < tensor rank)
• indices - Array of indices specifying which elements to select along dim

Returns

A new tensor with the same shape as the input except along dim, where the size is indices.len()

Examples

Basic Index Selection

use train_station::Tensor;

// Create a 2x3 tensor: [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
let tensor = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0], vec![2, 3]).unwrap();

// Select columns 2 and 0 from dimension 1
let result = tensor.index_select(1, &[2, 0]);

// Result shape is [2, 2] (same as input except dim 1 is now 2)
assert_eq!(result.shape().dims(), vec![2, 2]);

// Row 0: selected columns [2, 0] -> [2.0, 0.0]
assert_eq!(result.get(&[0, 0]), 2.0);
assert_eq!(result.get(&[0, 1]), 0.0);

// Row 1: selected columns [2, 0] -> [5.0, 3.0]
assert_eq!(result.get(&[1, 0]), 5.0);
assert_eq!(result.get(&[1, 1]), 3.0);

Index Selection with Gradient Tracking

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap()
    .with_requires_grad();

// Select specific elements with gradient tracking enabled
let mut result = tensor.index_select(1, &[1, 2]);
result.backward(None);

// Verify gradients are computed correctly
let grad = tensor.grad_owned().expect("gradient missing");
assert_eq!(grad.shape().dims(), vec![2, 3]);

Selecting Rows from a Matrix

use train_station::Tensor;

// Create a 3x2 matrix
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]).unwrap();

// Select rows 2 and 0 (dimension 0)
let result = tensor.index_select(0, &[2, 0]);

// Result shape is [2, 2]
assert_eq!(result.shape().dims(), vec![2, 2]);

// Selected rows: row 2 [5.0, 6.0], row 0 [1.0, 2.0]
assert_eq!(result.get(&[0, 0]), 5.0); // First row of result (was row 2)
assert_eq!(result.get(&[0, 1]), 6.0);
assert_eq!(result.get(&[1, 0]), 1.0); // Second row of result (was row 0)
assert_eq!(result.get(&[1, 1]), 2.0);

Performance Characteristics

• Time Complexity: O(n) where n is the number of elements in the output tensor
• Memory Usage: Creates a new tensor with size equal to the output shape
• Optimization: Uses precomputed strides for efficient memory access
• GradTrack Overhead: Minimal overhead when gradient tracking is enabled
• Memory Layout: Output tensor is always contiguous for optimal performance

Implementation Details

The index_select operation works by:

1. Validating the dimension and index bounds
2. Computing the output shape (same as input except along dim)
3. Creating a new contiguous output tensor
4. Iterating through all positions in the output tensor using nested loops:
   • Outer loop: iterate over dimensions before dim
   • Middle loop: iterate over the selected indices
   • Inner loop: iterate over dimensions after dim
5. Computing source offsets using the input tensor’s strides
6. Copying values from input to output tensor
7. Registering the operation for gradient computation if needed

Safety

This function performs comprehensive bounds checking to ensure:

• The specified dimension is within the tensor’s rank
• All indices are within bounds for the specified dimension
• Memory access is safe through proper offset calculations

Panics

This function will panic if:

• dim is greater than or equal to the tensor’s rank
• Any index in indices is out of bounds for the specified dimension

Thread Safety

This function is thread-safe and can be called concurrently on different tensors. The operation does not modify the input tensor and creates a new output tensor.

impl Tensor

pub fn masked_fill(&self, mask: &[bool], value: f32) -> Tensor

Fill masked elements with a specified value

This operation returns a copy of the input tensor where elements are replaced by the specified value wherever the corresponding boolean mask is true. Elements where the mask is false retain their original values from the input tensor.

The masked_fill operation is commonly used in machine learning for operations like masking attention weights, zeroing out specific elements, and implementing dropout-like functionality.

Arguments

• mask - Boolean array with the same length as the number of tensor elements
• value - The value to fill masked positions with

Returns

A new tensor with the same shape as the input, where masked elements are replaced by value and unmasked elements retain their original values

Examples

Basic Masked Fill

use train_station::Tensor;

// Create a 2x3 tensor: [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
let tensor = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0], vec![2, 3]).unwrap();

// Create a mask: [false, true, false, true, false, true]
let mask = [false, true, false, true, false, true];
let result = tensor.masked_fill(&mask, -1.0);

// Result: [[0.0, -1.0, 2.0], [-1.0, 4.0, -1.0]]
assert_eq!(result.shape().dims(), vec![2, 3]);
assert_eq!(result.get(&[0, 0]), 0.0);   // Unmasked
assert_eq!(result.get(&[0, 1]), -1.0);  // Masked
assert_eq!(result.get(&[0, 2]), 2.0);   // Unmasked
assert_eq!(result.get(&[1, 0]), -1.0);  // Masked
assert_eq!(result.get(&[1, 1]), 4.0);   // Unmasked
assert_eq!(result.get(&[1, 2]), -1.0);  // Masked

Masked Fill with Gradient Tracking

use train_station::Tensor;

let tensor = Tensor::from_slice(&[0.0, 0.1, 0.2, 0.3, 0.4, 0.5], vec![2, 3]).unwrap()
    .with_requires_grad();

// Create a mask with some true values
let mask = [false, true, false, true, false, false];
let mut result = tensor.masked_fill(&mask, 5.0);

// Compute gradients
result.backward(None);
let grad = tensor.grad_owned().expect("gradient missing");

// Gradients should be zero where mask is true, 1 elsewhere
assert_eq!(grad.shape().dims(), vec![2, 3]);
assert!((grad.get(&[0, 0]) - 1.0).abs() < 1e-6);   // Unmasked: gradient flows
assert!((grad.get(&[0, 1]) - 0.0).abs() < 1e-6);   // Masked: no gradient
assert!((grad.get(&[0, 2]) - 1.0).abs() < 1e-6);   // Unmasked: gradient flows

Zeroing Out Specific Elements

use train_station::Tensor;

// Create a tensor with some values
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();

// Create a mask to zero out every other element
let mask = [true, false, true, false, true, false];
let result = tensor.masked_fill(&mask, 0.0);

// Result: [[0.0, 2.0, 0.0], [4.0, 0.0, 6.0]]
assert_eq!(result.get(&[0, 0]), 0.0);  // Zeroed
assert_eq!(result.get(&[0, 1]), 2.0);  // Kept
assert_eq!(result.get(&[0, 2]), 0.0);  // Zeroed
assert_eq!(result.get(&[1, 0]), 4.0);  // Kept
assert_eq!(result.get(&[1, 1]), 0.0);  // Zeroed
assert_eq!(result.get(&[1, 2]), 6.0);  // Kept

Performance Characteristics

• Time Complexity: O(n) where n is the number of elements in the tensor
• Memory Usage: Creates a new tensor with the same size as the input
• Optimization: Uses efficient stride-based iteration for non-contiguous tensors
• GradTrack Overhead: Minimal overhead when gradient tracking is enabled
• Memory Layout: Output tensor is always contiguous for optimal performance

Implementation Details

The masked_fill operation works by:

1. Validating that the mask length equals the number of tensor elements
2. Creating a new contiguous output tensor with the same shape
3. Iterating through all elements in logical order
4. For each element, checking the corresponding mask value:
   • If mask is true: use the fill value
   • If mask is false: copy the original value from input tensor
5. Computing source offsets using the input tensor’s shape for non-contiguous tensors
6. Registering the operation for gradient computation if needed

Safety

This function performs bounds checking to ensure:

• The mask length equals the number of tensor elements
• Memory access is safe through proper offset calculations
• The operation handles both contiguous and non-contiguous tensors correctly

Panics

This function will panic if:

• The mask length does not equal the number of tensor elements

Thread Safety

This function is thread-safe and can be called concurrently on different tensors. The operation does not modify the input tensor and creates a new output tensor.

GradTrack Behavior

When gradient tracking is enabled:

• Gradients do not flow through masked positions (they are zeroed)
• Gradients flow normally through unmasked positions
• This behavior is useful for implementing operations like dropout

Examples found in repository

examples/neural_networks/multi_head_attention.rs (line 101)

    pub fn forward(
        &self,
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        attn_mask: Option<&Tensor>,
    ) -> Tensor {
        let qkv = Self::project_qkv(query, key, value, &self.q_proj, &self.k_proj, &self.v_proj);
        let (q, k, v) = qkv;

        // Split heads: [b, t, e] -> [b, h, t, d]
        let (b, tq, _e) = Self::triple(query);
        let (_b2, tk, _e2) = Self::triple(key);
        let q = Self::split_heads(&q, b, tq, self.num_heads, self.head_dim);
        let k = Self::split_heads(&k, b, tk, self.num_heads, self.head_dim);
        let v = Self::split_heads(&v, b, tk, self.num_heads, self.head_dim);

        // Scaled dot-product attention
        // logits: [b, h, tq, tk]
        let k_t = k.transpose(2, 3);
        let mut logits = q.matmul(&k_t).div_scalar((self.head_dim as f32).sqrt());
        if let Some(mask) = attn_mask {
            let dims = mask.shape().dims().to_vec();
            // If boolean-like mask matching [b,h,tq,tk], apply masked_fill
            if dims.len() == 4 && dims[0] == b && dims[1] == self.num_heads && dims[2] == tq {
                // Interpret mask > 0.5 as keep; we invert to build masked positions
                let cond: Vec<bool> = mask.data().iter().map(|&v| v < 0.5).collect();
                // Apply masked fill on a flattened view, then reshape back
                let flat_logits = logits.view(vec![(b * self.num_heads * tq * tk) as i32]);
                let filled = flat_logits.masked_fill(&cond, f32::NEG_INFINITY);
                logits = filled.view(vec![b as i32, self.num_heads as i32, tq as i32, tk as i32]);
            } else {
                // Fallback: additive mask
                logits = logits.add_tensor(mask);
            }
        }
        let attn = logits.softmax(3);

        // context: [b, h, tq, d]
        let context = attn.matmul(&v);
        let context = context.permute(vec![0, 2, 1, 3]); // [b, tq, h, d]
        let context = context.contiguous().view(vec![
            b as i32,
            tq as i32,
            (self.num_heads * self.head_dim) as i32,
        ]);

        // Output projection (flatten to 2D, project, then restore 3D)
        let flat = context.view(vec![(b * tq) as i32, self.embed_dim as i32]);
        let out2d = self.out_proj.forward(&flat);
        out2d.view(vec![b as i32, tq as i32, self.embed_dim as i32])
    }

impl Tensor

pub fn select(&self, dim: usize, index: usize) -> Tensor

Select a slice along a given dimension at a specific index

This operation extracts a slice from the input tensor by fixing a specific dimension at a given index. The result is a tensor with one fewer dimension than the input, containing the selected slice.

The select operation returns a view (zero-copy) when the base offset is zero, otherwise it creates a contiguous copy to ensure correctness. This operation is commonly used for extracting specific rows, columns, or slices from tensors.

Arguments

• dim - The dimension along which to select (must be < tensor rank)
• index - The index within the specified dimension to select (must be < dim size)

Returns

A tensor with the selected slice. The result has the same shape as the input except with the specified dimension removed.

Performance

Returns a view when possible (base offset is zero) to avoid copying. On non-zero offsets, falls back to a contiguous copy for correctness. Gradients propagate back to the selected slice when GradTrack is enabled.

Examples

Basic Row Selection

use train_station::Tensor;

// Create a 2x3 tensor: [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
let tensor = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0], vec![2, 3]).unwrap();

// Select row 1 (dimension 0, index 1)
let result = tensor.select(0, 1);

// Result shape is [3] (dimension 0 removed)
assert_eq!(result.shape().dims(), vec![3]);
assert_eq!(result.get(&[0]), 3.0);  // First element of row 1
assert_eq!(result.get(&[1]), 4.0);  // Second element of row 1
assert_eq!(result.get(&[2]), 5.0);  // Third element of row 1

Column Selection

use train_station::Tensor;

// Create a 2x3 tensor: [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
let tensor = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0], vec![2, 3]).unwrap();

// Select column 1 (dimension 1, index 1)
let result = tensor.select(1, 1);

// Result shape is [2] (dimension 1 removed)
assert_eq!(result.shape().dims(), vec![2]);
assert_eq!(result.get(&[0]), 1.0);  // Column 1, row 0
assert_eq!(result.get(&[1]), 4.0);  // Column 1, row 1

Select with Gradient Tracking

use train_station::Tensor;

let tensor = Tensor::from_slice(&[0.0, 1.0, 2.0, 3.0], vec![2, 2]).unwrap()
    .with_requires_grad();

// Select row 1 with gradient tracking enabled
let mut result = tensor.select(0, 1);
result.backward(None);

// Verify gradients are computed correctly
let grad = tensor.grad_owned().expect("gradient missing");
assert_eq!(grad.shape().dims(), vec![2, 2]);
// Only row 1 receives gradients
assert_eq!(grad.get(&[0, 0]), 0.0);  // Row 0: no gradient
assert_eq!(grad.get(&[0, 1]), 0.0);  // Row 0: no gradient
assert_eq!(grad.get(&[1, 0]), 1.0);  // Row 1: gradient flows
assert_eq!(grad.get(&[1, 1]), 1.0);  // Row 1: gradient flows

Performance Characteristics

• Time Complexity: O(n) where n is the number of elements in the selected slice
• Memory Usage: Zero-copy view when base offset is zero, otherwise creates a copy
• Optimization: Uses efficient stride-based access for non-contiguous tensors
• GradTrack Overhead: Minimal overhead when gradient tracking is enabled
• Memory Layout: Result is contiguous when a copy is made, view otherwise

Implementation Details

The select operation works by:

1. Validating the dimension and index bounds
2. Computing the new shape by removing the selected dimension
3. Computing the new strides by removing the selected dimension’s stride
4. Calculating the base offset for the selected slice
5. If base offset is zero: creating a view with adjusted shape/strides
6. If base offset is non-zero: creating a contiguous copy of the slice
7. Registering the operation for gradient computation if needed

Safety

This function performs comprehensive bounds checking to ensure:

• The tensor has non-zero rank
• The specified dimension is within the tensor’s rank
• The index is within bounds for the specified dimension
• Memory access is safe through proper offset calculations

Panics

This function will panic if:

• The tensor has zero rank
• dim is greater than or equal to the tensor’s rank
• index is greater than or equal to the size of the specified dimension

Thread Safety

This function is thread-safe and can be called concurrently on different tensors. The operation does not modify the input tensor and creates either a view or a new tensor.

View vs Copy Behavior

• View (zero-copy): When the base offset is zero, returns a view that shares the same memory as the input tensor with adjusted shape and strides
• Copy: When the base offset is non-zero, creates a contiguous copy to ensure correctness across all operations

GradTrack Behavior

When gradient tracking is enabled:

• Gradients are scattered back to the selected slice in the input tensor
• Other positions in the input tensor receive zero gradients
• This behavior ensures correct gradient flow for the selected elements

impl Tensor

pub fn zeros(shape_dims: Vec<usize>) -> Self

Creates a new tensor filled with zeros

Convenience constructor that creates a tensor and initializes all elements to zero. Uses optimized SIMD operations for efficient zero initialization.

Arguments

• shape_dims - Vector of dimension sizes defining the tensor shape

Returns

A new tensor with all elements initialized to zero

Performance

• Memory Allocation: Single allocation with optimized alignment
• Initialization: SIMD-optimized zero filling for large tensors
• Thread Safe: Atomic ID generation for gradtrack tracking

Examples

use train_station::Tensor;

let tensor = Tensor::zeros(vec![2, 3]);
assert_eq!(tensor.size(), 6);
assert_eq!(tensor.shape().dims(), vec![2, 3]);

// Verify all elements are zero
assert_eq!(tensor.get(&[0, 0]), 0.0);
assert_eq!(tensor.get(&[1, 2]), 0.0);

Examples found in repository

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 60)

    pub fn new(input_size: usize, output_size: usize, seed: Option<u64>) -> Self {
        // Xavier/Glorot initialization: scale by sqrt(1/input_size)
        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,
        }
    }

examples/getting_started/tensor_basics.rs (line 46)

fn demonstrate_tensor_creation() {
    println!("--- Tensor Creation ---");

    // Create tensors with different initializations
    let zeros = Tensor::zeros(vec![2, 3]);
    println!(
        "Zeros tensor: shape {:?}, data: {:?}",
        zeros.shape().dims(),
        zeros.data()
    );

    let ones = Tensor::ones(vec![3, 2]);
    println!(
        "Ones tensor: shape {:?}, data: {:?}",
        ones.shape().dims(),
        ones.data()
    );

    // Create tensor from slice
    let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let from_slice = Tensor::from_slice(&data, vec![2, 3]).unwrap();
    println!(
        "From slice: shape {:?}, data: {:?}",
        from_slice.shape().dims(),
        from_slice.data()
    );

    // Create tensor with specific value
    let mut filled = Tensor::new(vec![2, 2]);
    {
        let data = filled.data_mut();
        for value in data.iter_mut() {
            *value = 42.0;
        }
    }
    println!("Filled with 42: {:?}", filled.data());

    // Create tensor with random data
    let random = Tensor::randn(vec![2, 2], Some(42));
    println!(
        "Random tensor: shape {:?}, data: {:?}",
        random.shape().dims(),
        random.data()
    );
}

examples/getting_started/optimizer_basics.rs (line 52)

fn demonstrate_basic_optimizer_setup() {
    println!("--- Basic Optimizer Setup ---");

    // Create parameters that require gradients
    let weight = Tensor::randn(vec![3, 2], Some(42)).with_requires_grad();
    let bias = Tensor::zeros(vec![2]).with_requires_grad();

    println!("Created parameters:");
    println!(
        "  Weight: shape {:?}, requires_grad: {}",
        weight.shape().dims(),
        weight.requires_grad()
    );
    println!(
        "  Bias: shape {:?}, requires_grad: {}",
        bias.shape().dims(),
        bias.requires_grad()
    );

    // Create Adam optimizer with default configuration
    let mut optimizer = Adam::new();
    println!(
        "Created Adam optimizer with learning rate: {}",
        optimizer.learning_rate()
    );

    // Add parameters to optimizer
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);
    println!(
        "Added {} parameters to optimizer",
        optimizer.parameter_count()
    );

    // Create optimizer with custom configuration
    let config = AdamConfig {
        learning_rate: 0.01,
        beta1: 0.9,
        beta2: 0.999,
        eps: 1e-8,
        weight_decay: 0.0,
        amsgrad: false,
    };

    let mut custom_optimizer = Adam::with_config(config);
    custom_optimizer.add_parameter(&weight);
    custom_optimizer.add_parameter(&bias);

    println!(
        "Created custom optimizer with learning rate: {}",
        custom_optimizer.learning_rate()
    );

    // Demonstrate parameter linking
    println!("Parameter linking completed successfully");
}

/// Demonstrate simple linear regression training
fn demonstrate_linear_regression() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Linear Regression Training ---");

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(43)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Create simple training data: y = 2*x + 1
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    println!("Training data:");
    println!("  X: {:?}", x_data.data());
    println!("  Y: {:?}", y_true.data());
    println!("  Target: y = 2*x + 1");

    // Training loop
    let num_epochs = 100;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass: y_pred = x * weight + bias
        let y_pred = x_data.matmul(&weight) + &bias;

        // Compute loss: MSE
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        losses.push(loss.value());

        // Print progress every 20 epochs
        if epoch % 20 == 0 || epoch == num_epochs - 1 {
            println!("Epoch {:3}: Loss = {:.6}", epoch, loss.value());
        }
    }

    // Evaluate final model
    let final_predictions = x_data.matmul(&weight) + &bias;
    println!("\nFinal model evaluation:");
    println!("  Learned weight: {:.6}", weight.value());
    println!("  Learned bias: {:.6}", bias.value());
    println!("  Predictions vs True:");

    for i in 0..5 {
        let x1 = x_data.data()[i];
        let pred = final_predictions.data()[i];
        let true_val = y_true.data()[i];
        println!(
            "    x={:.1}: pred={:.3}, true={:.1}, error={:.3}",
            x1,
            pred,
            true_val,
            (pred - true_val).abs()
        );
    }

    Ok(())
}

/// Demonstrate advanced training patterns
fn demonstrate_advanced_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Training Patterns ---");

    // Create a more complex model
    let mut weight = Tensor::randn(vec![1, 2], Some(44)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![2]).with_requires_grad();

    // Create optimizer with different learning rate
    let mut optimizer = Adam::with_learning_rate(0.005);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Create training data: y = 2*x + [1, 3]
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(
        &[3.0, 5.0, 7.0, 9.0, 11.0, 6.0, 8.0, 10.0, 12.0, 14.0],
        vec![5, 2],
    )
    .unwrap();

    println!("Advanced training with monitoring:");
    println!("  Initial learning rate: {}", optimizer.learning_rate());

    // Training loop with monitoring
    let num_epochs = 50;
    let mut losses = Vec::new();
    let mut weight_norms = Vec::new();
    let mut gradient_norms = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Compute gradient norm before optimizer step
        let gradient_norm = weight.grad_owned().unwrap().norm();

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Learning rate scheduling: reduce every 10 epochs
        if epoch > 0 && epoch % 10 == 0 {
            let current_lr = optimizer.learning_rate();
            let new_lr = current_lr * 0.5;
            optimizer.set_learning_rate(new_lr);
            println!(
                "Epoch {:2}: Reduced learning rate from {:.3} to {:.3}",
                epoch, current_lr, new_lr
            );
        }

        // Record metrics
        losses.push(loss.value());
        weight_norms.push(weight.norm().value());
        gradient_norms.push(gradient_norm.value());

        // Print detailed progress
        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, Weight Norm = {:.6}, Gradient Norm = {:.6}",
                epoch,
                loss.value(),
                weight.norm().value(),
                gradient_norm.value()
            );
        }
    }

    println!("Final learning rate: {}", optimizer.learning_rate());

    // Analyze training progression
    let initial_loss = losses[0];
    let final_loss = losses[losses.len() - 1];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final weight norm: {:.6}", weight.norm().value());
    println!("  Final bias: {:?}", bias.data());

    Ok(())
}

/// Demonstrate learning rate scheduling
fn demonstrate_learning_rate_scheduling() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Learning Rate Scheduling ---");

    // Create simple model
    let mut weight = Tensor::randn(vec![1, 1], Some(45)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer with high initial learning rate
    let mut optimizer = Adam::with_learning_rate(0.1);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Simple data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3, 1]).unwrap();
    let y_true = Tensor::from_slice(&[2.0, 4.0, 6.0], vec![3, 1]).unwrap();

    println!("Initial learning rate: {}", optimizer.learning_rate());

    // Training loop with learning rate scheduling
    let num_epochs = 50;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Learning rate scheduling: reduce every 10 epochs
        if epoch > 0 && epoch % 10 == 0 {
            let current_lr = optimizer.learning_rate();
            let new_lr = current_lr * 0.5;
            optimizer.set_learning_rate(new_lr);
            println!(
                "Epoch {:2}: Reduced learning rate from {:.3} to {:.3}",
                epoch, current_lr, new_lr
            );
        }

        losses.push(loss.value());

        // Print progress
        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, LR = {:.3}",
                epoch,
                loss.value(),
                optimizer.learning_rate()
            );
        }
    }

    println!("Final learning rate: {}", optimizer.learning_rate());

    Ok(())
}

/// Demonstrate training monitoring and analysis
fn demonstrate_training_monitoring() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Monitoring ---");

    // Create model
    let mut weight = Tensor::randn(vec![1, 1], Some(46)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Training data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0], vec![4, 1]).unwrap();

    // Training loop with comprehensive monitoring
    let num_epochs = 30;
    let mut losses = Vec::new();
    let mut weight_history = Vec::new();
    let mut bias_history = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Record history
        losses.push(loss.value());
        weight_history.push(weight.value());
        bias_history.push(bias.value());

        // Print detailed monitoring
        if epoch % 5 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, Weight = {:.6}, Bias = {:.6}",
                epoch,
                loss.value(),
                weight.value(),
                bias.value()
            );
        }
    }

    // Analyze training progression
    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", losses[0]);
    println!("  Final loss: {:.6}", losses[losses.len() - 1]);
    println!(
        "  Loss reduction: {:.1}%",
        (losses[0] - losses[losses.len() - 1]) / losses[0] * 100.0
    );

    // Compute statistics
    let loss_mean = compute_mean(&losses);
    let loss_std = compute_std(&losses);
    let weight_change = (weight_history[weight_history.len() - 1] - weight_history[0]).abs();
    let bias_change = (bias_history[bias_history.len() - 1] - bias_history[0]).abs();

    println!("  Average loss: {:.6} ± {:.6}", loss_mean, loss_std);
    println!("  Weight change: {:.6}", weight_change);
    println!("  Bias change: {:.6}", bias_change);
    println!("  Final weight norm: {:.6}", weight.norm().value());
    println!("  Final bias: {:.6}", bias.value());

    Ok(())
}

examples/neural_networks/basic_transformer.rs (line 85)

    pub fn infer_autoregressive(&self, src: &Tensor, max_steps: usize) -> Tensor {
        let (b, _s, e) = Self::triple(src);
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }

        let mut out_seq: Vec<Tensor> = Vec::new();
        // Start token: zeros
        let mut current = Tensor::zeros(vec![b, 1, e]);
        for _step in 0..max_steps {
            // Build causal mask for length t
            let t = current.shape().dims()[1];
            let mut causal = Tensor::ones(vec![b, self.num_heads, t, t]);
            // Upper triangle as false -> masked for all batches and heads
            for bb in 0..b {
                for hh in 0..self.num_heads {
                    for i in 0..t {
                        for j in (i + 1)..t {
                            let offset = causal.memory_offset(&[bb, hh, i, j]);
                            let data = causal.data_mut();
                            data[offset] = 0.0;
                        }
                    }
                }
            }
            let mut step_out = current.clone();
            for dec in &self.decoders {
                step_out = dec.forward(&step_out, &memory, Some(&causal), None);
            }
            // (Toy) append placeholder token; real models would project last token
            out_seq.push(step_out.clone());
            // Append a zero token to grow sequence by 1 for next causal computation
            current = Tensor::zeros(vec![b, t + 1, e]);
        }
        // Simple return of final sequence placeholder
        current
    }

    /// Non auto-regressive inference: single forward pass
    pub fn infer_non_autoregressive(&self, src: &Tensor, tgt_len: usize) -> Tensor {
        let (b, _s, e) = Self::triple(src);
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }
        let tgt = Tensor::zeros(vec![b, tgt_len, e]);
        let mut out = tgt.clone();
        for dec in &self.decoders {
            out = dec.forward(&out, &memory, None, None);
        }
        out
    }

examples/neural_networks/multi_head_attention.rs (line 181)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Multi-Head Attention Example ===");

    let batch = 2usize;
    let src_len = 5usize;
    let tgt_len = 4usize;
    let embed = 16usize;
    let heads = 4usize;

    let query = Tensor::randn(vec![batch, tgt_len, embed], Some(7));
    let key = Tensor::randn(vec![batch, src_len, embed], Some(8));
    let value = Tensor::randn(vec![batch, src_len, embed], Some(9));

    let mut mha = MultiHeadAttention::new(embed, heads, Some(42));

    // Simple causal mask for target self-attention shape [b, h, tq, tk]
    let mut mask = Tensor::zeros(vec![batch, heads, tgt_len, src_len]);
    // Disallow attending to future positions when tgt_len <= src_len by adding -1e9
    // Here, just demonstrate mask broadcast/add mechanics with a light mask on last head
    if src_len >= tgt_len {
        // set upper triangle to a large negative value for head 0
        for i in 0..tgt_len {
            for j in (i + 1)..src_len {
                let idx = [0usize, 0usize, i, j];
                // Quick set via data_mut using a slice view
                let offset = mask.memory_offset(&idx);
                let data = mask.data_mut();
                data[offset] = -1e9;
            }
        }
    }

    let out = mha.forward(&query, &key, &value, Some(&mask));
    println!("Output shape: {:?}", out.shape().dims());

    // Tiny training step to confirm gradients are wired
    let mut optimizer = Adam::with_learning_rate(0.01);
    let mut params = mha.parameters();
    for p in &params {
        optimizer.add_parameter(p);
    }

    // Dummy loss = mean of output
    let mut loss = out.mean();
    loss.backward(None);
    optimizer.step(&mut params);
    optimizer.zero_grad(&mut params);

    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/optimizers/adam_configurations.rs (line 93)

fn demonstrate_default_adam() -> Result<(), Box<dyn std::error::Error>> {
    println!("--- Default Adam Configuration ---");

    // Create a simple regression problem: y = 2*x + 1
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(42)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create Adam optimizer with default configuration
    let mut optimizer = Adam::new();
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    println!("Default Adam configuration:");
    println!("  Learning rate: {}", optimizer.learning_rate());
    println!("  Initial weight: {:.6}", weight.value());
    println!("  Initial bias: {:.6}", bias.value());

    // Training loop
    let num_epochs = 50;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        losses.push(loss.value());

        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!("Epoch {:3}: Loss = {:.6}", epoch, loss.value());
        }
    }

    // Evaluate final model
    let _final_predictions = x_data.matmul(&weight) + &bias;
    println!("\nFinal model:");
    println!("  Learned weight: {:.6} (target: 2.0)", weight.value());
    println!("  Learned bias: {:.6} (target: 1.0)", bias.value());
    println!("  Final loss: {:.6}", losses[losses.len() - 1]);

    Ok(())
}

/// Demonstrate learning rate comparison
fn demonstrate_learning_rate_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Learning Rate Comparison ---");

    let learning_rates = [0.001, 0.01, 0.1];
    let mut results = Vec::new();

    for &lr in &learning_rates {
        println!("\nTesting learning rate: {}", lr);

        let stats = train_with_config(TrainingConfig {
            learning_rate: lr,
            ..Default::default()
        })?;

        results.push((lr, stats.clone()));

        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Convergence epoch: {}", stats.convergence_epoch);
    }

    // Compare results
    println!("\nLearning Rate Comparison Summary:");
    for (lr, stats) in &results {
        println!(
            "  LR={:6}: Loss={:.6}, Converged@{}",
            lr, stats.final_loss, stats.convergence_epoch
        );
    }

    Ok(())
}

/// Demonstrate weight decay comparison
fn demonstrate_weight_decay_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Weight Decay Comparison ---");

    let weight_decays = [0.0, 0.001, 0.01];
    let mut results = Vec::new();

    for &wd in &weight_decays {
        println!("\nTesting weight decay: {}", wd);

        let stats = train_with_config(TrainingConfig {
            weight_decay: wd,
            ..Default::default()
        })?;

        results.push((wd, stats.clone()));

        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Final weight norm: {:.6}", stats.weight_norm);
    }

    // Compare results
    println!("\nWeight Decay Comparison Summary:");
    for (wd, stats) in &results {
        println!(
            "  WD={:6}: Loss={:.6}, Weight Norm={:.6}",
            wd, stats.final_loss, stats.weight_norm
        );
    }

    Ok(())
}

/// Demonstrate beta parameter tuning
fn demonstrate_beta_parameter_tuning() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Beta Parameter Tuning ---");

    let beta_configs = [
        (0.9, 0.999),  // Default
        (0.8, 0.999),  // More aggressive momentum
        (0.95, 0.999), // Less aggressive momentum
        (0.9, 0.99),   // Faster second moment decay
    ];

    let mut results = Vec::new();

    for (i, (beta1, beta2)) in beta_configs.iter().enumerate() {
        println!(
            "\nTesting beta configuration {}: beta1={}, beta2={}",
            i + 1,
            beta1,
            beta2
        );

        let config = TrainingConfig {
            beta1: *beta1,
            beta2: *beta2,
            ..Default::default()
        };

        let stats = train_with_config(config)?;
        results.push(((*beta1, *beta2), stats.clone()));

        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Convergence epoch: {}", stats.convergence_epoch);
    }

    // Compare results
    println!("\nBeta Parameter Comparison Summary:");
    for ((beta1, beta2), stats) in &results {
        println!(
            "  B1={:4}, B2={:5}: Loss={:.6}, Converged@{}",
            beta1, beta2, stats.final_loss, stats.convergence_epoch
        );
    }

    Ok(())
}

/// Demonstrate configuration benchmarking
fn demonstrate_configuration_benchmarking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Configuration Benchmarking ---");

    // Define configurations to benchmark
    let configs = vec![
        (
            "Conservative",
            TrainingConfig {
                learning_rate: 0.001,
                weight_decay: 0.001,
                beta1: 0.95,
                ..Default::default()
            },
        ),
        (
            "Balanced",
            TrainingConfig {
                learning_rate: 0.01,
                weight_decay: 0.0,
                beta1: 0.9,
                ..Default::default()
            },
        ),
        (
            "Aggressive",
            TrainingConfig {
                learning_rate: 0.1,
                weight_decay: 0.0,
                beta1: 0.8,
                ..Default::default()
            },
        ),
    ];

    let mut benchmark_results = Vec::new();

    for (name, config) in configs {
        println!("\nBenchmarking {} configuration:", name);

        let start_time = std::time::Instant::now();
        let stats = train_with_config(config.clone())?;
        let elapsed = start_time.elapsed();

        println!("  Training time: {:.2}ms", elapsed.as_millis());
        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Convergence: {} epochs", stats.convergence_epoch);

        benchmark_results.push((name.to_string(), stats, elapsed));
    }

    // Summary
    println!("\nBenchmarking Summary:");
    for (name, stats, elapsed) in &benchmark_results {
        println!(
            "  {:12}: Loss={:.6}, Time={:4}ms, Converged@{}",
            name,
            stats.final_loss,
            elapsed.as_millis(),
            stats.convergence_epoch
        );
    }

    Ok(())
}

/// Helper function to train with specific configuration
fn train_with_config(config: TrainingConfig) -> Result<TrainingStats, Box<dyn std::error::Error>> {
    // Create training data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(123)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer with custom configuration
    let adam_config = AdamConfig {
        learning_rate: config.learning_rate,
        beta1: config.beta1,
        beta2: config.beta2,
        eps: 1e-8,
        weight_decay: config.weight_decay,
        amsgrad: false,
    };

    let mut optimizer = Adam::with_config(adam_config);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Training loop
    let mut losses = Vec::new();
    let mut convergence_epoch = config.epochs;

    for epoch in 0..config.epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        let loss_value = loss.value();
        losses.push(loss_value);

        // Check for convergence (loss < 0.01)
        if loss_value < 0.01 && convergence_epoch == config.epochs {
            convergence_epoch = epoch;
        }
    }

    Ok(TrainingStats {
        config,
        final_loss: losses[losses.len() - 1],
        loss_history: losses,
        convergence_epoch,
        weight_norm: weight.norm().value(),
    })
}

Additional examples can be found in:

• examples/optimizers/learning_rate_scheduling.rs

pub fn ones(shape_dims: Vec<usize>) -> Self

Creates a new tensor filled with ones

Convenience constructor that creates a tensor and initializes all elements to one. Uses optimized SIMD operations for efficient initialization.

Arguments

• shape_dims - Vector of dimension sizes defining the tensor shape

Returns

A new tensor with all elements initialized to one

Performance

• Memory Allocation: Single allocation with optimized alignment
• Initialization: SIMD-optimized one filling for large tensors
• Thread Safe: Atomic ID generation for gradtrack tracking

Examples

use train_station::Tensor;

let tensor = Tensor::ones(vec![2, 3]);
assert_eq!(tensor.size(), 6);
assert_eq!(tensor.shape().dims(), vec![2, 3]);

// Verify all elements are one
assert_eq!(tensor.get(&[0, 0]), 1.0);
assert_eq!(tensor.get(&[1, 2]), 1.0);

Examples found in repository

examples/getting_started/tensor_basics.rs (line 53)

fn demonstrate_tensor_creation() {
    println!("--- Tensor Creation ---");

    // Create tensors with different initializations
    let zeros = Tensor::zeros(vec![2, 3]);
    println!(
        "Zeros tensor: shape {:?}, data: {:?}",
        zeros.shape().dims(),
        zeros.data()
    );

    let ones = Tensor::ones(vec![3, 2]);
    println!(
        "Ones tensor: shape {:?}, data: {:?}",
        ones.shape().dims(),
        ones.data()
    );

    // Create tensor from slice
    let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let from_slice = Tensor::from_slice(&data, vec![2, 3]).unwrap();
    println!(
        "From slice: shape {:?}, data: {:?}",
        from_slice.shape().dims(),
        from_slice.data()
    );

    // Create tensor with specific value
    let mut filled = Tensor::new(vec![2, 2]);
    {
        let data = filled.data_mut();
        for value in data.iter_mut() {
            *value = 42.0;
        }
    }
    println!("Filled with 42: {:?}", filled.data());

    // Create tensor with random data
    let random = Tensor::randn(vec![2, 2], Some(42));
    println!(
        "Random tensor: shape {:?}, data: {:?}",
        random.shape().dims(),
        random.data()
    );
}

examples/neural_networks/basic_transformer.rs (line 89)

    pub fn infer_autoregressive(&self, src: &Tensor, max_steps: usize) -> Tensor {
        let (b, _s, e) = Self::triple(src);
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }

        let mut out_seq: Vec<Tensor> = Vec::new();
        // Start token: zeros
        let mut current = Tensor::zeros(vec![b, 1, e]);
        for _step in 0..max_steps {
            // Build causal mask for length t
            let t = current.shape().dims()[1];
            let mut causal = Tensor::ones(vec![b, self.num_heads, t, t]);
            // Upper triangle as false -> masked for all batches and heads
            for bb in 0..b {
                for hh in 0..self.num_heads {
                    for i in 0..t {
                        for j in (i + 1)..t {
                            let offset = causal.memory_offset(&[bb, hh, i, j]);
                            let data = causal.data_mut();
                            data[offset] = 0.0;
                        }
                    }
                }
            }
            let mut step_out = current.clone();
            for dec in &self.decoders {
                step_out = dec.forward(&step_out, &memory, Some(&causal), None);
            }
            // (Toy) append placeholder token; real models would project last token
            out_seq.push(step_out.clone());
            // Append a zero token to grow sequence by 1 for next causal computation
            current = Tensor::zeros(vec![b, t + 1, e]);
        }
        // Simple return of final sequence placeholder
        current
    }

    /// Non auto-regressive inference: single forward pass
    pub fn infer_non_autoregressive(&self, src: &Tensor, tgt_len: usize) -> Tensor {
        let (b, _s, e) = Self::triple(src);
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }
        let tgt = Tensor::zeros(vec![b, tgt_len, e]);
        let mut out = tgt.clone();
        for dec in &self.decoders {
            out = dec.forward(&out, &memory, None, None);
        }
        out
    }

    /// Helper: build boolean-like causal mask [b, heads, t, t] with 1.0 keep, 0.0 masked
    fn build_causal_mask_static(batch: usize, heads: usize, t: usize) -> Tensor {
        let mut mask = Tensor::ones(vec![batch, heads, t, t]);
        for bb in 0..batch {
            for hh in 0..heads {
                for i in 0..t {
                    for j in (i + 1)..t {
                        let offset = mask.memory_offset(&[bb, hh, i, j]);
                        let data = mask.data_mut();
                        data[offset] = 0.0;
                    }
                }
            }
        }
        mask
    }

examples/RL_training/dqn.rs (line 471)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== DQN Example (YardEnv discrete) ===");

    // Dims
    let state_dim = 3usize;
    let action_dim = 3usize;

    // Hparams
    let gamma = 0.99f32;
    let batch_size = 64usize;
    let start_steps = 200usize;
    let target_update_interval = 200usize; // hard update cadence
    let max_grad_norm = 1.0f32;
    let mut epsilon = 1.0f32;
    let eps_min = 0.05f32;
    let eps_decay_steps = 2_000usize; // linear decay
    let total_steps = std::env::var("DQN_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3000usize);

    // Models
    let mut q_net = QNet::new(state_dim, action_dim, Some(7));
    let mut q_targ = QNet::new(state_dim, action_dim, Some(8));
    q_targ.net.copy_from(&q_net.net);
    q_targ.set_requires_grad_all(false);

    // Optimizer
    let mut q_opt = Adam::with_learning_rate(3e-4);
    for p in q_net.parameters() {
        q_opt.add_parameter(p);
    }

    // Replay + env
    let mut rb = ReplayBuffer::new(100_000, state_dim);
    let mut env = YardEnv::new(12345);
    let mut rng = SmallRng::new(999_111);

    // Metrics
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    for t in 0..total_steps {
        // Epsilon-greedy action
        let action_index = if t < start_steps || rng.next_f32() < epsilon {
            rng.sample_index(action_dim)
        } else {
            let _ng = NoGradTrack::new();
            let q_vals = q_net.forward(&state);
            let row = q_vals.data();
            let mut best_i = 0usize;
            let mut best_v = row[0];
            for (i, &r) in row.iter().enumerate().take(action_dim).skip(1) {
                if r > best_v {
                    best_v = r;
                    best_i = i;
                }
            }
            best_i
        };

        // Env step
        let (next_state, reward, done) = env.step(action_index);
        episode_return += reward;

        // Store
        let s_slice = state.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            action_index,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        // Reset on done
        state = if done {
            let st = env.reset();
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Epsilon linear decay
        if t < eps_decay_steps {
            epsilon = (1.0 - (t as f32) / (eps_decay_steps as f32)) * (1.0 - eps_min) + eps_min;
        }

        // Train
        if rb.can_sample(batch_size) {
            let (s, a_idx, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Double DQN target: a* = argmax_a Q_online(s2,a); y = r + (1-d)*gamma*Q_target(s2, a*)
            let target_q = {
                let _ng = NoGradTrack::new();
                let q_online_s2 = q_net.forward(&s2);
                // argmax per row (manual on CPU)
                let row_stride = action_dim;
                let qd = q_online_s2.data();
                let mut next_actions: Vec<usize> = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let base = i * row_stride;
                    let mut bi = 0usize;
                    let mut bv = qd[base];
                    for j in 1..action_dim {
                        let v = qd[base + j];
                        if v > bv {
                            bv = v;
                            bi = j;
                        }
                    }
                    next_actions.push(bi);
                }
                let q_targ_s2 = q_targ.forward(&s2);
                let q_targ_g = q_targ_s2.gather(1, &next_actions, &[batch_size, 1]);
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&q_targ_g))
            };

            // Q(s,a) for current actions
            // Zero grads first
            {
                let mut params = q_net.parameters();
                q_opt.zero_grad(&mut params);
            }

            let q_all = q_net.forward(&s);
            let q_sa = q_all.gather(1, &a_idx, &[batch_size, 1]);
            let diff = q_sa.sub_tensor(&target_q);
            let mut loss = pseudo_huber_mean(&diff);
            loss.backward(None);

            // Step (filter only params with grads)
            {
                let params = q_net.parameters();
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    let gn = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    q_opt.step(&mut with_grads);
                    q_opt.zero_grad(&mut with_grads);
                    if t % 100 == 0 {
                        let mut pn = q_net.parameters();
                        let pn_l2 = params_l2_norm(&mut pn);
                        let q_mean = q_all.mean().value();
                        println!(
                            "t={:5} | loss={:.4} | q_mean={:.3} | grad_norm={:.3} | param_norm={:.3} | eps={:.3}",
                            t, loss.value(), q_mean, gn, pn_l2, epsilon
                        );
                    }
                }
            }

            // Target hard update
            if t % target_update_interval == 0 {
                q_targ.net.copy_from(&q_net.net);
            }

            // Clear graphs
            clear_all_graphs_known();
        }
    }

    println!("=== DQN training finished ===");
    Ok(())
}

examples/RL_training/td3.rs (line 545)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== TD3 Example (YardEnv) ===");

    // Environment / problem dims
    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hyperparameters (small for demo)
    let gamma = 0.99f32;
    let tau = 0.005f32; // Polyak
    let policy_noise = 0.2f32; // target smoothing noise stddev
    let exploration_noise = 0.1f32; // behavior policy noise stddev
    let policy_delay = 2usize;
    let batch_size = 64usize;
    let start_steps = 500usize; // random exploration steps
    let total_steps = 1500usize;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(11));
    let mut actor_targ = Actor::new(state_dim, action_dim, Some(12));
    actor_targ.net.copy_from(&actor.net);
    actor_targ.set_requires_grad_all(false);

    let mut critic1 = Critic::new(state_dim, action_dim, Some(21));
    let mut critic2 = Critic::new(state_dim, action_dim, Some(22));
    let mut critic1_targ = Critic::new(state_dim, action_dim, Some(23));
    let mut critic2_targ = Critic::new(state_dim, action_dim, Some(24));
    critic1_targ.net.copy_from(&critic1.net);
    critic2_targ.net.copy_from(&critic2.net);
    critic1_targ.set_requires_grad_all(false);
    critic2_targ.set_requires_grad_all(false);

    // Optimizers
    let mut actor_opt = Adam::with_learning_rate(1e-3);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }

    let mut critic_opt = Adam::with_learning_rate(1e-4);
    for p in critic1.parameters() {
        critic_opt.add_parameter(p);
    }
    for p in critic2.parameters() {
        critic_opt.add_parameter(p);
    }

    // Replay buffer and env
    let mut rb = ReplayBuffer::new(100_000, state_dim, action_dim);
    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(987654321);

    // Reset & metric trackers
    let mut state = env.reset(); // [1, state_dim]
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32; // smooth short-term
    let mut best_return = f32::NEG_INFINITY;
    let mut policy_updates: usize = 0;

    for t in 0..total_steps {
        // Select action
        let action_tensor = if t < start_steps {
            let a = rng.uniform(-1.0, 1.0);
            Tensor::from_slice(&[a], vec![1, action_dim]).unwrap()
        } else {
            // Behavior policy with exploration noise
            let _ng = NoGradTrack::new();
            let det = actor.forward(&state);
            let noise = Tensor::randn(vec![1, action_dim], None).mul_scalar(exploration_noise);
            tanh_bounded(&det.add_tensor(&noise))
        };
        let action_value = action_tensor.data()[0];

        // Environment step
        let (next_state, reward, done) = env.step(action_value);
        episode_return += reward;

        // Store transition
        let s_slice = state.data().to_vec();
        let a_slice = action_tensor.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            &a_slice,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        state = if done {
            let st = env.reset();
            // Metrics: update EMA and best
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={} | policy_updates={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size,
                policy_updates
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Training
        if rb.can_sample(batch_size) {
            // Sample batch
            let (s, a, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Compute target values y = r + (1-d)*gamma*min(Q1', Q2') using target networks (no grad)
            let target_q = {
                let _ng = NoGradTrack::new();
                // Target actions with smoothing noise (tanh bounds)
                let noise =
                    Tensor::randn(vec![batch_size, action_dim], None).mul_scalar(policy_noise);
                let a_targ = tanh_bounded(&actor_targ.forward(&s2).add_tensor(&noise));
                let q1_t = critic1_targ.forward(&s2, &a_targ);
                let q2_t = critic2_targ.forward(&s2, &a_targ);

                // Elementwise min via data() since this path is no-grad
                let q1d = q1_t.data();
                let q2d = q2_t.data();
                let mut min_vec = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let v1 = q1d[i];
                    let v2 = q2d[i];
                    min_vec.push(v1.min(v2));
                }
                let min_q = Tensor::from_slice(&min_vec, vec![batch_size, 1]).unwrap();
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&min_q))
            };

            // Critic update (both critics)
            // Zero grads in a short scope, then drop borrows before forward
            {
                let mut params = {
                    let c_params = critic1.parameters();
                    let c2_params = critic2.parameters();
                    let mut tmp: Vec<&mut Tensor> = Vec::new();
                    tmp.extend(c_params);
                    tmp.extend(c2_params);
                    tmp
                };
                critic_opt.zero_grad(&mut params);
            }

            // Forward current Q estimates
            let q1 = critic1.forward(&s, &a);
            let q2 = critic2.forward(&s, &a);
            let diff1 = q1.sub_tensor(&target_q);
            let diff2 = q2.sub_tensor(&target_q);
            let mut critic_loss = diff1
                .pow_scalar(2.0)
                .mean()
                .add_tensor(&diff2.pow_scalar(2.0).mean());

            // Backward
            critic_loss.backward(None);

            // Optional gradient clipping + step (only for params that received grads)
            {
                let params = {
                    let c_params = critic1.parameters();
                    let c2_params = critic2.parameters();
                    let mut tmp: Vec<&mut Tensor> = Vec::new();
                    tmp.extend(c_params);
                    tmp.extend(c2_params);
                    tmp
                };
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    // Pre-step metrics
                    let grad_norm_before = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    critic_opt.step(&mut with_grads);
                    critic_opt.zero_grad(&mut with_grads);

                    // Post-step metrics (param norm)
                    let mut for_norm_params = {
                        let c_params = critic1.parameters();
                        let c2_params = critic2.parameters();
                        let mut tmp: Vec<&mut Tensor> = Vec::new();
                        tmp.extend(c_params);
                        tmp.extend(c2_params);
                        tmp
                    };
                    let param_norm = params_l2_norm(&mut for_norm_params);

                    // Print compact critic metrics occasionally
                    if t % 100 == 0 {
                        let q1_mean = q1.mean().value();
                        let q2_mean = q2.mean().value();
                        let tq_mean = target_q.mean().value();
                        println!(
                            "t={:5} | critic_loss={:.4} | q1_mean={:.3} q2_mean={:.3} tq_mean={:.3} | grad_norm={:.3} | crit_param_norm={:.3}",
                            t,
                            critic_loss.value(),
                            q1_mean,
                            q2_mean,
                            tq_mean,
                            grad_norm_before,
                            param_norm
                        );
                    }
                }
            }

            // Delayed policy update
            if t % policy_delay == 0 {
                // Actor update: maximize Q1(s, actor(s)) -> minimize -Q1
                // Zero actor grads before backward
                {
                    let mut a_params: Vec<&mut Tensor> = actor.parameters();
                    actor_opt.zero_grad(&mut a_params);
                }

                let a_pred = actor.forward(&s);
                let q_for_actor = critic1.forward(&s, &a_pred);
                let mut actor_loss = q_for_actor.mul_scalar(-1.0).mean();
                actor_loss.backward(None);

                {
                    let a_params: Vec<&mut Tensor> = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in a_params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let grad_norm_before = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);

                        // Post-step param norm
                        let mut for_norm_params = actor.parameters();
                        let param_norm = params_l2_norm(&mut for_norm_params);

                        policy_updates += 1;
                        if t % 200 == 0 {
                            println!(
                                "t={:5} | actor_loss={:.4} | act_grad_norm={:.3} | act_param_norm={:.3} | lr_a={:.4e} lr_c={:.4e} | policy_updates={}",
                                t,
                                actor_loss.value(),
                                grad_norm_before,
                                param_norm,
                                actor_opt.learning_rate(),
                                critic_opt.learning_rate(),
                                policy_updates
                            );
                        }
                    }
                }

                // Target updates (Polyak averaging, no grad)
                actor_targ.net.soft_update_from(&actor.net, tau);
                critic1_targ.net.soft_update_from(&critic1.net, tau);
                critic2_targ.net.soft_update_from(&critic2.net, tau);
            }

            // Clear entire graphs to avoid stale accumulation across iterations
            clear_all_graphs_known();
        }
    }

    println!("=== TD3 training finished ===");
    Ok(())
}

pub fn zeros_on_device(shape_dims: Vec<usize>, device: Device) -> Self

Creates a new tensor filled with zeros on a specific device

Convenience constructor that creates a tensor on the specified device and initializes all elements to zero. Uses optimized SIMD operations for efficient zero initialization.

Arguments

• shape_dims - Vector of dimension sizes defining the tensor shape
• device - The device where the tensor should be allocated

Returns

A new tensor with all elements initialized to zero

Performance

• Memory Allocation: Device-specific allocation with optimized alignment
• Initialization: SIMD-optimized zero filling for large tensors
• Thread Safe: Atomic ID generation for gradtrack tracking

Examples

use train_station::Tensor;
use train_station::Device;

let tensor = Tensor::zeros_on_device(vec![2, 2], Device::cpu());
assert_eq!(tensor.device(), Device::cpu());
assert_eq!(tensor.size(), 4);

// Verify all elements are zero
assert_eq!(tensor.get(&[0, 0]), 0.0);
assert_eq!(tensor.get(&[1, 1]), 0.0);

Examples found in repository

examples/getting_started/tensor_basics.rs (line 201)

fn demonstrate_utility_functions() {
    println!("\n--- Utility Functions ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

    // Basic properties
    println!("Shape: {:?}", tensor.shape().dims());
    println!("Size: {}", tensor.size());
    println!("Is contiguous: {}", tensor.is_contiguous());
    println!("Device: {:?}", tensor.device());

    // Mathematical operations
    let sum = tensor.sum();
    println!("Sum: {}", sum.value());

    let mean = tensor.mean();
    println!("Mean: {}", mean.value());

    let norm = tensor.norm();
    println!("Norm: {}", norm.value());

    // Device placement
    let cpu_tensor = Tensor::zeros_on_device(vec![3, 3], train_station::Device::cpu());
    println!(
        "CPU tensor: shape {:?}, device: {:?}",
        cpu_tensor.shape().dims(),
        cpu_tensor.device()
    );
}

pub fn ones_on_device(shape_dims: Vec<usize>, device: Device) -> Self

Creates a new tensor filled with ones on a specific device

Convenience constructor that creates a tensor on the specified device and initializes all elements to one. Uses optimized SIMD operations for efficient initialization.

Arguments

• shape_dims - Vector of dimension sizes defining the tensor shape
• device - The device where the tensor should be allocated

Returns

A new tensor with all elements initialized to one

Performance

• Memory Allocation: Device-specific allocation with optimized alignment
• Initialization: SIMD-optimized one filling for large tensors
• Thread Safe: Atomic ID generation for gradtrack tracking

Examples

use train_station::Tensor;
use train_station::Device;

let tensor = Tensor::ones_on_device(vec![2, 2], Device::cpu());
assert_eq!(tensor.device(), Device::cpu());
assert_eq!(tensor.size(), 4);

// Verify all elements are one
assert_eq!(tensor.get(&[0, 0]), 1.0);
assert_eq!(tensor.get(&[1, 1]), 1.0);

pub fn fill(&mut self, value: f32)

Fills the tensor with a constant value using SIMD optimization

Efficiently initializes all elements of the tensor to the specified value. Uses SIMD operations for large tensors to maximize performance.

Arguments

• value - The value to fill the tensor with

Performance

• SIMD Optimization: Uses AVX2 for large tensors when available
• Unrolled Loops: 4x unrolling for better instruction throughput
• Memory Bandwidth: Optimized for maximum memory bandwidth utilization

Examples

use train_station::Tensor;

let mut tensor = Tensor::new(vec![2, 3]);
tensor.fill(42.0);

// Verify all elements are 42.0
assert_eq!(tensor.get(&[0, 0]), 42.0);
assert_eq!(tensor.get(&[1, 2]), 42.0);

Zero-Sized Tensor Handling

use train_station::Tensor;

let mut empty_tensor = Tensor::new(vec![0]);
empty_tensor.fill(42.0); // Should not panic
assert_eq!(empty_tensor.size(), 0);

impl Tensor

pub fn from_slice(data: &[f32], shape_dims: Vec<usize>) -> Result<Self, String>

Creates a tensor from a slice of data

Creates a new tensor with the specified shape and copies data from the provided slice. Validates that the data size matches the tensor shape before performing the copy operation.

This method provides an efficient way to create tensors from existing data sources while ensuring data integrity and proper memory management.

Arguments

• data - Slice of f32 values to copy into the tensor
• shape_dims - Vector of dimension sizes defining the tensor shape

Returns

• Ok(Tensor) - Successfully created tensor with copied data
• Err(String) - Error if data size doesn’t match shape

Performance

• Memory Copy: Efficient non-overlapping copy using SIMD when possible
• Validation: Fast size validation before allocation
• Alignment: Proper memory alignment for optimal performance
• Large Data: Optimized handling of large datasets

Examples

Basic Usage

use train_station::Tensor;

let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let tensor = Tensor::from_slice(&data, vec![2, 3]).unwrap();
assert_eq!(tensor.size(), 6);
assert_eq!(tensor.get(&[0, 0]), 1.0);
assert_eq!(tensor.get(&[1, 2]), 6.0);

Multi-Dimensional Data

use train_station::Tensor;

// 1D tensor
let data_1d = [1.0, 2.0, 3.0];
let tensor_1d = Tensor::from_slice(&data_1d, vec![3]).unwrap();
assert_eq!(tensor_1d.shape().dims(), vec![3]);
assert_eq!(tensor_1d.get(&[1]), 2.0);

// 3D tensor
let data_3d = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
let tensor_3d = Tensor::from_slice(&data_3d, vec![2, 2, 2]).unwrap();
assert_eq!(tensor_3d.shape().dims(), vec![2, 2, 2]);
assert_eq!(tensor_3d.get(&[0, 0, 0]), 1.0);
assert_eq!(tensor_3d.get(&[1, 1, 1]), 8.0);

Error Handling

use train_station::Tensor;

// Size mismatch error
let data = [1.0, 2.0, 3.0];
let result = Tensor::from_slice(&data, vec![2, 2]);
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.contains("Data size 3 doesn't match shape size 4"));

Zero-Sized Tensors

use train_station::Tensor;

// Handle empty tensors gracefully
let data: [f32; 0] = [];
let tensor = Tensor::from_slice(&data, vec![0]).unwrap();
assert_eq!(tensor.size(), 0);
assert_eq!(tensor.shape().dims(), vec![0]);

Large Data Sets

use train_station::Tensor;

// Efficient handling of large datasets
let size = 1000;
let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
let tensor = Tensor::from_slice(&data, vec![size]).unwrap();

assert_eq!(tensor.size(), size);
assert_eq!(tensor.get(&[0]), 0.0);
assert_eq!(tensor.get(&[100]), 100.0);
assert_eq!(tensor.get(&[999]), 999.0);

Implementation Details

This method performs the following steps:

1. Shape Validation: Creates a Shape object and validates dimensions
2. Size Check: Ensures data length matches the calculated tensor size
3. Memory Allocation: Allocates tensor memory with proper alignment
4. Data Copy: Uses efficient non-overlapping memory copy operation
5. Return: Returns the created tensor or descriptive error message

The memory copy operation uses std::ptr::copy_nonoverlapping for maximum performance and safety, ensuring no data corruption occurs during the copy process.

Examples found in repository

examples/RL_training/dqn.rs (line 185)

    fn state_tensor(&self) -> Tensor {
        Tensor::from_slice(&[self.pos, self.vel, 0.0], vec![1, 3]).unwrap()
    }

    fn step(&mut self, action_index: usize) -> (Tensor, f32, bool) {
        let a = Self::ACTIONS[action_index.min(2)];
        self.vel += 0.1 * a - 0.01 * self.pos;
        self.pos += self.vel;
        self.steps += 1;
        let reward = -(self.pos * self.pos) - 0.05 * (a * a);
        let done = self.pos.abs() > 3.0 || self.steps >= self.max_steps;
        (self.state_tensor(), reward, done)
    }
}

// -------------------------------
// Replay Buffer
// -------------------------------

struct ReplayBuffer {
    capacity: usize,
    size: usize,
    pos: usize,
    state_dim: usize,
    states: Vec<f32>,
    actions: Vec<usize>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    next_states: Vec<f32>,
}

impl ReplayBuffer {
    fn new(capacity: usize, state_dim: usize) -> Self {
        Self {
            capacity,
            size: 0,
            pos: 0,
            state_dim,
            states: vec![0.0; capacity * state_dim],
            actions: vec![0usize; capacity],
            rewards: vec![0.0; capacity],
            dones: vec![0.0; capacity],
            next_states: vec![0.0; capacity * state_dim],
        }
    }

    fn push(&mut self, s: &[f32], a_idx: usize, r: f32, d: f32, s2: &[f32]) {
        let i = self.pos;
        let so = i * self.state_dim;
        self.states[so..so + self.state_dim].copy_from_slice(s);
        self.actions[i] = a_idx;
        self.rewards[i] = r;
        self.dones[i] = d;
        self.next_states[so..so + self.state_dim].copy_from_slice(s2);
        self.pos = (self.pos + 1) % self.capacity;
        self.size = self.size.saturating_add(1).min(self.capacity);
    }

    fn can_sample(&self, batch_size: usize) -> bool {
        self.size >= batch_size
    }

    fn sample(
        &self,
        batch_size: usize,
        rng: &mut SmallRng,
    ) -> (Tensor, Vec<usize>, Tensor, Tensor, Tensor) {
        let mut s_vec = Vec::with_capacity(batch_size * self.state_dim);
        let mut a_idx = Vec::with_capacity(batch_size);
        let mut r_vec = Vec::with_capacity(batch_size);
        let mut d_vec = Vec::with_capacity(batch_size);
        let mut s2_vec = Vec::with_capacity(batch_size * self.state_dim);
        for _ in 0..batch_size {
            let idx = rng.sample_index(self.size);
            let so = idx * self.state_dim;
            s_vec.extend_from_slice(&self.states[so..so + self.state_dim]);
            a_idx.push(self.actions[idx]);
            r_vec.push(self.rewards[idx]);
            d_vec.push(self.dones[idx]);
            s2_vec.extend_from_slice(&self.next_states[so..so + self.state_dim]);
        }
        let s = Tensor::from_slice(&s_vec, vec![batch_size, self.state_dim]).unwrap();
        let r = Tensor::from_slice(&r_vec, vec![batch_size, 1]).unwrap();
        let d = Tensor::from_slice(&d_vec, vec![batch_size, 1]).unwrap();
        let s2 = Tensor::from_slice(&s2_vec, vec![batch_size, self.state_dim]).unwrap();
        (s, a_idx, r, d, s2)
    }

examples/RL_training/ppo_discrete.rs (line 151)

    fn state_tensor(&self) -> Tensor {
        Tensor::from_slice(&[self.pos, self.vel, 0.0], vec![1, 3]).unwrap()
    }
    fn step(&mut self, action_idx: usize) -> (Tensor, f32, bool) {
        let a = Self::ACTIONS[action_idx.min(2)];
        self.vel += 0.1 * a - 0.01 * self.pos;
        self.pos += self.vel;
        self.steps += 1;
        let reward = -(self.pos * self.pos) - 0.05 * (a * a);
        let done = self.pos.abs() > 3.0 || self.steps >= self.max_steps;
        (self.state_tensor(), reward, done)
    }
}

// -------------------------------
// Rollout storage
// -------------------------------

struct RolloutBatch {
    states: Vec<f32>,
    actions: Vec<usize>,
    old_logps: Vec<f32>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    values: Vec<f32>,
    next_states: Vec<f32>,
    _state_dim: usize,
}
impl RolloutBatch {
    fn new(cap: usize, sd: usize) -> Self {
        Self {
            states: Vec::with_capacity(cap * sd),
            actions: Vec::with_capacity(cap),
            old_logps: Vec::with_capacity(cap),
            rewards: Vec::with_capacity(cap),
            dones: Vec::with_capacity(cap),
            values: Vec::with_capacity(cap),
            next_states: Vec::with_capacity(cap * sd),
            _state_dim: sd,
        }
    }
    #[allow(clippy::too_many_arguments)]
    fn push(&mut self, s: &[f32], a: usize, lp: f32, r: f32, d: f32, v: f32, s2: &[f32]) {
        self.states.extend_from_slice(s);
        self.actions.push(a);
        self.old_logps.push(lp);
        self.rewards.push(r);
        self.dones.push(d);
        self.values.push(v);
        self.next_states.extend_from_slice(s2);
    }
    fn len(&self) -> usize {
        self.actions.len()
    }
}

// -------------------------------
// Helpers
// -------------------------------

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

// log-softmax for selected actions: given logits [B,A] and actions Vec<usize> -> log_prob [B,1]
fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

// probability ratio = exp(new_logp - old_logp)
fn ratio_from_logps(new_logp: &Tensor, old_logp: &Tensor) -> Tensor {
    new_logp.sub_tensor(old_logp).exp()
}

// Clamp ratio to [1-clip, 1+clip] using ReLU-based clamp (no custom ops)
fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/RL_training/ppo_continuous.rs (line 101)

    fn new(state_dim: usize, action_dim: usize, seed: Option<u64>) -> Self {
        let net = Mlp::new(&[state_dim, 64, 64, action_dim], seed);
        let log_std = Tensor::from_slice(&vec![0.0; action_dim], vec![action_dim])
            .unwrap()
            .with_requires_grad();
        Self { net, log_std }
    }
    fn forward(&self, state: &Tensor) -> (Tensor, Tensor) {
        // Returns (mean [B, A], log_std [A])
        let mean = self.net.forward(state);
        (
            mean,
            self.log_std
                .view(vec![1, self.log_std.shape().dims()[0] as i32]),
        )
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        let mut ps = self.net.parameters();
        ps.push(&mut self.log_std);
        ps
    }
}

// -------------------------------
// Critic: value function V(s)
// -------------------------------

struct Critic {
    net: Mlp,
}
impl Critic {
    fn new(state_dim: usize, seed: Option<u64>) -> Self {
        Self {
            net: Mlp::new(&[state_dim, 64, 64, 1], seed),
        }
    }
    fn forward(&self, state: &Tensor) -> Tensor {
        self.net.forward(state)
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.net.parameters()
    }
}

// -------------------------------
// Continuous YardEnv (same dynamics as TD3 env)
// -------------------------------

struct YardEnv {
    pos: f32,
    vel: f32,
    steps: usize,
    max_steps: usize,
    rng: SmallRng,
}
impl YardEnv {
    fn new(seed: u64) -> Self {
        let mut e = Self {
            pos: 0.0,
            vel: 0.0,
            steps: 0,
            max_steps: 200,
            rng: SmallRng::new(seed),
        };
        e.reset();
        e
    }
    fn reset(&mut self) -> Tensor {
        self.pos = (self.rng.next_f32() * 1.0) - 0.5;
        self.vel = (self.rng.next_f32() * 0.2) - 0.1;
        self.steps = 0;
        self.state_tensor()
    }
    fn state_tensor(&self) -> Tensor {
        Tensor::from_slice(&[self.pos, self.vel, 0.0], vec![1, 3]).unwrap()
    }
    fn step(&mut self, action_value: f32) -> (Tensor, f32, bool) {
        let a = action_value.clamp(-1.0, 1.0);
        self.vel += 0.1 * a - 0.01 * self.pos;
        self.pos += self.vel;
        self.steps += 1;
        let reward = -(self.pos * self.pos) - 0.1 * (a * a);
        let done = self.pos.abs() > 3.0 || self.steps >= self.max_steps;
        (self.state_tensor(), reward, done)
    }
}

// -------------------------------
// Trajectory storage
// -------------------------------

struct RolloutBatch {
    states: Vec<f32>,
    actions: Vec<f32>,
    log_probs: Vec<f32>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    values: Vec<f32>,
    next_states: Vec<f32>,
    _state_dim: usize,
}
impl RolloutBatch {
    fn new(capacity: usize, state_dim: usize) -> Self {
        Self {
            states: Vec::with_capacity(capacity * state_dim),
            actions: Vec::with_capacity(capacity),
            log_probs: Vec::with_capacity(capacity),
            rewards: Vec::with_capacity(capacity),
            dones: Vec::with_capacity(capacity),
            values: Vec::with_capacity(capacity),
            next_states: Vec::with_capacity(capacity * state_dim),
            _state_dim: state_dim,
        }
    }

    #[allow(clippy::too_many_arguments)]
    fn push(&mut self, s: &[f32], a: f32, lp: f32, r: f32, d: f32, v: f32, s2: &[f32]) {
        self.states.extend_from_slice(s);
        self.actions.push(a);
        self.log_probs.push(lp);
        self.rewards.push(r);
        self.dones.push(d);
        self.values.push(v);
        self.next_states.extend_from_slice(s2);
    }

    fn len(&self) -> usize {
        self.actions.len()
    }
}

// -------------------------------
// Math helpers
// -------------------------------

fn gaussian_log_prob(action: &Tensor, mean: &Tensor, log_std: &Tensor) -> Tensor {
    // All tensors shaped [B, A] (log_std is broadcastable)
    let std = log_std.exp();
    let var = std.pow_scalar(2.0);
    let log_scale = log_std;
    let diff = action.sub_tensor(mean);
    let log_prob = diff
        .pow_scalar(2.0)
        .div_tensor(&var)
        .add_scalar(std::f32::consts::LN_2 + std::f32::consts::PI)
        .add_tensor(&log_scale.mul_scalar(2.0))
        .mul_scalar(0.5)
        .mul_scalar(-1.0);
    // Sum across action dim (dim=1) -> [B,1]
    log_prob.sum_dims(&[1], true)
}

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

examples/RL_training/td3.rs (line 250)

    fn state_tensor(&self) -> Tensor {
        // Normalize to keep critic inputs bounded:
        // - Position is bounded by termination at |pos|>3 → scale by 3 to [-1,1]
        // - Velocity scaled by 1.0 and clamped to [-1,1]
        let pos_n = self.pos / 3.0;
        let vel_n = self.vel.clamp(-1.0, 1.0);
        Tensor::from_slice(&[pos_n, vel_n, 0.0], vec![1, 3]).unwrap()
    }

    fn step(&mut self, action_value: f32) -> (Tensor, f32, bool) {
        let a = action_value.clamp(-1.0, 1.0);
        self.vel += 0.1 * a - 0.01 * self.pos;
        self.pos += self.vel;
        self.steps += 1;

        let reward = -(self.pos * self.pos) - 0.1 * (a * a);
        let done = self.pos.abs() > 3.0 || self.steps >= self.max_steps;
        (self.state_tensor(), reward, done)
    }
}

// -------------------------------
// Replay Buffer
// -------------------------------

struct ReplayBuffer {
    capacity: usize,
    size: usize,
    pos: usize,
    state_dim: usize,
    action_dim: usize,
    states: Vec<f32>,
    actions: Vec<f32>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    next_states: Vec<f32>,
}

impl ReplayBuffer {
    fn new(capacity: usize, state_dim: usize, action_dim: usize) -> Self {
        Self {
            capacity,
            size: 0,
            pos: 0,
            state_dim,
            action_dim,
            states: vec![0.0; capacity * state_dim],
            actions: vec![0.0; capacity * action_dim],
            rewards: vec![0.0; capacity],
            dones: vec![0.0; capacity],
            next_states: vec![0.0; capacity * state_dim],
        }
    }

    fn push(&mut self, s: &[f32], a: &[f32], r: f32, d: f32, s2: &[f32]) {
        let i = self.pos;
        let so = i * self.state_dim;
        let ao = i * self.action_dim;
        self.states[so..so + self.state_dim].copy_from_slice(s);
        self.actions[ao..ao + self.action_dim].copy_from_slice(a);
        self.rewards[i] = r;
        self.dones[i] = d;
        self.next_states[so..so + self.state_dim].copy_from_slice(s2);

        self.pos = (self.pos + 1) % self.capacity;
        self.size = self.size.saturating_add(1).min(self.capacity);
    }

    fn can_sample(&self, batch_size: usize) -> bool {
        self.size >= batch_size
    }

    fn sample(
        &self,
        batch_size: usize,
        rng: &mut SmallRng,
    ) -> (Tensor, Tensor, Tensor, Tensor, Tensor) {
        let mut s_vec = Vec::with_capacity(batch_size * self.state_dim);
        let mut a_vec = Vec::with_capacity(batch_size * self.action_dim);
        let mut r_vec = Vec::with_capacity(batch_size);
        let mut d_vec = Vec::with_capacity(batch_size);
        let mut s2_vec = Vec::with_capacity(batch_size * self.state_dim);

        for _ in 0..batch_size {
            let idx = rng.sample_index(self.size);
            let so = idx * self.state_dim;
            let ao = idx * self.action_dim;
            s_vec.extend_from_slice(&self.states[so..so + self.state_dim]);
            a_vec.extend_from_slice(&self.actions[ao..ao + self.action_dim]);
            r_vec.push(self.rewards[idx]);
            d_vec.push(self.dones[idx]);
            s2_vec.extend_from_slice(&self.next_states[so..so + self.state_dim]);
        }

        let s = Tensor::from_slice(&s_vec, vec![batch_size, self.state_dim]).unwrap();
        let a = Tensor::from_slice(&a_vec, vec![batch_size, self.action_dim]).unwrap();
        let r = Tensor::from_slice(&r_vec, vec![batch_size, 1]).unwrap();
        let d = Tensor::from_slice(&d_vec, vec![batch_size, 1]).unwrap();
        let s2 = Tensor::from_slice(&s2_vec, vec![batch_size, self.state_dim]).unwrap();
        (s, a, r, d, s2)
    }
}

// -------------------------------
// Helper: gradient clipping by global norm
// -------------------------------

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    // Compute global L2 norm of all grads
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                let scaled = g.mul_scalar(scale);
                p.set_grad(scaled);
            }
        }
    }
}

// Compute global L2 norm of gradients across a parameter list (read-only)
fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// Compute L2 norm of parameters (weights/biases) across a parameter list
fn params_l2_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let _ng = NoGradTrack::new();
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        for &v in p.data() {
            total_sq += v * v;
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main: TD3 training on YardEnv
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== TD3 Example (YardEnv) ===");

    // Environment / problem dims
    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hyperparameters (small for demo)
    let gamma = 0.99f32;
    let tau = 0.005f32; // Polyak
    let policy_noise = 0.2f32; // target smoothing noise stddev
    let exploration_noise = 0.1f32; // behavior policy noise stddev
    let policy_delay = 2usize;
    let batch_size = 64usize;
    let start_steps = 500usize; // random exploration steps
    let total_steps = 1500usize;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(11));
    let mut actor_targ = Actor::new(state_dim, action_dim, Some(12));
    actor_targ.net.copy_from(&actor.net);
    actor_targ.set_requires_grad_all(false);

    let mut critic1 = Critic::new(state_dim, action_dim, Some(21));
    let mut critic2 = Critic::new(state_dim, action_dim, Some(22));
    let mut critic1_targ = Critic::new(state_dim, action_dim, Some(23));
    let mut critic2_targ = Critic::new(state_dim, action_dim, Some(24));
    critic1_targ.net.copy_from(&critic1.net);
    critic2_targ.net.copy_from(&critic2.net);
    critic1_targ.set_requires_grad_all(false);
    critic2_targ.set_requires_grad_all(false);

    // Optimizers
    let mut actor_opt = Adam::with_learning_rate(1e-3);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }

    let mut critic_opt = Adam::with_learning_rate(1e-4);
    for p in critic1.parameters() {
        critic_opt.add_parameter(p);
    }
    for p in critic2.parameters() {
        critic_opt.add_parameter(p);
    }

    // Replay buffer and env
    let mut rb = ReplayBuffer::new(100_000, state_dim, action_dim);
    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(987654321);

    // Reset & metric trackers
    let mut state = env.reset(); // [1, state_dim]
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32; // smooth short-term
    let mut best_return = f32::NEG_INFINITY;
    let mut policy_updates: usize = 0;

    for t in 0..total_steps {
        // Select action
        let action_tensor = if t < start_steps {
            let a = rng.uniform(-1.0, 1.0);
            Tensor::from_slice(&[a], vec![1, action_dim]).unwrap()
        } else {
            // Behavior policy with exploration noise
            let _ng = NoGradTrack::new();
            let det = actor.forward(&state);
            let noise = Tensor::randn(vec![1, action_dim], None).mul_scalar(exploration_noise);
            tanh_bounded(&det.add_tensor(&noise))
        };
        let action_value = action_tensor.data()[0];

        // Environment step
        let (next_state, reward, done) = env.step(action_value);
        episode_return += reward;

        // Store transition
        let s_slice = state.data().to_vec();
        let a_slice = action_tensor.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            &a_slice,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        state = if done {
            let st = env.reset();
            // Metrics: update EMA and best
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={} | policy_updates={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size,
                policy_updates
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Training
        if rb.can_sample(batch_size) {
            // Sample batch
            let (s, a, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Compute target values y = r + (1-d)*gamma*min(Q1', Q2') using target networks (no grad)
            let target_q = {
                let _ng = NoGradTrack::new();
                // Target actions with smoothing noise (tanh bounds)
                let noise =
                    Tensor::randn(vec![batch_size, action_dim], None).mul_scalar(policy_noise);
                let a_targ = tanh_bounded(&actor_targ.forward(&s2).add_tensor(&noise));
                let q1_t = critic1_targ.forward(&s2, &a_targ);
                let q2_t = critic2_targ.forward(&s2, &a_targ);

                // Elementwise min via data() since this path is no-grad
                let q1d = q1_t.data();
                let q2d = q2_t.data();
                let mut min_vec = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let v1 = q1d[i];
                    let v2 = q2d[i];
                    min_vec.push(v1.min(v2));
                }
                let min_q = Tensor::from_slice(&min_vec, vec![batch_size, 1]).unwrap();
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&min_q))
            };

            // Critic update (both critics)
            // Zero grads in a short scope, then drop borrows before forward
            {
                let mut params = {
                    let c_params = critic1.parameters();
                    let c2_params = critic2.parameters();
                    let mut tmp: Vec<&mut Tensor> = Vec::new();
                    tmp.extend(c_params);
                    tmp.extend(c2_params);
                    tmp
                };
                critic_opt.zero_grad(&mut params);
            }

            // Forward current Q estimates
            let q1 = critic1.forward(&s, &a);
            let q2 = critic2.forward(&s, &a);
            let diff1 = q1.sub_tensor(&target_q);
            let diff2 = q2.sub_tensor(&target_q);
            let mut critic_loss = diff1
                .pow_scalar(2.0)
                .mean()
                .add_tensor(&diff2.pow_scalar(2.0).mean());

            // Backward
            critic_loss.backward(None);

            // Optional gradient clipping + step (only for params that received grads)
            {
                let params = {
                    let c_params = critic1.parameters();
                    let c2_params = critic2.parameters();
                    let mut tmp: Vec<&mut Tensor> = Vec::new();
                    tmp.extend(c_params);
                    tmp.extend(c2_params);
                    tmp
                };
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    // Pre-step metrics
                    let grad_norm_before = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    critic_opt.step(&mut with_grads);
                    critic_opt.zero_grad(&mut with_grads);

                    // Post-step metrics (param norm)
                    let mut for_norm_params = {
                        let c_params = critic1.parameters();
                        let c2_params = critic2.parameters();
                        let mut tmp: Vec<&mut Tensor> = Vec::new();
                        tmp.extend(c_params);
                        tmp.extend(c2_params);
                        tmp
                    };
                    let param_norm = params_l2_norm(&mut for_norm_params);

                    // Print compact critic metrics occasionally
                    if t % 100 == 0 {
                        let q1_mean = q1.mean().value();
                        let q2_mean = q2.mean().value();
                        let tq_mean = target_q.mean().value();
                        println!(
                            "t={:5} | critic_loss={:.4} | q1_mean={:.3} q2_mean={:.3} tq_mean={:.3} | grad_norm={:.3} | crit_param_norm={:.3}",
                            t,
                            critic_loss.value(),
                            q1_mean,
                            q2_mean,
                            tq_mean,
                            grad_norm_before,
                            param_norm
                        );
                    }
                }
            }

            // Delayed policy update
            if t % policy_delay == 0 {
                // Actor update: maximize Q1(s, actor(s)) -> minimize -Q1
                // Zero actor grads before backward
                {
                    let mut a_params: Vec<&mut Tensor> = actor.parameters();
                    actor_opt.zero_grad(&mut a_params);
                }

                let a_pred = actor.forward(&s);
                let q_for_actor = critic1.forward(&s, &a_pred);
                let mut actor_loss = q_for_actor.mul_scalar(-1.0).mean();
                actor_loss.backward(None);

                {
                    let a_params: Vec<&mut Tensor> = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in a_params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let grad_norm_before = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);

                        // Post-step param norm
                        let mut for_norm_params = actor.parameters();
                        let param_norm = params_l2_norm(&mut for_norm_params);

                        policy_updates += 1;
                        if t % 200 == 0 {
                            println!(
                                "t={:5} | actor_loss={:.4} | act_grad_norm={:.3} | act_param_norm={:.3} | lr_a={:.4e} lr_c={:.4e} | policy_updates={}",
                                t,
                                actor_loss.value(),
                                grad_norm_before,
                                param_norm,
                                actor_opt.learning_rate(),
                                critic_opt.learning_rate(),
                                policy_updates
                            );
                        }
                    }
                }

                // Target updates (Polyak averaging, no grad)
                actor_targ.net.soft_update_from(&actor.net, tau);
                critic1_targ.net.soft_update_from(&critic1.net, tau);
                critic2_targ.net.soft_update_from(&critic2.net, tau);
            }

            // Clear entire graphs to avoid stale accumulation across iterations
            clear_all_graphs_known();
        }
    }

    println!("=== TD3 training finished ===");
    Ok(())
}

examples/getting_started/tensor_operators.rs (line 49)

fn demonstrate_basic_operators() {
    println!("--- Basic Tensor-Tensor Operators ---");

    let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();

    println!("Tensor A: {:?}", a.data());
    println!("Tensor B: {:?}", b.data());

    // Addition
    let c = &a + &b;
    println!("A + B: {:?}", c.data());

    // Subtraction
    let d = &a - &b;
    println!("A - B: {:?}", d.data());

    // Multiplication
    let e = &a * &b;
    println!("A * B: {:?}", e.data());

    // Division
    let f = &a / &b;
    println!("A / B: {:?}", f.data());
}

/// Demonstrate tensor-scalar operators
fn demonstrate_scalar_operators() {
    println!("\n--- Tensor-Scalar Operators ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    println!("Original tensor: {:?}", tensor.data());

    // Tensor + scalar
    let result1 = &tensor + 5.0;
    println!("Tensor + 5.0: {:?}", result1.data());

    // Scalar + tensor
    let result2 = 5.0 + &tensor;
    println!("5.0 + Tensor: {:?}", result2.data());

    // Tensor - scalar
    let result3 = &tensor - 2.0;
    println!("Tensor - 2.0: {:?}", result3.data());

    // Tensor * scalar
    let result4 = &tensor * 3.0;
    println!("Tensor * 3.0: {:?}", result4.data());

    // Scalar * tensor
    let result5 = 3.0 * &tensor;
    println!("3.0 * Tensor: {:?}", result5.data());

    // Tensor / scalar
    let result6 = &tensor / 2.0;
    println!("Tensor / 2.0: {:?}", result6.data());
}

/// Demonstrate assignment operators
fn demonstrate_operator_assignment() {
    println!("\n--- Assignment Operators ---");

    let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    println!("Original tensor: {:?}", tensor.data());

    // In-place addition
    tensor += 5.0;
    println!("After += 5.0: {:?}", tensor.data());

    // In-place subtraction
    tensor -= 2.0;
    println!("After -= 2.0: {:?}", tensor.data());

    // In-place multiplication
    tensor *= 3.0;
    println!("After *= 3.0: {:?}", tensor.data());

    // In-place division
    tensor /= 2.0;
    println!("After /= 2.0: {:?}", tensor.data());
}

/// Demonstrate operator chaining and complex expressions
fn demonstrate_operator_chaining() {
    println!("\n--- Operator Chaining ---");

    let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();
    let c = Tensor::from_slice(&[9.0, 10.0, 11.0, 12.0], vec![2, 2]).unwrap();

    println!("Tensor A: {:?}", a.data());
    println!("Tensor B: {:?}", b.data());
    println!("Tensor C: {:?}", c.data());

    // Complex expression: (A + B) * C - 5
    let result = (&a + &b) * &c - 5.0;
    println!("(A + B) * C - 5: {:?}", result.data());

    // Another complex expression: A * 2 + B / 2
    let result2 = &a * 2.0 + &b / 2.0;
    println!("A * 2 + B / 2: {:?}", result2.data());

    // Negation and addition: -A + B * C
    let result3 = -&a + &b * &c;
    println!("-A + B * C: {:?}", result3.data());

    // Division with parentheses: (A + B) / (C - 1)
    let result4 = (&a + &b) / (&c - 1.0);
    println!("(A + B) / (C - 1): {:?}", result4.data());
}

/// Demonstrate broadcasting behavior
fn demonstrate_broadcasting() {
    println!("\n--- Broadcasting ---");

    // 2D tensor
    let tensor_2d = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    println!(
        "2D tensor: shape {:?}, data: {:?}",
        tensor_2d.shape().dims(),
        tensor_2d.data()
    );

    // 1D tensor (will be broadcasted)
    let tensor_1d = Tensor::from_slice(&[10.0, 20.0], vec![2]).unwrap();
    println!(
        "1D tensor: shape {:?}, data: {:?}",
        tensor_1d.shape().dims(),
        tensor_1d.data()
    );

    // Broadcasting addition
    let broadcast_sum = &tensor_2d + &tensor_1d;
    println!(
        "Broadcast sum: shape {:?}, data: {:?}",
        broadcast_sum.shape().dims(),
        broadcast_sum.data()
    );

    // Broadcasting multiplication
    let broadcast_mul = &tensor_2d * &tensor_1d;
    println!(
        "Broadcast multiplication: shape {:?}, data: {:?}",
        broadcast_mul.shape().dims(),
        broadcast_mul.data()
    );

    // Broadcasting with scalar
    let broadcast_scalar = &tensor_2d + 100.0;
    println!(
        "Broadcast scalar: shape {:?}, data: {:?}",
        broadcast_scalar.shape().dims(),
        broadcast_scalar.data()
    );
}

/// Demonstrate equivalence between operators and method calls
fn demonstrate_method_equivalence() {
    println!("\n--- Operator vs Method Call Equivalence ---");

    let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();

    // Addition: operator vs method
    let operator_result = &a + &b;
    let method_result = a.add_tensor(&b);

    println!("A + B (operator): {:?}", operator_result.data());
    println!("A.add_tensor(B): {:?}", method_result.data());
    println!(
        "Results are equal: {}",
        operator_result.data() == method_result.data()
    );

    // Multiplication: operator vs method
    let operator_result = &a * &b;
    let method_result = a.mul_tensor(&b);

    println!("A * B (operator): {:?}", operator_result.data());
    println!("A.mul_tensor(B): {:?}", method_result.data());
    println!(
        "Results are equal: {}",
        operator_result.data() == method_result.data()
    );

    // Scalar addition: operator vs method
    let operator_result = &a + 5.0;
    let method_result = a.add_scalar(5.0);

    println!("A + 5.0 (operator): {:?}", operator_result.data());
    println!("A.add_scalar(5.0): {:?}", method_result.data());
    println!(
        "Results are equal: {}",
        operator_result.data() == method_result.data()
    );
}

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 175)

fn demonstrate_forward_pass() {
    println!("\n--- Forward Pass (with gradients) ---");

    let layer = LinearLayer::new(3, 2, Some(43));

    // Single input
    let input = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3]).unwrap();
    let output = layer.forward(&input);

    println!("Single input:");
    println!("  Input: {:?}", input.data());
    println!("  Output: {:?}", output.data());
    println!("  Output requires grad: {}", output.requires_grad());

    // Batch input
    let batch_input = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
    let batch_output = layer.forward(&batch_input);

    println!("Batch input:");
    println!("  Input shape: {:?}", batch_input.shape().dims());
    println!("  Output shape: {:?}", batch_output.shape().dims());
    println!("  Output requires grad: {}", batch_output.requires_grad());
}

/// Demonstrate forward pass without gradient tracking
fn demonstrate_forward_pass_no_grad() {
    println!("\n--- Forward Pass (no gradients) ---");

    let layer = LinearLayer::new(3, 2, Some(44));

    // Single input
    let input = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3]).unwrap();
    let output = layer.forward_no_grad(&input);

    println!("Single input (no grad):");
    println!("  Input: {:?}", input.data());
    println!("  Output: {:?}", output.data());
    println!("  Output requires grad: {}", output.requires_grad());

    // Compare with grad version
    let output_with_grad = layer.forward(&input);
    println!("Comparison:");
    println!(
        "  Same values: {}",
        output.data() == output_with_grad.data()
    );
    println!("  No grad requires grad: {}", output.requires_grad());
    println!(
        "  With grad requires grad: {}",
        output_with_grad.requires_grad()
    );
}

/// Demonstrate complete training loop
fn demonstrate_training_loop() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Loop ---");

    // Create layer and training data
    let mut layer = LinearLayer::new(2, 1, Some(45));

    // Simple regression task: y = 2*x1 + 3*x2 + 1
    let x_data = Tensor::from_slice(
        &[
            1.0, 1.0, // x1=1, x2=1 -> y=6
            2.0, 1.0, // x1=2, x2=1 -> y=8
            1.0, 2.0, // x1=1, x2=2 -> y=9
            2.0, 2.0, // x1=2, x2=2 -> y=11
        ],
        vec![4, 2],
    )
    .unwrap();

    let y_true = Tensor::from_slice(&[6.0, 8.0, 9.0, 11.0], vec![4, 1]).unwrap();

    println!("Training data:");
    println!("  X shape: {:?}", x_data.shape().dims());
    println!("  Y shape: {:?}", y_true.shape().dims());
    println!("  Target function: y = 2*x1 + 3*x2 + 1");

    // Create optimizer
    let config = AdamConfig {
        learning_rate: 0.01,
        beta1: 0.9,
        beta2: 0.999,
        eps: 1e-8,
        weight_decay: 0.0,
        amsgrad: false,
    };

    let mut optimizer = Adam::with_config(config);
    let params = layer.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    println!("Optimizer setup complete. Starting training...");

    // Training loop
    let num_epochs = 100;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = layer.forward(&x_data);

        // Compute loss: MSE
        let diff = y_pred.sub_tensor(&y_true);
        let mut loss = diff.pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        let mut params = layer.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        losses.push(loss.value());

        // Print progress
        if epoch % 20 == 0 || epoch == num_epochs - 1 {
            println!("Epoch {:3}: Loss = {:.6}", epoch, loss.value());
        }
    }

    // Evaluate final model
    let final_predictions = layer.forward_no_grad(&x_data);

    println!("\nFinal model evaluation:");
    println!("  Learned weights: {:?}", layer.weight.data());
    println!("  Learned bias: {:?}", layer.bias.data());
    println!("  Target weights: [2.0, 3.0]");
    println!("  Target bias: [1.0]");

    println!("  Predictions vs True:");
    for i in 0..4 {
        let pred = final_predictions.data()[i];
        let true_val = y_true.data()[i];
        println!(
            "    Sample {}: pred={:.3}, true={:.1}, error={:.3}",
            i + 1,
            pred,
            true_val,
            (pred - true_val).abs()
        );
    }

    // Training analysis
    let initial_loss = losses[0];
    let final_loss = losses[losses.len() - 1];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);

    Ok(())
}

/// Demonstrate single vs batch inference
fn demonstrate_single_vs_batch_inference() {
    println!("\n--- Single vs Batch Inference ---");

    let layer = LinearLayer::new(4, 3, Some(46));

    // Single inference
    println!("Single inference:");
    let single_input = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![1, 4]).unwrap();
    let single_output = layer.forward_no_grad(&single_input);
    println!("  Input shape: {:?}", single_input.shape().dims());
    println!("  Output shape: {:?}", single_output.shape().dims());
    println!("  Output: {:?}", single_output.data());

    // Batch inference
    println!("Batch inference:");
    let batch_input = Tensor::from_slice(
        &[
            1.0, 2.0, 3.0, 4.0, // Sample 1
            5.0, 6.0, 7.0, 8.0, // Sample 2
            9.0, 10.0, 11.0, 12.0, // Sample 3
        ],
        vec![3, 4],
    )
    .unwrap();
    let batch_output = layer.forward_no_grad(&batch_input);
    println!("  Input shape: {:?}", batch_input.shape().dims());
    println!("  Output shape: {:?}", batch_output.shape().dims());

    // Verify batch consistency - first sample should match single inference
    let _first_batch_sample = batch_output.view(vec![3, 3]); // Reshape to access first sample
    let first_sample_data = &batch_output.data()[0..3]; // First 3 elements
    let single_sample_data = single_output.data();

    println!("Consistency check:");
    println!("  Single output: {:?}", single_sample_data);
    println!("  First batch sample: {:?}", first_sample_data);
    println!(
        "  Match: {}",
        single_sample_data
            .iter()
            .zip(first_sample_data.iter())
            .all(|(a, b)| (a - b).abs() < 1e-6)
    );
}

/// Demonstrate serialization and loading
fn demonstrate_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Serialization ---");

    // Create and train a simple layer
    let mut original_layer = LinearLayer::new(2, 1, Some(47));

    // Simple training data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let y_true = Tensor::from_slice(&[5.0, 11.0], vec![2, 1]).unwrap();

    let mut optimizer = Adam::with_learning_rate(0.01);
    let params = original_layer.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    // Train for a few epochs
    for _ in 0..10 {
        let y_pred = original_layer.forward(&x_data);
        let mut loss = (y_pred.sub_tensor(&y_true)).pow_scalar(2.0).mean();
        loss.backward(None);

        let mut params = original_layer.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);
    }

    println!("Original layer trained");
    println!("  Weight: {:?}", original_layer.weight.data());
    println!("  Bias: {:?}", original_layer.bias.data());

    // Save layer
    original_layer.save_json("temp_linear_layer")?;

    // Load layer
    let loaded_layer = LinearLayer::load_json("temp_linear_layer", 2, 1)?;

    println!("Loaded layer");
    println!("  Weight: {:?}", loaded_layer.weight.data());
    println!("  Bias: {:?}", loaded_layer.bias.data());

    // Verify consistency
    let test_input = Tensor::from_slice(&[1.0, 1.0], vec![1, 2]).unwrap();
    let original_output = original_layer.forward_no_grad(&test_input);
    let loaded_output = loaded_layer.forward_no_grad(&test_input);

    println!("Consistency check:");
    println!("  Original output: {:?}", original_output.data());
    println!("  Loaded output: {:?}", loaded_output.data());
    println!(
        "  Match: {}",
        original_output
            .data()
            .iter()
            .zip(loaded_output.data().iter())
            .all(|(a, b)| (a - b).abs() < 1e-6)
    );

    println!("Serialization verification: PASSED");

    Ok(())
}

Additional examples can be found in:

• examples/iterators/element_iteration.rs
• examples/getting_started/tensor_basics.rs
• examples/neural_networks/feedforward_network.rs
• examples/getting_started/serialization_basics.rs
• examples/optimizers/adam_configurations.rs
• examples/optimizers/learning_rate_scheduling.rs
• examples/getting_started/optimizer_basics.rs
• examples/iterators/advanced_patterns.rs
• examples/iterators/performance_optimization.rs
• examples/supervised_training/supervised_bce.rs
• examples/supervised_training/supervised_classification.rs
• examples/supervised_training/supervised_regression.rs

impl Tensor

pub fn randn(shape_dims: Vec<usize>, seed: Option<u64>) -> Self

Creates a tensor with normally distributed random values (mean=0, std=1)

Similar to PyTorch’s torch.randn(), creates a tensor filled with random values drawn from a standard normal distribution (mean=0, standard deviation=1). Uses Box-Muller transform for efficient normal distribution generation.

This method provides high-quality random number generation with optional reproducibility through seed-based generation. The generated values follow a standard normal distribution suitable for machine learning applications.

Arguments

• shape_dims - Vector of dimension sizes defining the tensor shape
• seed - Optional seed for reproducible random generation

Returns

A new tensor with normally distributed random values

Performance

• Box-Muller Transform: Efficient normal distribution generation
• SIMD Optimization: Vectorized operations for large tensors
• Memory Efficient: Single-pass generation with optimized allocation
• Thread Safe: Uses thread-local random state

Examples

Basic Usage

use train_station::Tensor;

// Create a 2x3 tensor with random normal values
let tensor = Tensor::randn(vec![2, 3], None);
assert_eq!(tensor.size(), 6);
assert_eq!(tensor.shape().dims(), vec![2, 3]);

// Verify random values are generated
let first_value = tensor.get(&[0, 0]);
assert!(first_value != 0.0); // Should be random

Reproducible Generation

use train_station::Tensor;

// Create with fixed seed for reproducible results
let tensor1 = Tensor::randn(vec![100], Some(42));
let tensor2 = Tensor::randn(vec![100], Some(42));

// tensor1 and tensor2 will have identical values
for i in 0..tensor1.size() {
    assert!((tensor1.get(&[i]) - tensor2.get(&[i])).abs() < 1e-6);
}

Statistical Properties

use train_station::Tensor;

// Generate large tensor for statistical analysis
let tensor = Tensor::randn(vec![1000], Some(42));
assert_eq!(tensor.size(), 1000);

// Check that values are reasonable (within 4 standard deviations)
let mut min_val = f32::INFINITY;
let mut max_val = f32::NEG_INFINITY;
let mut sum = 0.0;

for i in 0..tensor.size() {
    let val = tensor.get(&[i]);
    min_val = min_val.min(val);
    max_val = max_val.max(val);
    sum += val;
}

let mean = sum / tensor.size() as f32;

// Mean should be close to 0, values should be within reasonable bounds
assert!(mean.abs() < 0.1, "Mean should be close to 0, got {}", mean);
assert!(min_val > -4.0, "Values should not be too negative, min: {}", min_val);
assert!(max_val < 4.0, "Values should not be too positive, max: {}", max_val);

Zero-Sized Tensors

use train_station::Tensor;

// Handle empty tensors gracefully
let tensor = Tensor::randn(vec![0], Some(42));
assert_eq!(tensor.size(), 0);
assert_eq!(tensor.shape().dims(), vec![0]);

Implementation Details

This method uses the Box-Muller transform to generate normally distributed random variables from uniform random variables. The process involves:

1. Random Number Generation: Uses Xorshift algorithm for uniform random numbers
2. Box-Muller Transform: Converts uniform random variables to normal distribution
3. SIMD Optimization: Vectorized operations for large tensors when available
4. Numerical Stability: Robust handling of edge cases and potential NaN values

The Box-Muller transform ensures that the generated values follow a true normal distribution with mean=0 and standard deviation=1, making it suitable for machine learning applications requiring normally distributed random values.

Examples found in repository

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 57)

    pub fn new(input_size: usize, output_size: usize, seed: Option<u64>) -> Self {
        // Xavier/Glorot initialization: scale by sqrt(1/input_size)
        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,
        }
    }

examples/neural_networks/basic_encoder.rs (line 81)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Encoder Example ===");

    let batch = 2usize;
    let seq = 6usize;
    let embed = 32usize;
    let heads = 4usize;

    let input = Tensor::randn(vec![batch, seq, embed], Some(11));
    let mut enc = EncoderBlock::new(embed, heads, Some(123));

    // Example: no mask (set Some(mask) to use masking)
    let out = enc.forward(&input, None);
    println!("Output shape: {:?}", out.shape().dims());

    // Verify gradients/optimization
    let mut opt = Adam::with_learning_rate(0.01);
    let mut params = enc.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/neural_networks/basic_decoder.rs (line 93)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Decoder Example ===");

    let batch = 2usize;
    let src = 7usize;
    let tgt = 5usize;
    let embed = 32usize;
    let heads = 4usize;

    let memory = Tensor::randn(vec![batch, src, embed], Some(21));
    let tgt_in = Tensor::randn(vec![batch, tgt, embed], Some(22));

    let mut dec = DecoderBlock::new(embed, heads, Some(456));
    let out = dec.forward(&tgt_in, &memory, None, None);
    println!("Output shape: {:?}", out.shape().dims());

    let mut opt = Adam::with_learning_rate(0.01);
    let mut params = dec.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/getting_started/tensor_basics.rs (line 80)

fn demonstrate_tensor_creation() {
    println!("--- Tensor Creation ---");

    // Create tensors with different initializations
    let zeros = Tensor::zeros(vec![2, 3]);
    println!(
        "Zeros tensor: shape {:?}, data: {:?}",
        zeros.shape().dims(),
        zeros.data()
    );

    let ones = Tensor::ones(vec![3, 2]);
    println!(
        "Ones tensor: shape {:?}, data: {:?}",
        ones.shape().dims(),
        ones.data()
    );

    // Create tensor from slice
    let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let from_slice = Tensor::from_slice(&data, vec![2, 3]).unwrap();
    println!(
        "From slice: shape {:?}, data: {:?}",
        from_slice.shape().dims(),
        from_slice.data()
    );

    // Create tensor with specific value
    let mut filled = Tensor::new(vec![2, 2]);
    {
        let data = filled.data_mut();
        for value in data.iter_mut() {
            *value = 42.0;
        }
    }
    println!("Filled with 42: {:?}", filled.data());

    // Create tensor with random data
    let random = Tensor::randn(vec![2, 2], Some(42));
    println!(
        "Random tensor: shape {:?}, data: {:?}",
        random.shape().dims(),
        random.data()
    );
}

examples/neural_networks/feedforward_network.rs (line 385)

fn demonstrate_configurable_architectures() {
    println!("\n--- Configurable Architectures ---");

    let architectures = vec![
        ("Shallow", vec![8]),
        ("Medium", vec![16, 8]),
        ("Deep", vec![32, 16, 8, 4]),
        ("Wide", vec![64, 32]),
        ("Bottleneck", vec![16, 4, 16]),
    ];

    for (name, hidden_sizes) in architectures {
        let config = FeedForwardConfig {
            input_size: 10,
            hidden_sizes,
            output_size: 3,
            use_bias: true,
        };

        let network = FeedForwardNetwork::new(config.clone(), Some(44));

        // Test forward pass
        let test_input = Tensor::randn(vec![5, 10], Some(45)); // Batch of 5
        let output = network.forward_no_grad(&test_input);

        println!("{} network:", name);
        println!("  Architecture: 10 -> {:?} -> 3", config.hidden_sizes);
        println!("  Parameters: {}", network.parameter_count());
        println!("  Test output shape: {:?}", output.shape().dims());
        println!(
            "  Output range: [{:.3}, {:.3}]",
            output.data().iter().fold(f32::INFINITY, |a, &b| a.min(b)),
            output
                .data()
                .iter()
                .fold(f32::NEG_INFINITY, |a, &b| a.max(b))
        );
    }
}

examples/neural_networks/basic_transformer.rs (line 241)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Transformer Example ===");

    let batch = 2usize;
    let src_len = 8usize;
    let tgt_len = 6usize;
    let embed = 32usize;
    let heads = 4usize;
    let layers = 2usize;

    let src = Tensor::randn(vec![batch, src_len, embed], Some(1001));
    let tgt = Tensor::randn(vec![batch, tgt_len, embed], Some(1002));

    let mut trf = BasicTransformer::new(embed, heads, layers, Some(999));
    let out = trf.forward(&src, &tgt);
    println!("Output shape: {:?}", out.shape().dims());

    // Quick optimization step
    let mut opt = Adam::with_learning_rate(0.005);
    let mut params = trf.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());

    // Demo: non auto-regressive inference (single pass)
    let nar = trf.infer_non_autoregressive(&src, tgt_len);
    println!("NAR output shape: {:?}", nar.shape().dims());

    // Demo: auto-regressive inference (toy)
    let ar = trf.infer_autoregressive(&src, 3);
    println!("AR output shape: {:?}", ar.shape().dims());

    // NAR training demo
    let nar_tgt = tgt.clone();
    trf.train_non_autoregressive_steps(&src, &nar_tgt, 3, 0.01);

    // AR training demo (teacher-forced)
    let ar_tgt = tgt.clone();
    trf.train_autoregressive_steps(&src, &ar_tgt, 3, 0.01);
    println!("=== Done ===");
    Ok(())
}

Additional examples can be found in:

• examples/getting_started/optimizer_basics.rs
• examples/getting_started/serialization_basics.rs
• examples/neural_networks/multi_head_attention.rs
• examples/optimizers/adam_configurations.rs
• examples/optimizers/learning_rate_scheduling.rs
• examples/RL_training/td3.rs

pub fn fill_randn(&mut self, seed: Option<u64>)

Fills the tensor with normally distributed random values

Internal method that fills an existing tensor with random values from a standard normal distribution. Uses Box-Muller transform for efficiency and provides SIMD optimization for large tensors.

This method is used internally by randn() and provides the core random number generation functionality with optimized performance characteristics.

Arguments

• seed - Optional seed for reproducible random generation

Performance

• Box-Muller Transform: Generates pairs of normal random variables
• SIMD Optimization: Vectorized operations when possible
• Memory Efficient: Single-pass generation
• Unrolled Loops: 4x unrolling for better instruction throughput

Implementation Details

The method performs the following steps:

1. Zero-sized Check: Returns early for empty tensors
2. RNG Initialization: Creates Xorshift RNG with seed or system time
3. SIMD Detection: Checks for AVX2 availability for optimized path
4. Generation: Uses SIMD or scalar path based on hardware support
5. Completion: Fills all tensor elements with normal random values

The method automatically handles hardware capabilities and falls back to scalar operations when SIMD is not available, ensuring compatibility across different CPU architectures.

impl Tensor

pub fn chunks(&self, chunk_size: usize) -> TensorChunksIterator<'_>

Standard slice-like chunks iterator. Use this instead of iter_chunks.

Iterates over contiguous or view-backed slices of the tensor with the specified chunk size. In no-grad fast mode, a single contiguous owner may be materialized to optimize subsequent views.

Arguments

• chunk_size - Number of elements per chunk (must be > 0)

Examples

use train_station::tensor::TensorCollectExt;
use train_station::Tensor;

let t = Tensor::from_slice(&(1..=6).map(|i| i as f32).collect::<Vec<_>>(), vec![6]).unwrap();
let y = t.chunks(2).map(|c| c.mul_scalar(2.0)).collect_shape(vec![6]);
assert_eq!(y.data(), &[2.0, 4.0, 6.0, 8.0, 10.0, 12.0]);

Examples found in repository

examples/iterators/performance_optimization.rs (line 178)

fn demonstrate_memory_optimization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Memory Optimization ---");

    // Create a large tensor for memory testing
    let size = 10000;
    let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
    let tensor = Tensor::from_slice(&data, vec![size])?;

    println!("Processing tensor of size: {}", size);

    // Pattern 1: Streaming processing with iterator chunks (process in blocks, collect with shape)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();
    let flattened = tensor.view(vec![size as i32]);
    let _streamed_result: Tensor = flattened
        .chunks(chunk_size)
        .map(|c| c.pow_scalar(2.0).sqrt())
        .collect_shape(vec![size]);
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter_elements()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect_shape(vec![size]);
    let full_time = start.elapsed();

    println!("  Streaming time: {:?}", streamed_time);
    println!("  Full processing time: {:?}", full_time);
    println!(
        "  Memory efficiency ratio: {:.2}x",
        full_time.as_nanos() as f64 / streamed_time.as_nanos() as f64
    );

    // Pattern 3: Lazy evaluation with take
    println!("\nPattern 2: Lazy Evaluation");
    let start = Instant::now();
    let lazy_result: Tensor = tensor
        .iter_elements()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect_shape(vec![1000]);
    let lazy_time = start.elapsed();

    println!("  Lazy processing (1000 elements): {:?}", lazy_time);
    println!("  Lazy result size: {}", lazy_result.size());

    // Pattern 4: Memory-efficient filtering
    println!("\nPattern 3: Memory-Efficient Filtering");
    let start = Instant::now();
    let filtered_result: Tensor = tensor
        .iter_elements()
        .filter(|elem| elem.value() > size as f32 / 2.0) // Keep only large values
        .map(|elem| elem.mul_scalar(2.0))
        .collect();
    let filtered_time = start.elapsed();

    println!("  Filtered processing: {:?}", filtered_time);
    println!(
        "  Filtered result size: {} (reduced from {})",
        filtered_result.size(),
        size
    );

    Ok(())
}

pub fn chunks_exact(&self, chunk_size: usize) -> TensorChunksExactIterator<'_>

Standard slice-like exact chunks iterator. Use this instead of iter_chunks_exact.

Yields only the exact chunks of size chunk_size, exposing any remainder via remainder(). See chunks() for a variant that yields the remainder as the last (smaller) chunk.

Examples

use train_station::Tensor;

let t = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5]).unwrap();
let mut it = t.chunks_exact(2);
assert_eq!(it.next().unwrap().data(), &[1.0, 2.0]);
assert_eq!(it.next().unwrap().data(), &[3.0, 4.0]);
assert_eq!(it.remainder().data(), &[5.0]);

pub fn iter_chunks(&self, chunk_size: usize) -> TensorChunksIterator<'_>

👎Deprecated: Use Tensor::chunks(…) instead. This alias will be removed before 1.0.

pub fn iter_chunks_exact( &self, chunk_size: usize, ) -> TensorChunksExactIterator<'_>

👎Deprecated: Use Tensor::chunks_exact(…) instead. This alias will be removed before 1.0.

impl Tensor

pub fn collect_into_shape<I: IntoIterator<Item = Tensor>>( iter: I, dims: Vec<usize>, ) -> Tensor

Collect tensors into a single tensor with target shape, copying data in iterator order. Optimizes copy using SIMD when available; asserts total size matches.

impl Tensor

pub fn collect_values_shape<I: IntoIterator<Item = f32>>( iter: I, dims: Vec<usize>, ) -> Tensor

Inherent helper to collect any iterator of f32 into a shaped tensor. This mirrors the ValuesCollectExt::collect_shape functionality but does not require importing the extension trait.

impl Tensor

pub fn iter_elements(&self) -> TensorElementIterator<'_>

Create an iterator over scalar elements (flattened view)

Each yielded item is a [1]-shaped Tensor view that shares storage with the source. This iterator is GradTrack-aware; element operations propagate gradients to the original tensor when gradients are enabled.

Returns

An iterator producing scalar view tensors in row-major order.

Examples

Collect transformed elements back to the original shape using collect_shape:

use train_station::tensor::TensorCollectExt;
use train_station::Tensor;

let x = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let y = x
    .iter_elements()
    .map(|e| e.mul_scalar(2.0))
    .collect_shape(vec![2, 2]);
assert_eq!(y.data(), &[2.0, 4.0, 6.0, 8.0]);

Examples found in repository

examples/iterators/performance_optimization.rs (line 108)

fn demonstrate_performance_benchmarking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Performance Benchmarking ---");

    // Create test data of different sizes
    let sizes = vec![100, 1000, 10000];

    for size in sizes {
        println!("\nBenchmarking with tensor size: {}", size);

        // Generate test data
        let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
        let tensor = Tensor::from_slice(&data, vec![size])?;

        // Benchmark 1: Direct tensor operations
        let start = Instant::now();
        let direct_result = tensor.mul_scalar(2.0).add_scalar(1.0);
        let direct_time = start.elapsed();

        // Benchmark 2: Iterator-based operations (grad-enabled views, flatten + collect_shape)
        let start = Instant::now();
        let iterator_result: Tensor = tensor
            .iter_elements()
            .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0))
            .collect_shape(vec![size]);
        let iterator_time = start.elapsed();

        // Benchmark 3: Chained iterator operations
        let start = Instant::now();
        let _chained_result: Tensor = tensor
            .iter_elements()
            .map(|elem| elem.mul_scalar(2.0))
            .filter(|elem| elem.value() > size as f32)
            .map(|elem| elem.add_scalar(1.0))
            .collect();
        let chained_time = start.elapsed();

        // Benchmark 4: NoGrad raw data streaming + collect_shape
        let start = Instant::now();
        let _streamed: Tensor = with_no_grad(|| {
            tensor
                .data()
                .iter()
                .copied()
                .map(|x| 2.0 * x + 1.0)
                .collect_shape(vec![size])
        });
        let streaming_time = start.elapsed();

        // Report results
        println!("  Direct operations: {:?}", direct_time);
        println!("  Iterator operations: {:?}", iterator_time);
        println!("  Chained operations: {:?}", chained_time);
        println!("  NoGrad streaming (data.iter): {:?}", streaming_time);

        // Verify correctness
        assert_eq!(direct_result.data(), iterator_result.data());
        println!(
            "  Results match: {}",
            direct_result.data() == iterator_result.data()
        );

        // Performance ratios
        let ratio = iterator_time.as_nanos() as f64 / direct_time.as_nanos() as f64;
        let ratio_stream = streaming_time.as_nanos() as f64 / direct_time.as_nanos() as f64;
        println!("  Iterator/Direct ratio: {:.2}x", ratio);
        println!("  Streaming/Direct ratio: {:.2}x", ratio_stream);
    }

    Ok(())
}

/// Demonstrate memory optimization patterns
///
/// Shows memory-efficient processing patterns and techniques
/// for minimizing memory usage while maintaining performance.
fn demonstrate_memory_optimization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Memory Optimization ---");

    // Create a large tensor for memory testing
    let size = 10000;
    let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
    let tensor = Tensor::from_slice(&data, vec![size])?;

    println!("Processing tensor of size: {}", size);

    // Pattern 1: Streaming processing with iterator chunks (process in blocks, collect with shape)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();
    let flattened = tensor.view(vec![size as i32]);
    let _streamed_result: Tensor = flattened
        .chunks(chunk_size)
        .map(|c| c.pow_scalar(2.0).sqrt())
        .collect_shape(vec![size]);
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter_elements()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect_shape(vec![size]);
    let full_time = start.elapsed();

    println!("  Streaming time: {:?}", streamed_time);
    println!("  Full processing time: {:?}", full_time);
    println!(
        "  Memory efficiency ratio: {:.2}x",
        full_time.as_nanos() as f64 / streamed_time.as_nanos() as f64
    );

    // Pattern 3: Lazy evaluation with take
    println!("\nPattern 2: Lazy Evaluation");
    let start = Instant::now();
    let lazy_result: Tensor = tensor
        .iter_elements()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect_shape(vec![1000]);
    let lazy_time = start.elapsed();

    println!("  Lazy processing (1000 elements): {:?}", lazy_time);
    println!("  Lazy result size: {}", lazy_result.size());

    // Pattern 4: Memory-efficient filtering
    println!("\nPattern 3: Memory-Efficient Filtering");
    let start = Instant::now();
    let filtered_result: Tensor = tensor
        .iter_elements()
        .filter(|elem| elem.value() > size as f32 / 2.0) // Keep only large values
        .map(|elem| elem.mul_scalar(2.0))
        .collect();
    let filtered_time = start.elapsed();

    println!("  Filtered processing: {:?}", filtered_time);
    println!(
        "  Filtered result size: {} (reduced from {})",
        filtered_result.size(),
        size
    );

    Ok(())
}

pub fn iter_range(&self, start: usize, end: usize) -> TensorElementIterator<'_>

Create an iterator over a clamped range of elements

Produces scalar view tensors from start..end (clamped to [0, size]).

Arguments

• start - Start index (inclusive)
• end - End index (exclusive)

Examples

use train_station::Tensor;

let x = Tensor::from_slice(&(0..6).map(|i| i as f32).collect::<Vec<_>>(), vec![6]).unwrap();
let vals: Vec<f32> = x.iter_range(2, 5).map(|e| e.value()).collect();
assert_eq!(vals, vec![2.0, 3.0, 4.0]);

Examples found in repository

examples/iterators/performance_optimization.rs (line 262)

fn demonstrate_large_scale_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Large-Scale Processing ---");

    // Simulate large dataset processing
    let sizes = vec![10000, 50000, 100000];

    for size in sizes {
        println!("\nProcessing dataset of size: {}", size);

        // Generate large dataset
        let data: Vec<f32> = (0..size)
            .map(|i| {
                let x = i as f32 / size as f32;
                x * x + 0.1 * (i % 10) as f32 // Quadratic with noise
            })
            .collect();

        let tensor = Tensor::from_slice(&data, vec![size])?;

        // Technique 1: Batch processing
        let batch_size = 1000;
        let start = Instant::now();

        let mut batch_results = Vec::new();
        for batch_start in (0..size).step_by(batch_size) {
            let batch_end = (batch_start + batch_size).min(size);
            let batch: Tensor = tensor
                .iter_range(batch_start, batch_end)
                .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
                .collect();
            batch_results.push(batch);
        }
        let batch_time = start.elapsed();

        // Technique 2: Parallel-like processing with stride
        let start = Instant::now();
        let stride = 4;
        let strided_result: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % stride == 0)
            .map(|(_, elem)| elem.pow_scalar(2.0).add_scalar(1.0))
            .collect();
        let strided_time = start.elapsed();

        // Technique 3: Hierarchical processing
        let start = Instant::now();
        let coarse: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % 10 == 0) // Every 10th element
            .map(|(_, elem)| elem.pow_scalar(2.0).add_scalar(1.0))
            .collect();
        let fine: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % 10 != 0) // Rest of elements
            .map(|(_, elem)| elem.pow_scalar(1.5).add_scalar(0.5))
            .collect();
        let hierarchical_time = start.elapsed();

        // Report performance
        println!("  Batch processing: {:?}", batch_time);
        println!("  Strided processing: {:?}", strided_time);
        println!("  Hierarchical processing: {:?}", hierarchical_time);

        // Memory usage analysis
        let total_batches = size.div_ceil(batch_size);
        println!("  Batch count: {}", total_batches);
        println!("  Strided result size: {}", strided_result.size());
        println!(
            "  Hierarchical: coarse={}, fine={}",
            coarse.size(),
            fine.size()
        );
    }

    Ok(())
}

examples/iterators/advanced_patterns.rs (line 363)

fn demonstrate_real_world_scenarios() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Real-world Scenarios ---");

    // Scenario 1: Time series analysis
    println!("\nScenario 1: Time Series Analysis");
    let time_series: Vec<f32> = (0..24)
        .map(|hour| {
            let base = 20.0 + 10.0 * (hour as f32 * std::f32::consts::PI / 12.0).sin();
            base + (hour % 3) as f32 * 2.0 // Add some noise
        })
        .collect();

    let series = Tensor::from_slice(&time_series, vec![24])?;
    println!("  Time series (24 hours): {:?}", series.data());

    // Calculate moving average with view-based iteration
    let window_size = 3;
    let moving_avg: Tensor = series
        .iter()
        .enumerate()
        .map(|(i, _)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(series.size());
            let window = series.iter_range(start, end);
            window.fold(0.0, |acc, elem| acc + elem.value()) / (end - start) as f32
        })
        .map(|val| Tensor::from_slice(&[val], vec![1]).unwrap())
        .collect();
    println!(
        "  Moving average (window={}): {:?}",
        window_size,
        moving_avg.data()
    );

    // Inference pipeline with NoGrad + streaming
    println!("\nInference pipeline (NoGrad + streaming)");
    let features = Tensor::from_slice(
        &(0..48).map(|i| i as f32 * 0.125).collect::<Vec<_>>(),
        vec![6, 8],
    )?;
    let fast = with_no_grad(|| {
        // Stream values directly, apply light affine, and collect back to same shape
        features
            .data()
            .iter()
            .copied()
            .map(|x| 0.75 * x + 0.1)
            .collect_shape(vec![6, 8])
    });
    println!(
        "  NoGrad streamed transform shape: {:?}",
        fast.shape().dims()
    );

    // Row-wise iteration with shape-preserving collection (GradTrack-friendly)
    let per_row: Tensor = features
        .iter()
        .map(|row| row.mul_scalar(0.5).add_scalar(2.0))
        .collect_shape(vec![6, 8]);
    println!("  Row-wise mapped shape: {:?}", per_row.shape().dims());

    // Scenario 2: Feature engineering
    println!("\nScenario 2: Feature Engineering");
    let features = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;
    println!("  Original features: {:?}", features.data());

    // Create polynomial features
    let poly_features: Tensor = features
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // x^1
                elem.pow_scalar(2.0), // x^2
                elem.pow_scalar(3.0), // x^3
            ]
        })
        .collect();
    println!(
        "  Polynomial features (x, x^2, x^3): {:?}",
        poly_features.data()
    );

    // Scenario 3: Data augmentation
    println!("\nScenario 3: Data Augmentation");
    let original = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?;
    println!("  Original data: {:?}", original.data());

    // Augment with noise and scaling
    let augmented: Tensor = original
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // Original
                elem.add_scalar(0.1), // Add noise
                elem.sub_scalar(0.1), // Subtract noise
                elem.mul_scalar(1.1), // Scale up
                elem.mul_scalar(0.9), // Scale down
            ]
        })
        .collect();
    println!("  Augmented data: {:?}", augmented.data());

    // Scenario 4: Statistical analysis
    println!("\nScenario 4: Statistical Analysis");
    let sample_data = Tensor::from_slice(&[1.1, 2.3, 1.8, 2.1, 1.9, 2.0, 1.7, 2.2], vec![8])?;
    println!("  Sample data: {:?}", sample_data.data());

    // Calculate various statistics
    let mean = sample_data.mean().value();
    let std = sample_data.std().value();
    let min = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);

    // Z-score normalization
    let z_scores: Tensor = sample_data
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();

    println!(
        "  Statistics: mean={:.3}, std={:.3}, min={:.3}, max={:.3}",
        mean, std, min, max
    );
    println!("  Z-scores: {:?}", z_scores.data());

    Ok(())
}

impl Tensor

pub fn iter_dim(&self, dim: usize) -> TensorDimIterator<'_>

Iterate over sub-tensors along a specific dimension.

Produces view tensors by slicing along the given dimension; each item has that dimension removed (rank - 1). Views share storage and preserve gradient tracking semantics.

Arguments

• dim - Dimension to iterate over

Examples

use train_station::tensor::TensorCollectExt;
use train_station::Tensor;

let t = Tensor::from_slice(&(1..=6).map(|i| i as f32).collect::<Vec<_>>(), vec![2, 3]).unwrap();
let out = t.iter_dim(0).map(|row| row.add_scalar(1.0)).collect_shape(vec![2, 3]);
assert_eq!(out.data(), &[2.0, 3.0, 4.0, 5.0, 6.0, 7.0]);

pub fn iter(&self) -> TensorDimIterator<'_>

Default iterator over the outermost dimension, yielding sub-tensors (N-D) or scalar views (1-D).

This is equivalent to iter_dim(0) with an important optimization:

• For 1-D tensors, it yields scalar element views of shape [1] (same as iter_flat()) to maximize GradTrack cooperation and collection performance.
• For N-D tensors (rank > 1), it yields sub-tensors with the outermost dimension removed (rank − 1), suitable for row/batch-wise processing.

Views share storage with the source tensor and preserve gradient tracking semantics. Use collect_shape([..]) to reconstruct shape efficiently after per-item transforms.

Returns

An iterator producing view tensors for each slice along the outermost dimension (or scalar views for 1-D).

Examples

1-D: element views (shape [1]) and shape-preserving collection

use train_station::tensor::TensorCollectExt;
use train_station::Tensor;

let v = Tensor::from_slice(&(0..6).map(|i| i as f32).collect::<Vec<_>>(), vec![6]).unwrap();
let out = v.iter().map(|e| e.add_scalar(1.0)).collect_shape(vec![6]);
assert_eq!(out.data(), &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);

2-D: row-wise transforms and shape-preserving collection

use train_station::tensor::TensorCollectExt;
use train_station::Tensor;

let m = Tensor::from_slice(&(1..=6).map(|i| i as f32).collect::<Vec<_>>(), vec![2, 3]).unwrap();
let y = m.iter().map(|row| row.mul_scalar(2.0)).collect_shape(vec![2, 3]);
assert_eq!(y.data(), &[2.0, 4.0, 6.0, 8.0, 10.0, 12.0]);

Examples found in repository

examples/iterators/element_iteration.rs (line 102)

fn demonstrate_basic_iteration() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Basic Element Iteration ---");

    // Create a simple tensor for demonstration
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;
    println!("Original tensor: {:?}", tensor.data());

    // Basic iteration with for loop
    println!("\nBasic iteration with for loop:");
    for (i, element) in tensor.iter().enumerate() {
        println!(
            "  Element {}: value = {:.1}, shape = {:?}",
            i,
            element.value(),
            element.shape().dims()
        );
    }

    // Element-wise transformation
    println!("\nElement-wise transformation (2x + 1):");
    let transformed: Tensor = tensor
        .iter()
        .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0))
        .collect();
    println!("  Result: {:?}", transformed.data());

    // Filtering elements
    println!("\nFiltering elements (values > 3.0):");
    let filtered: Tensor = tensor.iter().filter(|elem| elem.value() > 3.0).collect();
    println!("  Filtered: {:?}", filtered.data());

    Ok(())
}

/// Demonstrate standard iterator trait methods
///
/// Shows compatibility with Rust's standard library iterator methods
/// and demonstrates various functional programming patterns.
fn demonstrate_standard_methods() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Standard Iterator Methods ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;

    // Using map for transformations
    println!("\nMap transformation (square each element):");
    let squared: Tensor = tensor.iter().map(|elem| elem.pow_scalar(2.0)).collect();
    println!("  Squared: {:?}", squared.data());

    // Using enumerate for indexed operations
    println!("\nEnumerate with indexed operations:");
    let indexed: Tensor = tensor
        .iter()
        .enumerate()
        .map(|(i, elem)| elem.add_scalar(i as f32))
        .collect();
    println!("  Indexed: {:?}", indexed.data());

    // Using fold for reduction
    println!("\nFold for sum calculation:");
    let sum: f32 = tensor.iter().fold(0.0, |acc, elem| acc + elem.value());
    println!("  Sum: {:.1}", sum);

    // Using find for element search
    println!("\nFind specific element:");
    if let Some(found) = tensor.iter().find(|elem| elem.value() == 3.0) {
        println!("  Found element with value 3.0: {:.1}", found.value());
    }

    // Using any/all for condition checking
    println!("\nCondition checking:");
    let all_positive = tensor.iter().all(|elem| elem.value() > 0.0);
    let any_large = tensor.iter().any(|elem| elem.value() > 4.0);
    println!("  All positive: {}", all_positive);
    println!("  Any > 4.0: {}", any_large);

    Ok(())
}

/// Demonstrate gradient tracking through element operations
///
/// Shows how gradient tracking works seamlessly through iterator
/// operations, maintaining the computational graph for backpropagation.
fn demonstrate_gradient_tracking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Gradient Tracking ---");

    // Create a tensor with gradient tracking enabled
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?.with_requires_grad();
    println!("Input tensor (requires_grad): {:?}", tensor.data());

    // Perform element-wise operations through iteration
    let result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Apply a complex transformation: (x^2 + 1) * 2
            elem.pow_scalar(2.0).add_scalar(1.0).mul_scalar(2.0)
        })
        .collect();

    println!("Result tensor: {:?}", result.data());
    println!("Result requires_grad: {}", result.requires_grad());

    // Compute gradients
    let mut loss = result.sum();
    loss.backward(None);

    println!("Loss: {:.6}", loss.value());
    println!("Input gradients: {:?}", tensor.grad().map(|g| g.data()));

    Ok(())
}

/// Demonstrate advanced iterator patterns
///
/// Shows complex iterator chains and advanced functional programming
/// patterns for sophisticated data processing workflows.
fn demonstrate_advanced_patterns() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Iterator Patterns ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6])?;
    println!("Input tensor: {:?}", tensor.data());

    // Complex chain: enumerate -> filter -> map -> collect
    println!("\nComplex chain (even indices only, add index to value):");
    let result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 2 == 0) // Take even indices
        .map(|(i, elem)| elem.add_scalar(i as f32)) // Add index to value
        .collect();
    println!("  Result: {:?}", result.data());

    // Using take and skip for windowing
    println!("\nWindowing with take and skip:");
    let window1: Tensor = tensor.iter().take(3).collect();
    let window2: Tensor = tensor.iter().skip(2).take(3).collect();
    println!("  Window 1 (first 3): {:?}", window1.data());
    println!("  Window 2 (middle 3): {:?}", window2.data());

    // Using rev() for reverse iteration
    println!("\nReverse iteration:");
    let reversed: Tensor = tensor.iter().rev().collect();
    println!("  Reversed: {:?}", reversed.data());

    // Chaining with mathematical operations
    println!("\nMathematical operation chain:");
    let math_result: Tensor = tensor
        .iter()
        .map(|elem| elem.exp()) // e^x
        .filter(|elem| elem.value() < 50.0) // Filter large values
        .map(|elem| elem.log()) // ln(x)
        .collect();
    println!("  Math chain result: {:?}", math_result.data());

    // Using zip for element-wise combinations
    println!("\nElement-wise combination with zip:");
    let tensor2 = Tensor::from_slice(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0], vec![6])?;
    let combined: Tensor = tensor
        .iter()
        .zip(tensor2.iter())
        .map(|(a, b)| a.mul_tensor(&b)) // Element-wise multiplication
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

/// Demonstrate per-row transforms with shape-preserving collection
///
/// Shows how to use `iter()` over the outer dimension on a 2D tensor and
/// `collect_shape([..])` to maintain the original shape after mapping.
fn demonstrate_row_wise_collect_shape() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Row-wise iteration with collect_shape ---");
    let mat = Tensor::from_slice(&(1..=12).map(|x| x as f32).collect::<Vec<_>>(), vec![3, 4])?;
    println!("Input shape: {:?}", mat.shape().dims());

    // Map each row: 1.1*x + 0.5, then collect back to [3,4]
    let out: Tensor = mat
        .iter()
        .map(|row| row.mul_scalar(1.1).add_scalar(0.5))
        .collect_shape(vec![3, 4]);
    println!("  Output shape: {:?}", out.shape().dims());

    Ok(())
}

/// Demonstrate NoGrad fast paths and raw data streaming
///
/// Highlights how to get maximum performance in inference by:
/// - Disabling gradient tracking with `with_no_grad`
/// - Iterating raw values via `tensor.data().iter().copied()`
/// - Using `collect_shape` to stream directly into destination tensors
fn demonstrate_nograd_and_streaming() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- NoGrad & Streaming (Inference Fast Paths) ---");

    let input = Tensor::from_slice(
        &(0..24).map(|i| i as f32 * 0.25).collect::<Vec<_>>(),
        vec![4, 6],
    )?;
    println!("Input shape: {:?}", input.shape().dims());

    // NoGrad: stream values directly and reshape
    let out = with_no_grad(|| {
        input
            .data()
            .iter()
            .copied()
            .map(|x| 1.2 * x - 0.3)
            .collect_shape(vec![4, 6])
    });
    println!(
        "  NoGrad streamed map (1.2x-0.3) -> shape {:?}",
        out.shape().dims()
    );

    // Compare to view-based element iteration in NoGrad
    let out_view: Tensor = with_no_grad(|| {
        input
            .iter()
            .map(|e| e.mul_scalar(1.2).add_scalar(-0.3))
            .collect_shape(vec![4, 6])
    });
    println!(
        "  NoGrad view-based map shape {:?}",
        out_view.shape().dims()
    );

    // Quick parity check
    assert_eq!(out.data(), out_view.data());
    println!("  Parity check passed.");

    // Show simple flatten + collect back to a different shape
    let reshaped = with_no_grad(|| input.data().iter().copied().collect_shape(vec![6, 4]));
    println!(
        "  Reshaped via streaming collect_shape: {:?}",
        reshaped.shape().dims()
    );

    Ok(())
}

examples/iterators/advanced_patterns.rs (line 113)

fn demonstrate_data_pipeline() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Data Processing Pipeline ---");

    // Simulate raw sensor data with noise
    let raw_data: Vec<f32> = (0..20)
        .map(|i| {
            let base = i as f32 * 0.5;
            let noise = (i % 3) as f32 * 0.1;
            base + noise
        })
        .collect();

    let tensor = Tensor::from_slice(&raw_data, vec![20])?;
    println!("Raw sensor data: {:?}", tensor.data());

    // Multi-stage processing pipeline
    println!("\nProcessing pipeline:");
    println!("1. Normalize data (z-score)");
    println!("2. Apply smoothing filter");
    println!("3. Detect outliers");
    println!("4. Apply feature scaling");

    // Stage 1: Normalization
    let mean = tensor.mean().value();
    let std = tensor.std().value();
    let normalized: Tensor = tensor
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();
    println!(
        "  Normalized (mean={:.3}, std={:.3}): {:?}",
        mean,
        std,
        normalized.data()
    );

    // Stage 2: Smoothing (simple moving average)
    let smoothed: Tensor = normalized
        .iter()
        .enumerate()
        .map(|(i, elem)| {
            if i == 0 || i == normalized.size() - 1 {
                elem.clone()
            } else {
                // Simple 3-point average
                let prev = normalized.element_view(i - 1);
                let next = normalized.element_view(i + 1);
                elem.add_tensor(&prev).add_tensor(&next).div_scalar(3.0)
            }
        })
        .collect();
    println!("  Smoothed: {:?}", smoothed.data());

    // Stage 3: Outlier detection and removal
    let outlier_threshold = 2.0;
    let cleaned: Tensor = smoothed
        .iter()
        .filter(|elem| elem.value().abs() < outlier_threshold)
        .collect();
    println!(
        "  Outliers removed (threshold={}): {:?}",
        outlier_threshold,
        cleaned.data()
    );

    // Stage 4: Feature scaling to [0, 1] range
    let min_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);
    let scaled: Tensor = cleaned
        .iter()
        .map(|elem| elem.sub_scalar(min_val).div_scalar(max_val - min_val))
        .collect();
    println!("  Scaled to [0,1]: {:?}", scaled.data());

    Ok(())
}

/// Demonstrate conditional processing patterns
///
/// Shows how to implement dynamic filtering and transformation
/// based on data characteristics and conditions.
fn demonstrate_conditional_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Conditional Processing ---");

    // Create data with mixed characteristics
    let data = vec![1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0, 9.0, -10.0];
    let tensor = Tensor::from_slice(&data, vec![10])?;
    println!("Input data: {:?}", tensor.data());

    // Conditional transformation based on sign
    println!("\nConditional transformation (positive/negative handling):");
    let processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val > 0.0 {
                elem.pow_scalar(2.0) // Square positive values
            } else {
                elem.mul_scalar(-1.0).sqrt() // Square root of absolute negative values
            }
        })
        .collect();
    println!("  Processed: {:?}", processed.data());

    // Adaptive filtering based on local statistics
    println!("\nAdaptive filtering (remove values > 2 std from local mean):");
    let window_size = 3;
    let adaptive_filtered: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, elem)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(tensor.size());

            // Calculate local mean and std
            let local_values: Vec<f32> = (start..end)
                .map(|j| tensor.element_view(j).value())
                .collect();

            let local_mean = local_values.iter().sum::<f32>() / local_values.len() as f32;
            let local_variance = local_values
                .iter()
                .map(|v| (v - local_mean).powi(2))
                .sum::<f32>()
                / local_values.len() as f32;
            let local_std = local_variance.sqrt();

            let threshold = local_mean + 2.0 * local_std;
            elem.value() <= threshold
        })
        .map(|(_, elem)| elem)
        .collect();
    println!("  Adaptive filtered: {:?}", adaptive_filtered.data());

    // Multi-condition processing
    println!("\nMulti-condition processing:");
    let multi_processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            match () {
                _ if val > 5.0 => elem.mul_scalar(2.0), // Double large values
                _ if val < -5.0 => elem.div_scalar(2.0), // Halve small values
                _ if val.abs() < 2.0 => elem.add_scalar(1.0), // Add 1 to small values
                _ => elem.clone(),                      // Keep others unchanged
            }
        })
        .collect();
    println!("  Multi-condition: {:?}", multi_processed.data());

    Ok(())
}

/// Demonstrate batch processing operations
///
/// Shows efficient processing of large datasets using iterator
/// patterns and batch operations for performance optimization.
fn demonstrate_batch_operations() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Batch Operations ---");

    // Create a larger dataset for batch processing
    let size = 100;
    let data: Vec<f32> = (0..size)
        .map(|i| {
            let x = i as f32 / size as f32;
            x * x + 0.1 * (i % 7) as f32 // Quadratic with some noise
        })
        .collect();

    let tensor = Tensor::from_slice(&data, vec![size])?;
    println!("Dataset size: {}", tensor.size());

    // Batch processing with windowing (iterator views)
    println!("\nBatch processing with sliding windows:");
    let batch_size = 10;
    let batches: Vec<Tensor> = tensor
        .iter()
        .collect::<Vec<_>>()
        .chunks(batch_size)
        .map(|chunk| {
            // Process each batch independently
            chunk
                .iter()
                .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
                .collect()
        })
        .collect();

    println!(
        "  Processed {} batches of size {}",
        batches.len(),
        batch_size
    );
    for (i, batch) in batches.iter().enumerate() {
        println!(
            "    Batch {}: mean={:.3}, std={:.3}",
            i,
            batch.mean().value(),
            batch.std().value()
        );
    }

    // Parallel-like processing with stride
    println!("\nStrided processing (every nth element):");
    let stride = 5;
    let strided: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % stride == 0)
        .map(|(_, elem)| elem)
        .collect();
    println!("  Strided (every {}th): {:?}", stride, strided.data());

    // Hierarchical processing
    println!("\nHierarchical processing (coarse to fine):");
    let coarse: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 == 0) // Take every 4th element
        .map(|(_, elem)| elem)
        .collect();

    let fine: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 != 0) // Take the rest
        .map(|(_, elem)| elem)
        .collect();

    println!("  Coarse (every 4th): {:?}", coarse.data());
    println!("  Fine (rest): {:?}", fine.data());

    // Combine coarse and fine with different processing
    let combined: Tensor = coarse
        .iter()
        .map(|elem| elem.mul_scalar(2.0)) // Scale coarse
        .chain(fine.iter().map(|elem| elem.div_scalar(2.0))) // Scale fine
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

/// Demonstrate real-world processing scenarios
///
/// Shows practical applications of iterator patterns for
/// common data processing tasks in machine learning and analytics.
fn demonstrate_real_world_scenarios() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Real-world Scenarios ---");

    // Scenario 1: Time series analysis
    println!("\nScenario 1: Time Series Analysis");
    let time_series: Vec<f32> = (0..24)
        .map(|hour| {
            let base = 20.0 + 10.0 * (hour as f32 * std::f32::consts::PI / 12.0).sin();
            base + (hour % 3) as f32 * 2.0 // Add some noise
        })
        .collect();

    let series = Tensor::from_slice(&time_series, vec![24])?;
    println!("  Time series (24 hours): {:?}", series.data());

    // Calculate moving average with view-based iteration
    let window_size = 3;
    let moving_avg: Tensor = series
        .iter()
        .enumerate()
        .map(|(i, _)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(series.size());
            let window = series.iter_range(start, end);
            window.fold(0.0, |acc, elem| acc + elem.value()) / (end - start) as f32
        })
        .map(|val| Tensor::from_slice(&[val], vec![1]).unwrap())
        .collect();
    println!(
        "  Moving average (window={}): {:?}",
        window_size,
        moving_avg.data()
    );

    // Inference pipeline with NoGrad + streaming
    println!("\nInference pipeline (NoGrad + streaming)");
    let features = Tensor::from_slice(
        &(0..48).map(|i| i as f32 * 0.125).collect::<Vec<_>>(),
        vec![6, 8],
    )?;
    let fast = with_no_grad(|| {
        // Stream values directly, apply light affine, and collect back to same shape
        features
            .data()
            .iter()
            .copied()
            .map(|x| 0.75 * x + 0.1)
            .collect_shape(vec![6, 8])
    });
    println!(
        "  NoGrad streamed transform shape: {:?}",
        fast.shape().dims()
    );

    // Row-wise iteration with shape-preserving collection (GradTrack-friendly)
    let per_row: Tensor = features
        .iter()
        .map(|row| row.mul_scalar(0.5).add_scalar(2.0))
        .collect_shape(vec![6, 8]);
    println!("  Row-wise mapped shape: {:?}", per_row.shape().dims());

    // Scenario 2: Feature engineering
    println!("\nScenario 2: Feature Engineering");
    let features = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;
    println!("  Original features: {:?}", features.data());

    // Create polynomial features
    let poly_features: Tensor = features
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // x^1
                elem.pow_scalar(2.0), // x^2
                elem.pow_scalar(3.0), // x^3
            ]
        })
        .collect();
    println!(
        "  Polynomial features (x, x^2, x^3): {:?}",
        poly_features.data()
    );

    // Scenario 3: Data augmentation
    println!("\nScenario 3: Data Augmentation");
    let original = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?;
    println!("  Original data: {:?}", original.data());

    // Augment with noise and scaling
    let augmented: Tensor = original
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // Original
                elem.add_scalar(0.1), // Add noise
                elem.sub_scalar(0.1), // Subtract noise
                elem.mul_scalar(1.1), // Scale up
                elem.mul_scalar(0.9), // Scale down
            ]
        })
        .collect();
    println!("  Augmented data: {:?}", augmented.data());

    // Scenario 4: Statistical analysis
    println!("\nScenario 4: Statistical Analysis");
    let sample_data = Tensor::from_slice(&[1.1, 2.3, 1.8, 2.1, 1.9, 2.0, 1.7, 2.2], vec![8])?;
    println!("  Sample data: {:?}", sample_data.data());

    // Calculate various statistics
    let mean = sample_data.mean().value();
    let std = sample_data.std().value();
    let min = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);

    // Z-score normalization
    let z_scores: Tensor = sample_data
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();

    println!(
        "  Statistics: mean={:.3}, std={:.3}, min={:.3}, max={:.3}",
        mean, std, min, max
    );
    println!("  Z-scores: {:?}", z_scores.data());

    Ok(())
}

examples/iterators/performance_optimization.rs (line 273)

fn demonstrate_large_scale_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Large-Scale Processing ---");

    // Simulate large dataset processing
    let sizes = vec![10000, 50000, 100000];

    for size in sizes {
        println!("\nProcessing dataset of size: {}", size);

        // Generate large dataset
        let data: Vec<f32> = (0..size)
            .map(|i| {
                let x = i as f32 / size as f32;
                x * x + 0.1 * (i % 10) as f32 // Quadratic with noise
            })
            .collect();

        let tensor = Tensor::from_slice(&data, vec![size])?;

        // Technique 1: Batch processing
        let batch_size = 1000;
        let start = Instant::now();

        let mut batch_results = Vec::new();
        for batch_start in (0..size).step_by(batch_size) {
            let batch_end = (batch_start + batch_size).min(size);
            let batch: Tensor = tensor
                .iter_range(batch_start, batch_end)
                .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
                .collect();
            batch_results.push(batch);
        }
        let batch_time = start.elapsed();

        // Technique 2: Parallel-like processing with stride
        let start = Instant::now();
        let stride = 4;
        let strided_result: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % stride == 0)
            .map(|(_, elem)| elem.pow_scalar(2.0).add_scalar(1.0))
            .collect();
        let strided_time = start.elapsed();

        // Technique 3: Hierarchical processing
        let start = Instant::now();
        let coarse: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % 10 == 0) // Every 10th element
            .map(|(_, elem)| elem.pow_scalar(2.0).add_scalar(1.0))
            .collect();
        let fine: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % 10 != 0) // Rest of elements
            .map(|(_, elem)| elem.pow_scalar(1.5).add_scalar(0.5))
            .collect();
        let hierarchical_time = start.elapsed();

        // Report performance
        println!("  Batch processing: {:?}", batch_time);
        println!("  Strided processing: {:?}", strided_time);
        println!("  Hierarchical processing: {:?}", hierarchical_time);

        // Memory usage analysis
        let total_batches = size.div_ceil(batch_size);
        println!("  Batch count: {}", total_batches);
        println!("  Strided result size: {}", strided_result.size());
        println!(
            "  Hierarchical: coarse={}, fine={}",
            coarse.size(),
            fine.size()
        );
    }

    Ok(())
}

/// Demonstrate advanced optimization techniques
///
/// Shows sophisticated optimization strategies and techniques
/// for maximizing performance in tensor iterator operations.
fn demonstrate_optimization_techniques() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Optimization Techniques ---");

    let size = 50000;
    let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
    let tensor = Tensor::from_slice(&data, vec![size])?;

    println!("Optimizing processing for size: {}", size);

    // Technique 1: Operation fusion
    println!("\nTechnique 1: Operation Fusion");
    let start = Instant::now();
    let fused_result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Fuse multiple operations into single chain
            elem.mul_scalar(2.0).add_scalar(1.0).pow_scalar(2.0).sqrt()
        })
        .collect();
    let fused_time = start.elapsed();

    // Technique 2: Conditional optimization
    println!("\nTechnique 2: Conditional Optimization");
    let start = Instant::now();
    let conditional_result: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val < size as f32 / 2.0 {
                elem.mul_scalar(2.0) // Simple operation for small values
            } else {
                elem.pow_scalar(2.0).sqrt() // Complex operation for large values
            }
        })
        .collect();
    let conditional_time = start.elapsed();

    // Technique 3: Cache-friendly processing
    println!("\nTechnique 3: Cache-Friendly Processing");
    let start = Instant::now();
    let cache_friendly_result: Tensor = tensor
        .iter()
        .take(1000) // Process in cache-friendly chunks
        .map(|elem| elem.mul_scalar(2.0))
        .collect();
    let cache_friendly_time = start.elapsed();

    // Technique 4: Memory pooling simulation
    println!("\nTechnique 4: Memory Pooling Simulation");
    let start = Instant::now();
    let pooled_result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 100 == 0) // Process every 100th element
        .map(|(_, elem)| elem.pow_scalar(2.0))
        .collect();
    let pooled_time = start.elapsed();

    // Report optimization results
    println!("  Fused operations: {:?}", fused_time);
    println!("  Conditional optimization: {:?}", conditional_time);
    println!("  Cache-friendly processing: {:?}", cache_friendly_time);
    println!("  Memory pooling simulation: {:?}", pooled_time);

    // Performance analysis
    let fastest = fused_time
        .min(conditional_time)
        .min(cache_friendly_time)
        .min(pooled_time);
    println!("  Fastest technique: {:?}", fastest);

    // Memory efficiency analysis
    println!("  Fused result size: {}", fused_result.size());
    println!("  Conditional result size: {}", conditional_result.size());
    println!(
        "  Cache-friendly result size: {}",
        cache_friendly_result.size()
    );
    println!("  Pooled result size: {}", pooled_result.size());

    // Technique 5: Gradient optimization
    println!("\nTechnique 5: Gradient Optimization");
    let grad_tensor = tensor.with_requires_grad();
    let start = Instant::now();

    let grad_result: Tensor = grad_tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
        .collect();

    let mut loss = grad_result.sum();
    loss.backward(None);
    let grad_time = start.elapsed();

    println!("  Gradient computation: {:?}", grad_time);
    println!(
        "  Gradient tracking enabled: {}",
        grad_result.requires_grad()
    );

    Ok(())
}

pub fn outer_iter(&self) -> TensorDimIterator<'_>

Explicit alias for outermost-dimension iteration of sub-tensors. Equivalent to iter_dim(0).

impl Tensor

pub fn windows(&self, window_size: usize) -> TensorWindowsIterator<'_>

Overlapping windows iterator with step=1. Use this instead of iter_windows.

Produces overlapping linear windows as view tensors. In no-grad fast mode, a contiguous owner may be materialized once for faster subsequent views.

Arguments

• window_size - Length of each window (> 0)

Examples

use train_station::Tensor;

let t = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
let v: Vec<f32> = t.windows(3).map(|w| w.sum().value()).collect();
assert_eq!(v, vec![6.0, 9.0]);

pub fn windows_step( &self, window_size: usize, step: usize, ) -> TensorWindowsIterator<'_>

Overlapping windows iterator with custom step. Use this instead of iter_windows_step.

Produces windows starting at positions 0, step, 2*step, ... up to the last valid start. Reverse iteration yields the same sequence in reverse.

Arguments

• window_size - Length of each window (> 0)
• step - Step between consecutive window starts (> 0)

Examples

use train_station::Tensor;

let t = Tensor::from_slice(&(1..=8).map(|i| i as f32).collect::<Vec<_>>(), vec![8]).unwrap();
let wins: Vec<Tensor> = t.windows_step(3, 2).collect();
assert_eq!(wins[0].data(), &[1.0, 2.0, 3.0]);
assert_eq!(wins[1].data(), &[3.0, 4.0, 5.0]);
assert_eq!(wins[2].data(), &[5.0, 6.0, 7.0]);

pub fn iter_windows(&self, window_size: usize) -> TensorWindowsIterator<'_>

👎Deprecated: Use Tensor::windows(…) instead. This alias will be removed before 1.0.

pub fn iter_windows_step( &self, window_size: usize, step: usize, ) -> TensorWindowsIterator<'_>

👎Deprecated: Use Tensor::windows_step(…) instead. This alias will be removed before 1.0.

impl Tensor

pub fn add_tensor(&self, other: &Tensor) -> Tensor

Element-wise addition with another tensor with broadcasting support.

Performs element-wise addition with automatic broadcasting: output[i] = self[i] + other[i]

Broadcasting enables addition between tensors of different but compatible shapes. Compatible shapes follow NumPy broadcasting rules:

• Dimensions are aligned from the rightmost dimension
• Dimensions are compatible if they are equal, or one of them is 1
• Missing dimensions are treated as 1

Arguments

• other - Tensor to add. Shapes must be broadcast-compatible.

Returns

A new tensor containing the element-wise sum with broadcast result shape

Examples

Same Shape Addition

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let b = Tensor::from_slice(&[4.0, 5.0, 6.0], vec![3]).unwrap();
let c = a.add_tensor(&b);
assert_eq!(c.shape().dims(), vec![3]);
assert_eq!(c.get(&[0]), 5.0);
assert_eq!(c.get(&[1]), 7.0);
assert_eq!(c.get(&[2]), 9.0);

Broadcasting Addition

use train_station::Tensor;

// Broadcasting: [2, 1] + [1, 3] -> [2, 3]
let a = Tensor::from_slice(&[1.0, 2.0], vec![2, 1]).unwrap();
let b = Tensor::from_slice(&[10.0, 20.0, 30.0], vec![1, 3]).unwrap();
let c = a.add_tensor(&b);
assert_eq!(c.shape().dims(), vec![2, 3]);
assert_eq!(c.get(&[0, 0]), 11.0);
assert_eq!(c.get(&[0, 1]), 21.0);
assert_eq!(c.get(&[1, 0]), 12.0);
assert_eq!(c.get(&[1, 1]), 22.0);

Scalar Broadcasting

use train_station::Tensor;

// Scalar broadcasting: [2, 3] + scalar -> [2, 3]
let a = Tensor::ones(vec![2, 3]);
let b = Tensor::from_slice(&[5.0], vec![1]).unwrap();
let c = a.add_tensor(&b);
assert_eq!(c.shape().dims(), vec![2, 3]);
assert_eq!(c.get(&[0, 0]), 6.0);
assert_eq!(c.get(&[1, 2]), 6.0);

Panics

Panics if tensor shapes are not broadcast-compatible

Examples found in repository

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 75)

    pub fn forward(&self, input: &Tensor) -> Tensor {
        // Matrix multiplication: [batch_size, input_size] @ [input_size, output_size] = [batch_size, output_size]
        let output = input.matmul(&self.weight);
        // Add bias: [batch_size, output_size] + [output_size] = [batch_size, output_size]
        output.add_tensor(&self.bias)
    }

examples/supervised_training/supervised_bce.rs (line 63)

fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

examples/RL_training/ppo_discrete.rs (line 299)

fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/neural_networks/basic_encoder.rs (line 55)

    pub fn forward(&self, input: &Tensor, attn_mask: Option<&Tensor>) -> Tensor {
        let attn = self.mha.forward(input, input, input, attn_mask);
        let res1 = attn.add_tensor(input);

        // Feed-forward network with ReLU and residual
        let (b, t, e) = Self::triple(input);
        let x2d = res1.contiguous().view(vec![(b * t) as i32, e as i32]);
        let hidden = self.ffn_in.forward(&x2d).relu();
        let out2d = self.ffn_out.forward(&hidden);
        let out = out2d.view(vec![b as i32, t as i32, e as i32]);
        out.add_tensor(&res1)
    }

examples/RL_training/ppo_continuous.rs (line 244)

fn gaussian_log_prob(action: &Tensor, mean: &Tensor, log_std: &Tensor) -> Tensor {
    // All tensors shaped [B, A] (log_std is broadcastable)
    let std = log_std.exp();
    let var = std.pow_scalar(2.0);
    let log_scale = log_std;
    let diff = action.sub_tensor(mean);
    let log_prob = diff
        .pow_scalar(2.0)
        .div_tensor(&var)
        .add_scalar(std::f32::consts::LN_2 + std::f32::consts::PI)
        .add_tensor(&log_scale.mul_scalar(2.0))
        .mul_scalar(0.5)
        .mul_scalar(-1.0);
    // Sum across action dim (dim=1) -> [B,1]
    log_prob.sum_dims(&[1], true)
}

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

examples/neural_networks/basic_decoder.rs (line 64)

    pub fn forward(
        &self,
        tgt: &Tensor,
        memory: &Tensor,
        causal_mask: Option<&Tensor>,
        cross_mask: Option<&Tensor>,
    ) -> Tensor {
        let self_attn = self.self_attn.forward(tgt, tgt, tgt, causal_mask);
        let res1 = self_attn.add_tensor(tgt);

        let cross = self.cross_attn.forward(&res1, memory, memory, cross_mask);
        let res2 = cross.add_tensor(&res1);

        let (b, t, e) = Self::triple(tgt);
        let x2d = res2.contiguous().view(vec![(b * t) as i32, e as i32]);
        let hidden = self.ffn_in.forward(&x2d).relu();
        let out2d = self.ffn_out.forward(&hidden);
        let out = out2d.view(vec![b as i32, t as i32, e as i32]);
        out.add_tensor(&res2)
    }

Additional examples can be found in:

• examples/getting_started/tensor_basics.rs
• examples/RL_training/td3.rs
• examples/getting_started/tensor_operators.rs
• examples/neural_networks/multi_head_attention.rs
• examples/iterators/advanced_patterns.rs
• examples/RL_training/dqn.rs

pub fn add_scalar(&self, scalar: f32) -> Tensor

Broadcast addition with a scalar value.

Examples found in repository

examples/RL_training/dqn.rs (line 323)

fn pseudo_huber_mean(diff: &Tensor) -> Tensor {
    diff.pow_scalar(2.0)
        .add_scalar(1.0)
        .sqrt()
        .sub_scalar(1.0)
        .mean()
}

examples/supervised_training/supervised_bce.rs (line 64)

fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

examples/RL_training/ppo_continuous.rs (line 243)

fn gaussian_log_prob(action: &Tensor, mean: &Tensor, log_std: &Tensor) -> Tensor {
    // All tensors shaped [B, A] (log_std is broadcastable)
    let std = log_std.exp();
    let var = std.pow_scalar(2.0);
    let log_scale = log_std;
    let diff = action.sub_tensor(mean);
    let log_prob = diff
        .pow_scalar(2.0)
        .div_tensor(&var)
        .add_scalar(std::f32::consts::LN_2 + std::f32::consts::PI)
        .add_tensor(&log_scale.mul_scalar(2.0))
        .mul_scalar(0.5)
        .mul_scalar(-1.0);
    // Sum across action dim (dim=1) -> [B,1]
    log_prob.sum_dims(&[1], true)
}

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

examples/getting_started/tensor_basics.rs (line 112)

fn demonstrate_basic_operations() {
    println!("\n--- Basic Operations ---");

    let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();

    // Addition
    let sum = a.add_tensor(&b);
    println!("A + B: {:?}", sum.data());

    // Subtraction
    let diff = a.sub_tensor(&b);
    println!("A - B: {:?}", diff.data());

    // Multiplication
    let product = a.mul_tensor(&b);
    println!("A * B: {:?}", product.data());

    // Division
    let quotient = a.div_tensor(&b);
    println!("A / B: {:?}", quotient.data());

    // Scalar operations
    let scalar_add = a.add_scalar(5.0);
    println!("A + 5.0: {:?}", scalar_add.data());

    let scalar_mul = a.mul_scalar(2.0);
    println!("A * 2.0: {:?}", scalar_mul.data());
}

examples/iterators/element_iteration.rs (line 115)

fn demonstrate_basic_iteration() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Basic Element Iteration ---");

    // Create a simple tensor for demonstration
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;
    println!("Original tensor: {:?}", tensor.data());

    // Basic iteration with for loop
    println!("\nBasic iteration with for loop:");
    for (i, element) in tensor.iter().enumerate() {
        println!(
            "  Element {}: value = {:.1}, shape = {:?}",
            i,
            element.value(),
            element.shape().dims()
        );
    }

    // Element-wise transformation
    println!("\nElement-wise transformation (2x + 1):");
    let transformed: Tensor = tensor
        .iter()
        .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0))
        .collect();
    println!("  Result: {:?}", transformed.data());

    // Filtering elements
    println!("\nFiltering elements (values > 3.0):");
    let filtered: Tensor = tensor.iter().filter(|elem| elem.value() > 3.0).collect();
    println!("  Filtered: {:?}", filtered.data());

    Ok(())
}

/// Demonstrate standard iterator trait methods
///
/// Shows compatibility with Rust's standard library iterator methods
/// and demonstrates various functional programming patterns.
fn demonstrate_standard_methods() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Standard Iterator Methods ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;

    // Using map for transformations
    println!("\nMap transformation (square each element):");
    let squared: Tensor = tensor.iter().map(|elem| elem.pow_scalar(2.0)).collect();
    println!("  Squared: {:?}", squared.data());

    // Using enumerate for indexed operations
    println!("\nEnumerate with indexed operations:");
    let indexed: Tensor = tensor
        .iter()
        .enumerate()
        .map(|(i, elem)| elem.add_scalar(i as f32))
        .collect();
    println!("  Indexed: {:?}", indexed.data());

    // Using fold for reduction
    println!("\nFold for sum calculation:");
    let sum: f32 = tensor.iter().fold(0.0, |acc, elem| acc + elem.value());
    println!("  Sum: {:.1}", sum);

    // Using find for element search
    println!("\nFind specific element:");
    if let Some(found) = tensor.iter().find(|elem| elem.value() == 3.0) {
        println!("  Found element with value 3.0: {:.1}", found.value());
    }

    // Using any/all for condition checking
    println!("\nCondition checking:");
    let all_positive = tensor.iter().all(|elem| elem.value() > 0.0);
    let any_large = tensor.iter().any(|elem| elem.value() > 4.0);
    println!("  All positive: {}", all_positive);
    println!("  Any > 4.0: {}", any_large);

    Ok(())
}

/// Demonstrate gradient tracking through element operations
///
/// Shows how gradient tracking works seamlessly through iterator
/// operations, maintaining the computational graph for backpropagation.
fn demonstrate_gradient_tracking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Gradient Tracking ---");

    // Create a tensor with gradient tracking enabled
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?.with_requires_grad();
    println!("Input tensor (requires_grad): {:?}", tensor.data());

    // Perform element-wise operations through iteration
    let result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Apply a complex transformation: (x^2 + 1) * 2
            elem.pow_scalar(2.0).add_scalar(1.0).mul_scalar(2.0)
        })
        .collect();

    println!("Result tensor: {:?}", result.data());
    println!("Result requires_grad: {}", result.requires_grad());

    // Compute gradients
    let mut loss = result.sum();
    loss.backward(None);

    println!("Loss: {:.6}", loss.value());
    println!("Input gradients: {:?}", tensor.grad().map(|g| g.data()));

    Ok(())
}

/// Demonstrate advanced iterator patterns
///
/// Shows complex iterator chains and advanced functional programming
/// patterns for sophisticated data processing workflows.
fn demonstrate_advanced_patterns() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Iterator Patterns ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6])?;
    println!("Input tensor: {:?}", tensor.data());

    // Complex chain: enumerate -> filter -> map -> collect
    println!("\nComplex chain (even indices only, add index to value):");
    let result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 2 == 0) // Take even indices
        .map(|(i, elem)| elem.add_scalar(i as f32)) // Add index to value
        .collect();
    println!("  Result: {:?}", result.data());

    // Using take and skip for windowing
    println!("\nWindowing with take and skip:");
    let window1: Tensor = tensor.iter().take(3).collect();
    let window2: Tensor = tensor.iter().skip(2).take(3).collect();
    println!("  Window 1 (first 3): {:?}", window1.data());
    println!("  Window 2 (middle 3): {:?}", window2.data());

    // Using rev() for reverse iteration
    println!("\nReverse iteration:");
    let reversed: Tensor = tensor.iter().rev().collect();
    println!("  Reversed: {:?}", reversed.data());

    // Chaining with mathematical operations
    println!("\nMathematical operation chain:");
    let math_result: Tensor = tensor
        .iter()
        .map(|elem| elem.exp()) // e^x
        .filter(|elem| elem.value() < 50.0) // Filter large values
        .map(|elem| elem.log()) // ln(x)
        .collect();
    println!("  Math chain result: {:?}", math_result.data());

    // Using zip for element-wise combinations
    println!("\nElement-wise combination with zip:");
    let tensor2 = Tensor::from_slice(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0], vec![6])?;
    let combined: Tensor = tensor
        .iter()
        .zip(tensor2.iter())
        .map(|(a, b)| a.mul_tensor(&b)) // Element-wise multiplication
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

/// Demonstrate per-row transforms with shape-preserving collection
///
/// Shows how to use `iter()` over the outer dimension on a 2D tensor and
/// `collect_shape([..])` to maintain the original shape after mapping.
fn demonstrate_row_wise_collect_shape() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Row-wise iteration with collect_shape ---");
    let mat = Tensor::from_slice(&(1..=12).map(|x| x as f32).collect::<Vec<_>>(), vec![3, 4])?;
    println!("Input shape: {:?}", mat.shape().dims());

    // Map each row: 1.1*x + 0.5, then collect back to [3,4]
    let out: Tensor = mat
        .iter()
        .map(|row| row.mul_scalar(1.1).add_scalar(0.5))
        .collect_shape(vec![3, 4]);
    println!("  Output shape: {:?}", out.shape().dims());

    Ok(())
}

/// Demonstrate NoGrad fast paths and raw data streaming
///
/// Highlights how to get maximum performance in inference by:
/// - Disabling gradient tracking with `with_no_grad`
/// - Iterating raw values via `tensor.data().iter().copied()`
/// - Using `collect_shape` to stream directly into destination tensors
fn demonstrate_nograd_and_streaming() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- NoGrad & Streaming (Inference Fast Paths) ---");

    let input = Tensor::from_slice(
        &(0..24).map(|i| i as f32 * 0.25).collect::<Vec<_>>(),
        vec![4, 6],
    )?;
    println!("Input shape: {:?}", input.shape().dims());

    // NoGrad: stream values directly and reshape
    let out = with_no_grad(|| {
        input
            .data()
            .iter()
            .copied()
            .map(|x| 1.2 * x - 0.3)
            .collect_shape(vec![4, 6])
    });
    println!(
        "  NoGrad streamed map (1.2x-0.3) -> shape {:?}",
        out.shape().dims()
    );

    // Compare to view-based element iteration in NoGrad
    let out_view: Tensor = with_no_grad(|| {
        input
            .iter()
            .map(|e| e.mul_scalar(1.2).add_scalar(-0.3))
            .collect_shape(vec![4, 6])
    });
    println!(
        "  NoGrad view-based map shape {:?}",
        out_view.shape().dims()
    );

    // Quick parity check
    assert_eq!(out.data(), out_view.data());
    println!("  Parity check passed.");

    // Show simple flatten + collect back to a different shape
    let reshaped = with_no_grad(|| input.data().iter().copied().collect_shape(vec![6, 4]));
    println!(
        "  Reshaped via streaming collect_shape: {:?}",
        reshaped.shape().dims()
    );

    Ok(())
}

examples/getting_started/tensor_operators.rs (line 233)

fn demonstrate_method_equivalence() {
    println!("\n--- Operator vs Method Call Equivalence ---");

    let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();

    // Addition: operator vs method
    let operator_result = &a + &b;
    let method_result = a.add_tensor(&b);

    println!("A + B (operator): {:?}", operator_result.data());
    println!("A.add_tensor(B): {:?}", method_result.data());
    println!(
        "Results are equal: {}",
        operator_result.data() == method_result.data()
    );

    // Multiplication: operator vs method
    let operator_result = &a * &b;
    let method_result = a.mul_tensor(&b);

    println!("A * B (operator): {:?}", operator_result.data());
    println!("A.mul_tensor(B): {:?}", method_result.data());
    println!(
        "Results are equal: {}",
        operator_result.data() == method_result.data()
    );

    // Scalar addition: operator vs method
    let operator_result = &a + 5.0;
    let method_result = a.add_scalar(5.0);

    println!("A + 5.0 (operator): {:?}", operator_result.data());
    println!("A.add_scalar(5.0): {:?}", method_result.data());
    println!(
        "Results are equal: {}",
        operator_result.data() == method_result.data()
    );
}

Additional examples can be found in:

• examples/iterators/performance_optimization.rs
• examples/iterators/advanced_patterns.rs
• examples/RL_training/ppo_discrete.rs

impl Tensor

pub fn div_tensor(&self, other: &Tensor) -> Tensor

Element-wise division with another tensor with broadcasting support.

Performs element-wise division with automatic broadcasting: output[i] = self[i] / other[i]

Broadcasting enables division between tensors of different but compatible shapes. Compatible shapes follow NumPy broadcasting rules:

• Dimensions are aligned from the rightmost dimension
• Dimensions are compatible if they are equal, or one of them is 1
• Missing dimensions are treated as 1

Arguments

• other - Tensor to divide by. Shapes must be broadcast-compatible.

Returns

A new tensor containing the element-wise quotient with broadcast result shape

Examples

Same Shape Division

use train_station::Tensor;

let a = Tensor::from_slice(&[10.0, 20.0, 30.0], vec![3]).unwrap();
let b = Tensor::from_slice(&[2.0, 4.0, 5.0], vec![3]).unwrap();
let c = a.div_tensor(&b);
assert_eq!(c.shape().dims(), vec![3]);
assert_eq!(c.get(&[0]), 5.0);
assert_eq!(c.get(&[1]), 5.0);
assert_eq!(c.get(&[2]), 6.0);

Broadcasting Division

use train_station::Tensor;

// Broadcasting: [2, 1] / [1, 3] -> [2, 3]
let a = Tensor::from_slice(&[10.0, 20.0], vec![2, 1]).unwrap();
let b = Tensor::from_slice(&[1.0, 2.0, 5.0], vec![1, 3]).unwrap();
let c = a.div_tensor(&b);
assert_eq!(c.shape().dims(), vec![2, 3]);
assert_eq!(c.get(&[0, 0]), 10.0);
assert_eq!(c.get(&[0, 1]), 5.0);
assert_eq!(c.get(&[1, 0]), 20.0);
assert_eq!(c.get(&[1, 1]), 10.0);

Scalar Division

use train_station::Tensor;

// Scalar division: [2, 3] / scalar -> [2, 3]
let a = Tensor::ones(vec![2, 3]);
let b = Tensor::from_slice(&[2.0], vec![1]).unwrap();
let c = a.div_tensor(&b);
assert_eq!(c.shape().dims(), vec![2, 3]);
assert_eq!(c.get(&[0, 0]), 0.5);
assert_eq!(c.get(&[1, 2]), 0.5);

Panics

Panics if tensor shapes are not broadcast-compatible or division by zero

Examples found in repository

examples/RL_training/ppo_continuous.rs (line 242)

fn gaussian_log_prob(action: &Tensor, mean: &Tensor, log_std: &Tensor) -> Tensor {
    // All tensors shaped [B, A] (log_std is broadcastable)
    let std = log_std.exp();
    let var = std.pow_scalar(2.0);
    let log_scale = log_std;
    let diff = action.sub_tensor(mean);
    let log_prob = diff
        .pow_scalar(2.0)
        .div_tensor(&var)
        .add_scalar(std::f32::consts::LN_2 + std::f32::consts::PI)
        .add_tensor(&log_scale.mul_scalar(2.0))
        .mul_scalar(0.5)
        .mul_scalar(-1.0);
    // Sum across action dim (dim=1) -> [B,1]
    log_prob.sum_dims(&[1], true)
}

examples/getting_started/tensor_basics.rs (line 108)

fn demonstrate_basic_operations() {
    println!("\n--- Basic Operations ---");

    let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();

    // Addition
    let sum = a.add_tensor(&b);
    println!("A + B: {:?}", sum.data());

    // Subtraction
    let diff = a.sub_tensor(&b);
    println!("A - B: {:?}", diff.data());

    // Multiplication
    let product = a.mul_tensor(&b);
    println!("A * B: {:?}", product.data());

    // Division
    let quotient = a.div_tensor(&b);
    println!("A / B: {:?}", quotient.data());

    // Scalar operations
    let scalar_add = a.add_scalar(5.0);
    println!("A + 5.0: {:?}", scalar_add.data());

    let scalar_mul = a.mul_scalar(2.0);
    println!("A * 2.0: {:?}", scalar_mul.data());
}

pub fn div_scalar(&self, scalar: f32) -> Tensor

Broadcast division with a scalar value.

Divides every element by the scalar: output[i] = self[i] / scalar

Arguments

• scalar - Value to divide each element by (must not be zero)

Returns

A new tensor with each element divided by the scalar

Examples

Basic Scalar Division

use train_station::Tensor;

let a = Tensor::from_slice(&[10.0, 20.0, 30.0], vec![3]).unwrap();
let b = a.div_scalar(10.0);
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 1.0);
assert_eq!(b.get(&[1]), 2.0);
assert_eq!(b.get(&[2]), 3.0);

Multi-dimensional Scalar Division

use train_station::Tensor;

let a = Tensor::ones(vec![2, 3]);
let b = a.div_scalar(2.0);
assert_eq!(b.shape().dims(), vec![2, 3]);
assert_eq!(b.get(&[0, 0]), 0.5);
assert_eq!(b.get(&[1, 2]), 0.5);

Panics

Panics if scalar is zero

Examples found in repository

examples/neural_networks/multi_head_attention.rs (line 92)

    pub fn forward(
        &self,
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        attn_mask: Option<&Tensor>,
    ) -> Tensor {
        let qkv = Self::project_qkv(query, key, value, &self.q_proj, &self.k_proj, &self.v_proj);
        let (q, k, v) = qkv;

        // Split heads: [b, t, e] -> [b, h, t, d]
        let (b, tq, _e) = Self::triple(query);
        let (_b2, tk, _e2) = Self::triple(key);
        let q = Self::split_heads(&q, b, tq, self.num_heads, self.head_dim);
        let k = Self::split_heads(&k, b, tk, self.num_heads, self.head_dim);
        let v = Self::split_heads(&v, b, tk, self.num_heads, self.head_dim);

        // Scaled dot-product attention
        // logits: [b, h, tq, tk]
        let k_t = k.transpose(2, 3);
        let mut logits = q.matmul(&k_t).div_scalar((self.head_dim as f32).sqrt());
        if let Some(mask) = attn_mask {
            let dims = mask.shape().dims().to_vec();
            // If boolean-like mask matching [b,h,tq,tk], apply masked_fill
            if dims.len() == 4 && dims[0] == b && dims[1] == self.num_heads && dims[2] == tq {
                // Interpret mask > 0.5 as keep; we invert to build masked positions
                let cond: Vec<bool> = mask.data().iter().map(|&v| v < 0.5).collect();
                // Apply masked fill on a flattened view, then reshape back
                let flat_logits = logits.view(vec![(b * self.num_heads * tq * tk) as i32]);
                let filled = flat_logits.masked_fill(&cond, f32::NEG_INFINITY);
                logits = filled.view(vec![b as i32, self.num_heads as i32, tq as i32, tk as i32]);
            } else {
                // Fallback: additive mask
                logits = logits.add_tensor(mask);
            }
        }
        let attn = logits.softmax(3);

        // context: [b, h, tq, d]
        let context = attn.matmul(&v);
        let context = context.permute(vec![0, 2, 1, 3]); // [b, tq, h, d]
        let context = context.contiguous().view(vec![
            b as i32,
            tq as i32,
            (self.num_heads * self.head_dim) as i32,
        ]);

        // Output projection (flatten to 2D, project, then restore 3D)
        let flat = context.view(vec![(b * tq) as i32, self.embed_dim as i32]);
        let out2d = self.out_proj.forward(&flat);
        out2d.view(vec![b as i32, tq as i32, self.embed_dim as i32])
    }

examples/iterators/advanced_patterns.rs (line 114)

fn demonstrate_data_pipeline() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Data Processing Pipeline ---");

    // Simulate raw sensor data with noise
    let raw_data: Vec<f32> = (0..20)
        .map(|i| {
            let base = i as f32 * 0.5;
            let noise = (i % 3) as f32 * 0.1;
            base + noise
        })
        .collect();

    let tensor = Tensor::from_slice(&raw_data, vec![20])?;
    println!("Raw sensor data: {:?}", tensor.data());

    // Multi-stage processing pipeline
    println!("\nProcessing pipeline:");
    println!("1. Normalize data (z-score)");
    println!("2. Apply smoothing filter");
    println!("3. Detect outliers");
    println!("4. Apply feature scaling");

    // Stage 1: Normalization
    let mean = tensor.mean().value();
    let std = tensor.std().value();
    let normalized: Tensor = tensor
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();
    println!(
        "  Normalized (mean={:.3}, std={:.3}): {:?}",
        mean,
        std,
        normalized.data()
    );

    // Stage 2: Smoothing (simple moving average)
    let smoothed: Tensor = normalized
        .iter()
        .enumerate()
        .map(|(i, elem)| {
            if i == 0 || i == normalized.size() - 1 {
                elem.clone()
            } else {
                // Simple 3-point average
                let prev = normalized.element_view(i - 1);
                let next = normalized.element_view(i + 1);
                elem.add_tensor(&prev).add_tensor(&next).div_scalar(3.0)
            }
        })
        .collect();
    println!("  Smoothed: {:?}", smoothed.data());

    // Stage 3: Outlier detection and removal
    let outlier_threshold = 2.0;
    let cleaned: Tensor = smoothed
        .iter()
        .filter(|elem| elem.value().abs() < outlier_threshold)
        .collect();
    println!(
        "  Outliers removed (threshold={}): {:?}",
        outlier_threshold,
        cleaned.data()
    );

    // Stage 4: Feature scaling to [0, 1] range
    let min_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);
    let scaled: Tensor = cleaned
        .iter()
        .map(|elem| elem.sub_scalar(min_val).div_scalar(max_val - min_val))
        .collect();
    println!("  Scaled to [0,1]: {:?}", scaled.data());

    Ok(())
}

/// Demonstrate conditional processing patterns
///
/// Shows how to implement dynamic filtering and transformation
/// based on data characteristics and conditions.
fn demonstrate_conditional_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Conditional Processing ---");

    // Create data with mixed characteristics
    let data = vec![1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0, 9.0, -10.0];
    let tensor = Tensor::from_slice(&data, vec![10])?;
    println!("Input data: {:?}", tensor.data());

    // Conditional transformation based on sign
    println!("\nConditional transformation (positive/negative handling):");
    let processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val > 0.0 {
                elem.pow_scalar(2.0) // Square positive values
            } else {
                elem.mul_scalar(-1.0).sqrt() // Square root of absolute negative values
            }
        })
        .collect();
    println!("  Processed: {:?}", processed.data());

    // Adaptive filtering based on local statistics
    println!("\nAdaptive filtering (remove values > 2 std from local mean):");
    let window_size = 3;
    let adaptive_filtered: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, elem)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(tensor.size());

            // Calculate local mean and std
            let local_values: Vec<f32> = (start..end)
                .map(|j| tensor.element_view(j).value())
                .collect();

            let local_mean = local_values.iter().sum::<f32>() / local_values.len() as f32;
            let local_variance = local_values
                .iter()
                .map(|v| (v - local_mean).powi(2))
                .sum::<f32>()
                / local_values.len() as f32;
            let local_std = local_variance.sqrt();

            let threshold = local_mean + 2.0 * local_std;
            elem.value() <= threshold
        })
        .map(|(_, elem)| elem)
        .collect();
    println!("  Adaptive filtered: {:?}", adaptive_filtered.data());

    // Multi-condition processing
    println!("\nMulti-condition processing:");
    let multi_processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            match () {
                _ if val > 5.0 => elem.mul_scalar(2.0), // Double large values
                _ if val < -5.0 => elem.div_scalar(2.0), // Halve small values
                _ if val.abs() < 2.0 => elem.add_scalar(1.0), // Add 1 to small values
                _ => elem.clone(),                      // Keep others unchanged
            }
        })
        .collect();
    println!("  Multi-condition: {:?}", multi_processed.data());

    Ok(())
}

/// Demonstrate batch processing operations
///
/// Shows efficient processing of large datasets using iterator
/// patterns and batch operations for performance optimization.
fn demonstrate_batch_operations() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Batch Operations ---");

    // Create a larger dataset for batch processing
    let size = 100;
    let data: Vec<f32> = (0..size)
        .map(|i| {
            let x = i as f32 / size as f32;
            x * x + 0.1 * (i % 7) as f32 // Quadratic with some noise
        })
        .collect();

    let tensor = Tensor::from_slice(&data, vec![size])?;
    println!("Dataset size: {}", tensor.size());

    // Batch processing with windowing (iterator views)
    println!("\nBatch processing with sliding windows:");
    let batch_size = 10;
    let batches: Vec<Tensor> = tensor
        .iter()
        .collect::<Vec<_>>()
        .chunks(batch_size)
        .map(|chunk| {
            // Process each batch independently
            chunk
                .iter()
                .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
                .collect()
        })
        .collect();

    println!(
        "  Processed {} batches of size {}",
        batches.len(),
        batch_size
    );
    for (i, batch) in batches.iter().enumerate() {
        println!(
            "    Batch {}: mean={:.3}, std={:.3}",
            i,
            batch.mean().value(),
            batch.std().value()
        );
    }

    // Parallel-like processing with stride
    println!("\nStrided processing (every nth element):");
    let stride = 5;
    let strided: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % stride == 0)
        .map(|(_, elem)| elem)
        .collect();
    println!("  Strided (every {}th): {:?}", stride, strided.data());

    // Hierarchical processing
    println!("\nHierarchical processing (coarse to fine):");
    let coarse: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 == 0) // Take every 4th element
        .map(|(_, elem)| elem)
        .collect();

    let fine: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 != 0) // Take the rest
        .map(|(_, elem)| elem)
        .collect();

    println!("  Coarse (every 4th): {:?}", coarse.data());
    println!("  Fine (rest): {:?}", fine.data());

    // Combine coarse and fine with different processing
    let combined: Tensor = coarse
        .iter()
        .map(|elem| elem.mul_scalar(2.0)) // Scale coarse
        .chain(fine.iter().map(|elem| elem.div_scalar(2.0))) // Scale fine
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

/// Demonstrate real-world processing scenarios
///
/// Shows practical applications of iterator patterns for
/// common data processing tasks in machine learning and analytics.
fn demonstrate_real_world_scenarios() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Real-world Scenarios ---");

    // Scenario 1: Time series analysis
    println!("\nScenario 1: Time Series Analysis");
    let time_series: Vec<f32> = (0..24)
        .map(|hour| {
            let base = 20.0 + 10.0 * (hour as f32 * std::f32::consts::PI / 12.0).sin();
            base + (hour % 3) as f32 * 2.0 // Add some noise
        })
        .collect();

    let series = Tensor::from_slice(&time_series, vec![24])?;
    println!("  Time series (24 hours): {:?}", series.data());

    // Calculate moving average with view-based iteration
    let window_size = 3;
    let moving_avg: Tensor = series
        .iter()
        .enumerate()
        .map(|(i, _)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(series.size());
            let window = series.iter_range(start, end);
            window.fold(0.0, |acc, elem| acc + elem.value()) / (end - start) as f32
        })
        .map(|val| Tensor::from_slice(&[val], vec![1]).unwrap())
        .collect();
    println!(
        "  Moving average (window={}): {:?}",
        window_size,
        moving_avg.data()
    );

    // Inference pipeline with NoGrad + streaming
    println!("\nInference pipeline (NoGrad + streaming)");
    let features = Tensor::from_slice(
        &(0..48).map(|i| i as f32 * 0.125).collect::<Vec<_>>(),
        vec![6, 8],
    )?;
    let fast = with_no_grad(|| {
        // Stream values directly, apply light affine, and collect back to same shape
        features
            .data()
            .iter()
            .copied()
            .map(|x| 0.75 * x + 0.1)
            .collect_shape(vec![6, 8])
    });
    println!(
        "  NoGrad streamed transform shape: {:?}",
        fast.shape().dims()
    );

    // Row-wise iteration with shape-preserving collection (GradTrack-friendly)
    let per_row: Tensor = features
        .iter()
        .map(|row| row.mul_scalar(0.5).add_scalar(2.0))
        .collect_shape(vec![6, 8]);
    println!("  Row-wise mapped shape: {:?}", per_row.shape().dims());

    // Scenario 2: Feature engineering
    println!("\nScenario 2: Feature Engineering");
    let features = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;
    println!("  Original features: {:?}", features.data());

    // Create polynomial features
    let poly_features: Tensor = features
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // x^1
                elem.pow_scalar(2.0), // x^2
                elem.pow_scalar(3.0), // x^3
            ]
        })
        .collect();
    println!(
        "  Polynomial features (x, x^2, x^3): {:?}",
        poly_features.data()
    );

    // Scenario 3: Data augmentation
    println!("\nScenario 3: Data Augmentation");
    let original = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?;
    println!("  Original data: {:?}", original.data());

    // Augment with noise and scaling
    let augmented: Tensor = original
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // Original
                elem.add_scalar(0.1), // Add noise
                elem.sub_scalar(0.1), // Subtract noise
                elem.mul_scalar(1.1), // Scale up
                elem.mul_scalar(0.9), // Scale down
            ]
        })
        .collect();
    println!("  Augmented data: {:?}", augmented.data());

    // Scenario 4: Statistical analysis
    println!("\nScenario 4: Statistical Analysis");
    let sample_data = Tensor::from_slice(&[1.1, 2.3, 1.8, 2.1, 1.9, 2.0, 1.7, 2.2], vec![8])?;
    println!("  Sample data: {:?}", sample_data.data());

    // Calculate various statistics
    let mean = sample_data.mean().value();
    let std = sample_data.std().value();
    let min = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);

    // Z-score normalization
    let z_scores: Tensor = sample_data
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();

    println!(
        "  Statistics: mean={:.3}, std={:.3}, min={:.3}, max={:.3}",
        mean, std, min, max
    );
    println!("  Z-scores: {:?}", z_scores.data());

    Ok(())
}

impl Tensor

pub fn exp(&self) -> Tensor

Element-wise exponential function.

Computes e^x for each element: output[i] = e^(self[i])

Returns

A new tensor with the exponential of each element

Examples

Basic Exponential

use train_station::Tensor;

let a = Tensor::from_slice(&[0.0, 1.0, 2.0], vec![3]).unwrap();
let b = a.exp();
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 1.0); // e^0 = 1
assert!((b.get(&[1]) - 2.71828).abs() < 1e-5); // e^1 ≈ 2.71828
assert!((b.get(&[2]) - 7.38906).abs() < 1e-5); // e^2 ≈ 7.38906

Negative Values

use train_station::Tensor;

let a = Tensor::from_slice(&[-1.0, 0.0, 1.0], vec![3]).unwrap();
let b = a.exp();
assert_eq!(b.shape().dims(), vec![3]);
assert!((b.get(&[0]) - 0.36788).abs() < 1e-5); // e^(-1) ≈ 0.36788
assert_eq!(b.get(&[1]), 1.0); // e^0 = 1
assert!((b.get(&[2]) - 2.71828).abs() < 1e-5); // e^1 ≈ 2.71828

Examples found in repository

examples/supervised_training/supervised_bce.rs (line 64)

fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

examples/supervised_training/supervised_classification.rs (line 53)

fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

examples/RL_training/ppo_discrete.rs (line 281)

fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

// probability ratio = exp(new_logp - old_logp)
fn ratio_from_logps(new_logp: &Tensor, old_logp: &Tensor) -> Tensor {
    new_logp.sub_tensor(old_logp).exp()
}

examples/RL_training/ppo_continuous.rs (line 236)

fn gaussian_log_prob(action: &Tensor, mean: &Tensor, log_std: &Tensor) -> Tensor {
    // All tensors shaped [B, A] (log_std is broadcastable)
    let std = log_std.exp();
    let var = std.pow_scalar(2.0);
    let log_scale = log_std;
    let diff = action.sub_tensor(mean);
    let log_prob = diff
        .pow_scalar(2.0)
        .div_tensor(&var)
        .add_scalar(std::f32::consts::LN_2 + std::f32::consts::PI)
        .add_tensor(&log_scale.mul_scalar(2.0))
        .mul_scalar(0.5)
        .mul_scalar(-1.0);
    // Sum across action dim (dim=1) -> [B,1]
    log_prob.sum_dims(&[1], true)
}

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

examples/iterators/element_iteration.rs (line 240)

fn demonstrate_advanced_patterns() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Iterator Patterns ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6])?;
    println!("Input tensor: {:?}", tensor.data());

    // Complex chain: enumerate -> filter -> map -> collect
    println!("\nComplex chain (even indices only, add index to value):");
    let result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 2 == 0) // Take even indices
        .map(|(i, elem)| elem.add_scalar(i as f32)) // Add index to value
        .collect();
    println!("  Result: {:?}", result.data());

    // Using take and skip for windowing
    println!("\nWindowing with take and skip:");
    let window1: Tensor = tensor.iter().take(3).collect();
    let window2: Tensor = tensor.iter().skip(2).take(3).collect();
    println!("  Window 1 (first 3): {:?}", window1.data());
    println!("  Window 2 (middle 3): {:?}", window2.data());

    // Using rev() for reverse iteration
    println!("\nReverse iteration:");
    let reversed: Tensor = tensor.iter().rev().collect();
    println!("  Reversed: {:?}", reversed.data());

    // Chaining with mathematical operations
    println!("\nMathematical operation chain:");
    let math_result: Tensor = tensor
        .iter()
        .map(|elem| elem.exp()) // e^x
        .filter(|elem| elem.value() < 50.0) // Filter large values
        .map(|elem| elem.log()) // ln(x)
        .collect();
    println!("  Math chain result: {:?}", math_result.data());

    // Using zip for element-wise combinations
    println!("\nElement-wise combination with zip:");
    let tensor2 = Tensor::from_slice(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0], vec![6])?;
    let combined: Tensor = tensor
        .iter()
        .zip(tensor2.iter())
        .map(|(a, b)| a.mul_tensor(&b)) // Element-wise multiplication
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

impl Tensor

pub fn leaky_relu(&self, negative_slope: f32) -> Tensor

Element-wise Leaky ReLU activation.

Applies Leaky ReLU to each element: output[i] = max(0, x) + negative_slope * min(0, x)

Unlike standard ReLU, allows a small gradient when the unit is not active.

Arguments

• negative_slope - Slope for negative values (typically small, e.g., 0.01 or 0.1)

Returns

A new tensor with Leaky ReLU applied to each element

Examples

Basic Leaky ReLU

use train_station::Tensor;

let a = Tensor::from_slice(&[-2.0, -1.0, 0.0, 1.0], vec![4]).unwrap();
let b = a.leaky_relu(0.1);
assert_eq!(b.shape().dims(), vec![4]);
assert!((b.get(&[0]) - (-0.2)).abs() < 1e-6); // -2.0 * 0.1 = -0.2
assert!((b.get(&[1]) - (-0.1)).abs() < 1e-6); // -1.0 * 0.1 = -0.1
assert_eq!(b.get(&[2]), 0.0); // max(0, 0) = 0
assert_eq!(b.get(&[3]), 1.0); // max(0, 1) = 1

Different Negative Slopes

use train_station::Tensor;

let a = Tensor::from_slice(&[-1.0, 0.0, 1.0], vec![3]).unwrap();
let b = a.leaky_relu(0.01); // Smaller negative slope
assert_eq!(b.shape().dims(), vec![3]);
assert!((b.get(&[0]) - (-0.01)).abs() < 1e-6); // -1.0 * 0.01 = -0.01
assert_eq!(b.get(&[1]), 0.0); // max(0, 0) = 0
assert_eq!(b.get(&[2]), 1.0); // max(0, 1) = 1

impl Tensor

pub fn log(&self) -> Tensor

Element-wise natural logarithm.

Computes the natural logarithm for each element: output[i] = ln(self[i])

Returns

A new tensor with the natural logarithm of each element

Examples

Basic Natural Logarithm

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 2.71828, 7.38906], vec![3]).unwrap();
let b = a.log();
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 0.0); // ln(1) = 0
assert!((b.get(&[1]) - 1.0).abs() < 1e-5); // ln(e) ≈ 1
assert!((b.get(&[2]) - 2.0).abs() < 1e-5); // ln(e^2) ≈ 2

Mathematical Properties

use train_station::Tensor;

let a = Tensor::from_slice(&[4.0, 8.0, 16.0], vec![3]).unwrap();
let b = a.log();
assert_eq!(b.shape().dims(), vec![3]);
assert!((b.get(&[0]) - 1.38629).abs() < 1e-5); // ln(4) ≈ 1.38629
assert!((b.get(&[1]) - 2.07944).abs() < 1e-5); // ln(8) ≈ 2.07944
assert!((b.get(&[2]) - 2.77259).abs() < 1e-5); // ln(16) ≈ 2.77259

Panics

Panics if any element is non-positive (x <= 0)

Examples found in repository

examples/supervised_training/supervised_bce.rs (line 64)

fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

examples/supervised_training/supervised_classification.rs (line 55)

fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

examples/RL_training/ppo_discrete.rs (line 283)

fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

// probability ratio = exp(new_logp - old_logp)
fn ratio_from_logps(new_logp: &Tensor, old_logp: &Tensor) -> Tensor {
    new_logp.sub_tensor(old_logp).exp()
}

// Clamp ratio to [1-clip, 1+clip] using ReLU-based clamp (no custom ops)
fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/iterators/element_iteration.rs (line 242)

fn demonstrate_advanced_patterns() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Iterator Patterns ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6])?;
    println!("Input tensor: {:?}", tensor.data());

    // Complex chain: enumerate -> filter -> map -> collect
    println!("\nComplex chain (even indices only, add index to value):");
    let result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 2 == 0) // Take even indices
        .map(|(i, elem)| elem.add_scalar(i as f32)) // Add index to value
        .collect();
    println!("  Result: {:?}", result.data());

    // Using take and skip for windowing
    println!("\nWindowing with take and skip:");
    let window1: Tensor = tensor.iter().take(3).collect();
    let window2: Tensor = tensor.iter().skip(2).take(3).collect();
    println!("  Window 1 (first 3): {:?}", window1.data());
    println!("  Window 2 (middle 3): {:?}", window2.data());

    // Using rev() for reverse iteration
    println!("\nReverse iteration:");
    let reversed: Tensor = tensor.iter().rev().collect();
    println!("  Reversed: {:?}", reversed.data());

    // Chaining with mathematical operations
    println!("\nMathematical operation chain:");
    let math_result: Tensor = tensor
        .iter()
        .map(|elem| elem.exp()) // e^x
        .filter(|elem| elem.value() < 50.0) // Filter large values
        .map(|elem| elem.log()) // ln(x)
        .collect();
    println!("  Math chain result: {:?}", math_result.data());

    // Using zip for element-wise combinations
    println!("\nElement-wise combination with zip:");
    let tensor2 = Tensor::from_slice(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0], vec![6])?;
    let combined: Tensor = tensor
        .iter()
        .zip(tensor2.iter())
        .map(|(a, b)| a.mul_tensor(&b)) // Element-wise multiplication
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

impl Tensor

pub fn matmul(&self, other: &Tensor) -> Tensor

Matrix multiplication with intelligent kernel dispatch

Performs matrix multiplication using optimized SIMD kernels selected based on:

• Runtime SIMD capability (AVX512/AVX2/SSE2/Scalar)
• Matrix operation type (1D@1D, 1D@2D, 2D@1D, 2D@2D, ND@ND)
• Matrix size classification (Small/Medium/Large)
• Memory alignment characteristics

Arguments

• other - Right-hand side tensor for multiplication

Returns

Result tensor with appropriate shape based on operation type

Panics

Panics if tensor shapes are incompatible for matrix multiplication

Examples

use train_station::Tensor;

// 2D matrix multiplication
let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();
let result = a.matmul(&b);
assert_eq!(result.shape().dims(), vec![2, 2]);

// 1D dot product
let a = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let b = Tensor::from_slice(&[4.0, 5.0, 6.0], vec![3]).unwrap();
let result = a.matmul(&b);
assert_eq!(result.shape().dims(), vec![]); // Scalar

Examples found in repository

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 73)

    pub fn forward(&self, input: &Tensor) -> Tensor {
        // Matrix multiplication: [batch_size, input_size] @ [input_size, output_size] = [batch_size, output_size]
        let output = input.matmul(&self.weight);
        // Add bias: [batch_size, output_size] + [output_size] = [batch_size, output_size]
        output.add_tensor(&self.bias)
    }

examples/optimizers/adam_configurations.rs (line 111)

fn demonstrate_default_adam() -> Result<(), Box<dyn std::error::Error>> {
    println!("--- Default Adam Configuration ---");

    // Create a simple regression problem: y = 2*x + 1
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(42)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create Adam optimizer with default configuration
    let mut optimizer = Adam::new();
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    println!("Default Adam configuration:");
    println!("  Learning rate: {}", optimizer.learning_rate());
    println!("  Initial weight: {:.6}", weight.value());
    println!("  Initial bias: {:.6}", bias.value());

    // Training loop
    let num_epochs = 50;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        losses.push(loss.value());

        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!("Epoch {:3}: Loss = {:.6}", epoch, loss.value());
        }
    }

    // Evaluate final model
    let _final_predictions = x_data.matmul(&weight) + &bias;
    println!("\nFinal model:");
    println!("  Learned weight: {:.6} (target: 2.0)", weight.value());
    println!("  Learned bias: {:.6} (target: 1.0)", bias.value());
    println!("  Final loss: {:.6}", losses[losses.len() - 1]);

    Ok(())
}

/// Demonstrate learning rate comparison
fn demonstrate_learning_rate_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Learning Rate Comparison ---");

    let learning_rates = [0.001, 0.01, 0.1];
    let mut results = Vec::new();

    for &lr in &learning_rates {
        println!("\nTesting learning rate: {}", lr);

        let stats = train_with_config(TrainingConfig {
            learning_rate: lr,
            ..Default::default()
        })?;

        results.push((lr, stats.clone()));

        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Convergence epoch: {}", stats.convergence_epoch);
    }

    // Compare results
    println!("\nLearning Rate Comparison Summary:");
    for (lr, stats) in &results {
        println!(
            "  LR={:6}: Loss={:.6}, Converged@{}",
            lr, stats.final_loss, stats.convergence_epoch
        );
    }

    Ok(())
}

/// Demonstrate weight decay comparison
fn demonstrate_weight_decay_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Weight Decay Comparison ---");

    let weight_decays = [0.0, 0.001, 0.01];
    let mut results = Vec::new();

    for &wd in &weight_decays {
        println!("\nTesting weight decay: {}", wd);

        let stats = train_with_config(TrainingConfig {
            weight_decay: wd,
            ..Default::default()
        })?;

        results.push((wd, stats.clone()));

        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Final weight norm: {:.6}", stats.weight_norm);
    }

    // Compare results
    println!("\nWeight Decay Comparison Summary:");
    for (wd, stats) in &results {
        println!(
            "  WD={:6}: Loss={:.6}, Weight Norm={:.6}",
            wd, stats.final_loss, stats.weight_norm
        );
    }

    Ok(())
}

/// Demonstrate beta parameter tuning
fn demonstrate_beta_parameter_tuning() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Beta Parameter Tuning ---");

    let beta_configs = [
        (0.9, 0.999),  // Default
        (0.8, 0.999),  // More aggressive momentum
        (0.95, 0.999), // Less aggressive momentum
        (0.9, 0.99),   // Faster second moment decay
    ];

    let mut results = Vec::new();

    for (i, (beta1, beta2)) in beta_configs.iter().enumerate() {
        println!(
            "\nTesting beta configuration {}: beta1={}, beta2={}",
            i + 1,
            beta1,
            beta2
        );

        let config = TrainingConfig {
            beta1: *beta1,
            beta2: *beta2,
            ..Default::default()
        };

        let stats = train_with_config(config)?;
        results.push(((*beta1, *beta2), stats.clone()));

        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Convergence epoch: {}", stats.convergence_epoch);
    }

    // Compare results
    println!("\nBeta Parameter Comparison Summary:");
    for ((beta1, beta2), stats) in &results {
        println!(
            "  B1={:4}, B2={:5}: Loss={:.6}, Converged@{}",
            beta1, beta2, stats.final_loss, stats.convergence_epoch
        );
    }

    Ok(())
}

/// Demonstrate configuration benchmarking
fn demonstrate_configuration_benchmarking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Configuration Benchmarking ---");

    // Define configurations to benchmark
    let configs = vec![
        (
            "Conservative",
            TrainingConfig {
                learning_rate: 0.001,
                weight_decay: 0.001,
                beta1: 0.95,
                ..Default::default()
            },
        ),
        (
            "Balanced",
            TrainingConfig {
                learning_rate: 0.01,
                weight_decay: 0.0,
                beta1: 0.9,
                ..Default::default()
            },
        ),
        (
            "Aggressive",
            TrainingConfig {
                learning_rate: 0.1,
                weight_decay: 0.0,
                beta1: 0.8,
                ..Default::default()
            },
        ),
    ];

    let mut benchmark_results = Vec::new();

    for (name, config) in configs {
        println!("\nBenchmarking {} configuration:", name);

        let start_time = std::time::Instant::now();
        let stats = train_with_config(config.clone())?;
        let elapsed = start_time.elapsed();

        println!("  Training time: {:.2}ms", elapsed.as_millis());
        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Convergence: {} epochs", stats.convergence_epoch);

        benchmark_results.push((name.to_string(), stats, elapsed));
    }

    // Summary
    println!("\nBenchmarking Summary:");
    for (name, stats, elapsed) in &benchmark_results {
        println!(
            "  {:12}: Loss={:.6}, Time={:4}ms, Converged@{}",
            name,
            stats.final_loss,
            elapsed.as_millis(),
            stats.convergence_epoch
        );
    }

    Ok(())
}

/// Helper function to train with specific configuration
fn train_with_config(config: TrainingConfig) -> Result<TrainingStats, Box<dyn std::error::Error>> {
    // Create training data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(123)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer with custom configuration
    let adam_config = AdamConfig {
        learning_rate: config.learning_rate,
        beta1: config.beta1,
        beta2: config.beta2,
        eps: 1e-8,
        weight_decay: config.weight_decay,
        amsgrad: false,
    };

    let mut optimizer = Adam::with_config(adam_config);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Training loop
    let mut losses = Vec::new();
    let mut convergence_epoch = config.epochs;

    for epoch in 0..config.epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        let loss_value = loss.value();
        losses.push(loss_value);

        // Check for convergence (loss < 0.01)
        if loss_value < 0.01 && convergence_epoch == config.epochs {
            convergence_epoch = epoch;
        }
    }

    Ok(TrainingStats {
        config,
        final_loss: losses[losses.len() - 1],
        loss_history: losses,
        convergence_epoch,
        weight_norm: weight.norm().value(),
    })
}

examples/optimizers/learning_rate_scheduling.rs (line 343)

fn train_with_scheduler(
    scheduler: &mut dyn LearningRateScheduler,
    num_epochs: usize,
) -> Result<TrainingStats, Box<dyn std::error::Error>> {
    // Create training data: y = 2*x + 1
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(456)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer with initial learning rate
    let mut optimizer = Adam::with_learning_rate(0.05);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Training loop
    let mut losses = Vec::new();
    let mut lr_history = Vec::new();
    let mut convergence_epoch = num_epochs;

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Update learning rate using scheduler
        let current_lr = optimizer.learning_rate();
        let new_lr = scheduler.step(current_lr, epoch, loss.value());

        if (new_lr - current_lr).abs() > 1e-8 {
            optimizer.set_learning_rate(new_lr);
        }

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        let loss_value = loss.value();
        losses.push(loss_value);
        lr_history.push(new_lr);

        // Check for convergence
        if loss_value < 0.01 && convergence_epoch == num_epochs {
            convergence_epoch = epoch;
        }
    }

    Ok(TrainingStats {
        scheduler_name: scheduler.name().to_string(),
        final_loss: losses[losses.len() - 1],
        lr_history,
        loss_history: losses,
        convergence_epoch,
    })
}

examples/getting_started/optimizer_basics.rs (line 132)

fn demonstrate_linear_regression() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Linear Regression Training ---");

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(43)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Create simple training data: y = 2*x + 1
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    println!("Training data:");
    println!("  X: {:?}", x_data.data());
    println!("  Y: {:?}", y_true.data());
    println!("  Target: y = 2*x + 1");

    // Training loop
    let num_epochs = 100;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass: y_pred = x * weight + bias
        let y_pred = x_data.matmul(&weight) + &bias;

        // Compute loss: MSE
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        losses.push(loss.value());

        // Print progress every 20 epochs
        if epoch % 20 == 0 || epoch == num_epochs - 1 {
            println!("Epoch {:3}: Loss = {:.6}", epoch, loss.value());
        }
    }

    // Evaluate final model
    let final_predictions = x_data.matmul(&weight) + &bias;
    println!("\nFinal model evaluation:");
    println!("  Learned weight: {:.6}", weight.value());
    println!("  Learned bias: {:.6}", bias.value());
    println!("  Predictions vs True:");

    for i in 0..5 {
        let x1 = x_data.data()[i];
        let pred = final_predictions.data()[i];
        let true_val = y_true.data()[i];
        println!(
            "    x={:.1}: pred={:.3}, true={:.1}, error={:.3}",
            x1,
            pred,
            true_val,
            (pred - true_val).abs()
        );
    }

    Ok(())
}

/// Demonstrate advanced training patterns
fn demonstrate_advanced_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Training Patterns ---");

    // Create a more complex model
    let mut weight = Tensor::randn(vec![1, 2], Some(44)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![2]).with_requires_grad();

    // Create optimizer with different learning rate
    let mut optimizer = Adam::with_learning_rate(0.005);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Create training data: y = 2*x + [1, 3]
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(
        &[3.0, 5.0, 7.0, 9.0, 11.0, 6.0, 8.0, 10.0, 12.0, 14.0],
        vec![5, 2],
    )
    .unwrap();

    println!("Advanced training with monitoring:");
    println!("  Initial learning rate: {}", optimizer.learning_rate());

    // Training loop with monitoring
    let num_epochs = 50;
    let mut losses = Vec::new();
    let mut weight_norms = Vec::new();
    let mut gradient_norms = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Compute gradient norm before optimizer step
        let gradient_norm = weight.grad_owned().unwrap().norm();

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Learning rate scheduling: reduce every 10 epochs
        if epoch > 0 && epoch % 10 == 0 {
            let current_lr = optimizer.learning_rate();
            let new_lr = current_lr * 0.5;
            optimizer.set_learning_rate(new_lr);
            println!(
                "Epoch {:2}: Reduced learning rate from {:.3} to {:.3}",
                epoch, current_lr, new_lr
            );
        }

        // Record metrics
        losses.push(loss.value());
        weight_norms.push(weight.norm().value());
        gradient_norms.push(gradient_norm.value());

        // Print detailed progress
        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, Weight Norm = {:.6}, Gradient Norm = {:.6}",
                epoch,
                loss.value(),
                weight.norm().value(),
                gradient_norm.value()
            );
        }
    }

    println!("Final learning rate: {}", optimizer.learning_rate());

    // Analyze training progression
    let initial_loss = losses[0];
    let final_loss = losses[losses.len() - 1];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final weight norm: {:.6}", weight.norm().value());
    println!("  Final bias: {:?}", bias.data());

    Ok(())
}

/// Demonstrate learning rate scheduling
fn demonstrate_learning_rate_scheduling() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Learning Rate Scheduling ---");

    // Create simple model
    let mut weight = Tensor::randn(vec![1, 1], Some(45)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer with high initial learning rate
    let mut optimizer = Adam::with_learning_rate(0.1);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Simple data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3, 1]).unwrap();
    let y_true = Tensor::from_slice(&[2.0, 4.0, 6.0], vec![3, 1]).unwrap();

    println!("Initial learning rate: {}", optimizer.learning_rate());

    // Training loop with learning rate scheduling
    let num_epochs = 50;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Learning rate scheduling: reduce every 10 epochs
        if epoch > 0 && epoch % 10 == 0 {
            let current_lr = optimizer.learning_rate();
            let new_lr = current_lr * 0.5;
            optimizer.set_learning_rate(new_lr);
            println!(
                "Epoch {:2}: Reduced learning rate from {:.3} to {:.3}",
                epoch, current_lr, new_lr
            );
        }

        losses.push(loss.value());

        // Print progress
        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, LR = {:.3}",
                epoch,
                loss.value(),
                optimizer.learning_rate()
            );
        }
    }

    println!("Final learning rate: {}", optimizer.learning_rate());

    Ok(())
}

/// Demonstrate training monitoring and analysis
fn demonstrate_training_monitoring() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Monitoring ---");

    // Create model
    let mut weight = Tensor::randn(vec![1, 1], Some(46)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Training data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0], vec![4, 1]).unwrap();

    // Training loop with comprehensive monitoring
    let num_epochs = 30;
    let mut losses = Vec::new();
    let mut weight_history = Vec::new();
    let mut bias_history = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Record history
        losses.push(loss.value());
        weight_history.push(weight.value());
        bias_history.push(bias.value());

        // Print detailed monitoring
        if epoch % 5 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, Weight = {:.6}, Bias = {:.6}",
                epoch,
                loss.value(),
                weight.value(),
                bias.value()
            );
        }
    }

    // Analyze training progression
    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", losses[0]);
    println!("  Final loss: {:.6}", losses[losses.len() - 1]);
    println!(
        "  Loss reduction: {:.1}%",
        (losses[0] - losses[losses.len() - 1]) / losses[0] * 100.0
    );

    // Compute statistics
    let loss_mean = compute_mean(&losses);
    let loss_std = compute_std(&losses);
    let weight_change = (weight_history[weight_history.len() - 1] - weight_history[0]).abs();
    let bias_change = (bias_history[bias_history.len() - 1] - bias_history[0]).abs();

    println!("  Average loss: {:.6} ± {:.6}", loss_mean, loss_std);
    println!("  Weight change: {:.6}", weight_change);
    println!("  Bias change: {:.6}", bias_change);
    println!("  Final weight norm: {:.6}", weight.norm().value());
    println!("  Final bias: {:.6}", bias.value());

    Ok(())
}

examples/neural_networks/multi_head_attention.rs (line 92)

    pub fn forward(
        &self,
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        attn_mask: Option<&Tensor>,
    ) -> Tensor {
        let qkv = Self::project_qkv(query, key, value, &self.q_proj, &self.k_proj, &self.v_proj);
        let (q, k, v) = qkv;

        // Split heads: [b, t, e] -> [b, h, t, d]
        let (b, tq, _e) = Self::triple(query);
        let (_b2, tk, _e2) = Self::triple(key);
        let q = Self::split_heads(&q, b, tq, self.num_heads, self.head_dim);
        let k = Self::split_heads(&k, b, tk, self.num_heads, self.head_dim);
        let v = Self::split_heads(&v, b, tk, self.num_heads, self.head_dim);

        // Scaled dot-product attention
        // logits: [b, h, tq, tk]
        let k_t = k.transpose(2, 3);
        let mut logits = q.matmul(&k_t).div_scalar((self.head_dim as f32).sqrt());
        if let Some(mask) = attn_mask {
            let dims = mask.shape().dims().to_vec();
            // If boolean-like mask matching [b,h,tq,tk], apply masked_fill
            if dims.len() == 4 && dims[0] == b && dims[1] == self.num_heads && dims[2] == tq {
                // Interpret mask > 0.5 as keep; we invert to build masked positions
                let cond: Vec<bool> = mask.data().iter().map(|&v| v < 0.5).collect();
                // Apply masked fill on a flattened view, then reshape back
                let flat_logits = logits.view(vec![(b * self.num_heads * tq * tk) as i32]);
                let filled = flat_logits.masked_fill(&cond, f32::NEG_INFINITY);
                logits = filled.view(vec![b as i32, self.num_heads as i32, tq as i32, tk as i32]);
            } else {
                // Fallback: additive mask
                logits = logits.add_tensor(mask);
            }
        }
        let attn = logits.softmax(3);

        // context: [b, h, tq, d]
        let context = attn.matmul(&v);
        let context = context.permute(vec![0, 2, 1, 3]); // [b, tq, h, d]
        let context = context.contiguous().view(vec![
            b as i32,
            tq as i32,
            (self.num_heads * self.head_dim) as i32,
        ]);

        // Output projection (flatten to 2D, project, then restore 3D)
        let flat = context.view(vec![(b * tq) as i32, self.embed_dim as i32]);
        let out2d = self.out_proj.forward(&flat);
        out2d.view(vec![b as i32, tq as i32, self.embed_dim as i32])
    }

impl Tensor

pub fn mul_tensor(&self, other: &Tensor) -> Tensor

Element-wise multiplication with another tensor with broadcasting support.

Performs element-wise multiplication with automatic broadcasting: output[i] = self[i] * other[i]

Broadcasting enables multiplication between tensors of different but compatible shapes. Compatible shapes follow NumPy broadcasting rules:

• Dimensions are aligned from the rightmost dimension
• Dimensions are compatible if they are equal, or one of them is 1
• Missing dimensions are treated as 1

Arguments

• other - Tensor to multiply. Shapes must be broadcast-compatible.

Returns

A new tensor containing the element-wise product with broadcast result shape

Examples

Same Shape Multiplication

use train_station::Tensor;

let a = Tensor::from_slice(&[2.0, 3.0, 4.0], vec![3]).unwrap();
let b = Tensor::from_slice(&[5.0, 6.0, 7.0], vec![3]).unwrap();
let c = a.mul_tensor(&b);
assert_eq!(c.shape().dims(), vec![3]);
assert_eq!(c.get(&[0]), 10.0); // 2.0 * 5.0
assert_eq!(c.get(&[1]), 18.0); // 3.0 * 6.0
assert_eq!(c.get(&[2]), 28.0); // 4.0 * 7.0

Broadcasting Multiplication

use train_station::Tensor;

let a = Tensor::from_slice(&[2.0, 3.0], vec![2, 1]).unwrap();
let b = Tensor::from_slice(&[10.0, 20.0, 30.0], vec![1, 3]).unwrap();
let c = a.mul_tensor(&b);
assert_eq!(c.shape().dims(), vec![2, 3]);
// Result: [[20.0, 40.0, 60.0], [30.0, 60.0, 90.0]]
assert_eq!(c.get(&[0, 0]), 20.0); // 2.0 * 10.0
assert_eq!(c.get(&[0, 1]), 40.0); // 2.0 * 20.0
assert_eq!(c.get(&[1, 0]), 30.0); // 3.0 * 10.0

Panics

Panics if tensor shapes are not broadcast-compatible

Examples found in repository

examples/supervised_training/supervised_bce.rs (line 61)

fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

examples/getting_started/tensor_basics.rs (line 104)

fn demonstrate_basic_operations() {
    println!("\n--- Basic Operations ---");

    let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();

    // Addition
    let sum = a.add_tensor(&b);
    println!("A + B: {:?}", sum.data());

    // Subtraction
    let diff = a.sub_tensor(&b);
    println!("A - B: {:?}", diff.data());

    // Multiplication
    let product = a.mul_tensor(&b);
    println!("A * B: {:?}", product.data());

    // Division
    let quotient = a.div_tensor(&b);
    println!("A / B: {:?}", quotient.data());

    // Scalar operations
    let scalar_add = a.add_scalar(5.0);
    println!("A + 5.0: {:?}", scalar_add.data());

    let scalar_mul = a.mul_scalar(2.0);
    println!("A * 2.0: {:?}", scalar_mul.data());
}

examples/getting_started/tensor_operators.rs (line 222)

fn demonstrate_method_equivalence() {
    println!("\n--- Operator vs Method Call Equivalence ---");

    let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();

    // Addition: operator vs method
    let operator_result = &a + &b;
    let method_result = a.add_tensor(&b);

    println!("A + B (operator): {:?}", operator_result.data());
    println!("A.add_tensor(B): {:?}", method_result.data());
    println!(
        "Results are equal: {}",
        operator_result.data() == method_result.data()
    );

    // Multiplication: operator vs method
    let operator_result = &a * &b;
    let method_result = a.mul_tensor(&b);

    println!("A * B (operator): {:?}", operator_result.data());
    println!("A.mul_tensor(B): {:?}", method_result.data());
    println!(
        "Results are equal: {}",
        operator_result.data() == method_result.data()
    );

    // Scalar addition: operator vs method
    let operator_result = &a + 5.0;
    let method_result = a.add_scalar(5.0);

    println!("A + 5.0 (operator): {:?}", operator_result.data());
    println!("A.add_scalar(5.0): {:?}", method_result.data());
    println!(
        "Results are equal: {}",
        operator_result.data() == method_result.data()
    );
}

examples/iterators/element_iteration.rs (line 252)

fn demonstrate_advanced_patterns() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Iterator Patterns ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6])?;
    println!("Input tensor: {:?}", tensor.data());

    // Complex chain: enumerate -> filter -> map -> collect
    println!("\nComplex chain (even indices only, add index to value):");
    let result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 2 == 0) // Take even indices
        .map(|(i, elem)| elem.add_scalar(i as f32)) // Add index to value
        .collect();
    println!("  Result: {:?}", result.data());

    // Using take and skip for windowing
    println!("\nWindowing with take and skip:");
    let window1: Tensor = tensor.iter().take(3).collect();
    let window2: Tensor = tensor.iter().skip(2).take(3).collect();
    println!("  Window 1 (first 3): {:?}", window1.data());
    println!("  Window 2 (middle 3): {:?}", window2.data());

    // Using rev() for reverse iteration
    println!("\nReverse iteration:");
    let reversed: Tensor = tensor.iter().rev().collect();
    println!("  Reversed: {:?}", reversed.data());

    // Chaining with mathematical operations
    println!("\nMathematical operation chain:");
    let math_result: Tensor = tensor
        .iter()
        .map(|elem| elem.exp()) // e^x
        .filter(|elem| elem.value() < 50.0) // Filter large values
        .map(|elem| elem.log()) // ln(x)
        .collect();
    println!("  Math chain result: {:?}", math_result.data());

    // Using zip for element-wise combinations
    println!("\nElement-wise combination with zip:");
    let tensor2 = Tensor::from_slice(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0], vec![6])?;
    let combined: Tensor = tensor
        .iter()
        .zip(tensor2.iter())
        .map(|(a, b)| a.mul_tensor(&b)) // Element-wise multiplication
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

examples/RL_training/dqn.rs (line 472)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== DQN Example (YardEnv discrete) ===");

    // Dims
    let state_dim = 3usize;
    let action_dim = 3usize;

    // Hparams
    let gamma = 0.99f32;
    let batch_size = 64usize;
    let start_steps = 200usize;
    let target_update_interval = 200usize; // hard update cadence
    let max_grad_norm = 1.0f32;
    let mut epsilon = 1.0f32;
    let eps_min = 0.05f32;
    let eps_decay_steps = 2_000usize; // linear decay
    let total_steps = std::env::var("DQN_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3000usize);

    // Models
    let mut q_net = QNet::new(state_dim, action_dim, Some(7));
    let mut q_targ = QNet::new(state_dim, action_dim, Some(8));
    q_targ.net.copy_from(&q_net.net);
    q_targ.set_requires_grad_all(false);

    // Optimizer
    let mut q_opt = Adam::with_learning_rate(3e-4);
    for p in q_net.parameters() {
        q_opt.add_parameter(p);
    }

    // Replay + env
    let mut rb = ReplayBuffer::new(100_000, state_dim);
    let mut env = YardEnv::new(12345);
    let mut rng = SmallRng::new(999_111);

    // Metrics
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    for t in 0..total_steps {
        // Epsilon-greedy action
        let action_index = if t < start_steps || rng.next_f32() < epsilon {
            rng.sample_index(action_dim)
        } else {
            let _ng = NoGradTrack::new();
            let q_vals = q_net.forward(&state);
            let row = q_vals.data();
            let mut best_i = 0usize;
            let mut best_v = row[0];
            for (i, &r) in row.iter().enumerate().take(action_dim).skip(1) {
                if r > best_v {
                    best_v = r;
                    best_i = i;
                }
            }
            best_i
        };

        // Env step
        let (next_state, reward, done) = env.step(action_index);
        episode_return += reward;

        // Store
        let s_slice = state.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            action_index,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        // Reset on done
        state = if done {
            let st = env.reset();
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Epsilon linear decay
        if t < eps_decay_steps {
            epsilon = (1.0 - (t as f32) / (eps_decay_steps as f32)) * (1.0 - eps_min) + eps_min;
        }

        // Train
        if rb.can_sample(batch_size) {
            let (s, a_idx, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Double DQN target: a* = argmax_a Q_online(s2,a); y = r + (1-d)*gamma*Q_target(s2, a*)
            let target_q = {
                let _ng = NoGradTrack::new();
                let q_online_s2 = q_net.forward(&s2);
                // argmax per row (manual on CPU)
                let row_stride = action_dim;
                let qd = q_online_s2.data();
                let mut next_actions: Vec<usize> = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let base = i * row_stride;
                    let mut bi = 0usize;
                    let mut bv = qd[base];
                    for j in 1..action_dim {
                        let v = qd[base + j];
                        if v > bv {
                            bv = v;
                            bi = j;
                        }
                    }
                    next_actions.push(bi);
                }
                let q_targ_s2 = q_targ.forward(&s2);
                let q_targ_g = q_targ_s2.gather(1, &next_actions, &[batch_size, 1]);
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&q_targ_g))
            };

            // Q(s,a) for current actions
            // Zero grads first
            {
                let mut params = q_net.parameters();
                q_opt.zero_grad(&mut params);
            }

            let q_all = q_net.forward(&s);
            let q_sa = q_all.gather(1, &a_idx, &[batch_size, 1]);
            let diff = q_sa.sub_tensor(&target_q);
            let mut loss = pseudo_huber_mean(&diff);
            loss.backward(None);

            // Step (filter only params with grads)
            {
                let params = q_net.parameters();
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    let gn = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    q_opt.step(&mut with_grads);
                    q_opt.zero_grad(&mut with_grads);
                    if t % 100 == 0 {
                        let mut pn = q_net.parameters();
                        let pn_l2 = params_l2_norm(&mut pn);
                        let q_mean = q_all.mean().value();
                        println!(
                            "t={:5} | loss={:.4} | q_mean={:.3} | grad_norm={:.3} | param_norm={:.3} | eps={:.3}",
                            t, loss.value(), q_mean, gn, pn_l2, epsilon
                        );
                    }
                }
            }

            // Target hard update
            if t % target_update_interval == 0 {
                q_targ.net.copy_from(&q_net.net);
            }

            // Clear graphs
            clear_all_graphs_known();
        }
    }

    println!("=== DQN training finished ===");
    Ok(())
}

examples/RL_training/ppo_discrete.rs (line 503)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

Additional examples can be found in:

• examples/RL_training/ppo_continuous.rs
• examples/RL_training/td3.rs

pub fn mul_scalar(&self, scalar: f32) -> Tensor

Broadcast multiplication with a scalar value.

Multiplies every element by the scalar: output[i] = self[i] * scalar

Arguments

• scalar - Value to multiply with each element

Returns

A new tensor with each element multiplied by the scalar

Examples

Basic Scalar Multiplication

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let b = a.mul_scalar(10.0);
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 10.0); // 1.0 * 10.0
assert_eq!(b.get(&[1]), 20.0); // 2.0 * 10.0
assert_eq!(b.get(&[2]), 30.0); // 3.0 * 10.0

Negative Scalar Multiplication

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let b = a.mul_scalar(-2.0);
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), -2.0); // 1.0 * -2.0
assert_eq!(b.get(&[1]), -4.0); // 2.0 * -2.0
assert_eq!(b.get(&[2]), -6.0); // 3.0 * -2.0

Examples found in repository

examples/RL_training/../neural_networks/basic_linear_layer.rs (line 58)

    pub fn new(input_size: usize, output_size: usize, seed: Option<u64>) -> Self {
        // Xavier/Glorot initialization: scale by sqrt(1/input_size)
        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,
        }
    }

examples/RL_training/dqn.rs (line 291)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

fn params_l2_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let _ng = NoGradTrack::new();
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        for &v in p.data() {
            total_sq += v * v;
        }
    }
    total_sq.sqrt()
}

// Pseudo-Huber loss: sqrt(1 + diff^2) - 1 (smooth, robust)
fn pseudo_huber_mean(diff: &Tensor) -> Tensor {
    diff.pow_scalar(2.0)
        .add_scalar(1.0)
        .sqrt()
        .sub_scalar(1.0)
        .mean()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== DQN Example (YardEnv discrete) ===");

    // Dims
    let state_dim = 3usize;
    let action_dim = 3usize;

    // Hparams
    let gamma = 0.99f32;
    let batch_size = 64usize;
    let start_steps = 200usize;
    let target_update_interval = 200usize; // hard update cadence
    let max_grad_norm = 1.0f32;
    let mut epsilon = 1.0f32;
    let eps_min = 0.05f32;
    let eps_decay_steps = 2_000usize; // linear decay
    let total_steps = std::env::var("DQN_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3000usize);

    // Models
    let mut q_net = QNet::new(state_dim, action_dim, Some(7));
    let mut q_targ = QNet::new(state_dim, action_dim, Some(8));
    q_targ.net.copy_from(&q_net.net);
    q_targ.set_requires_grad_all(false);

    // Optimizer
    let mut q_opt = Adam::with_learning_rate(3e-4);
    for p in q_net.parameters() {
        q_opt.add_parameter(p);
    }

    // Replay + env
    let mut rb = ReplayBuffer::new(100_000, state_dim);
    let mut env = YardEnv::new(12345);
    let mut rng = SmallRng::new(999_111);

    // Metrics
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    for t in 0..total_steps {
        // Epsilon-greedy action
        let action_index = if t < start_steps || rng.next_f32() < epsilon {
            rng.sample_index(action_dim)
        } else {
            let _ng = NoGradTrack::new();
            let q_vals = q_net.forward(&state);
            let row = q_vals.data();
            let mut best_i = 0usize;
            let mut best_v = row[0];
            for (i, &r) in row.iter().enumerate().take(action_dim).skip(1) {
                if r > best_v {
                    best_v = r;
                    best_i = i;
                }
            }
            best_i
        };

        // Env step
        let (next_state, reward, done) = env.step(action_index);
        episode_return += reward;

        // Store
        let s_slice = state.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            action_index,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        // Reset on done
        state = if done {
            let st = env.reset();
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Epsilon linear decay
        if t < eps_decay_steps {
            epsilon = (1.0 - (t as f32) / (eps_decay_steps as f32)) * (1.0 - eps_min) + eps_min;
        }

        // Train
        if rb.can_sample(batch_size) {
            let (s, a_idx, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Double DQN target: a* = argmax_a Q_online(s2,a); y = r + (1-d)*gamma*Q_target(s2, a*)
            let target_q = {
                let _ng = NoGradTrack::new();
                let q_online_s2 = q_net.forward(&s2);
                // argmax per row (manual on CPU)
                let row_stride = action_dim;
                let qd = q_online_s2.data();
                let mut next_actions: Vec<usize> = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let base = i * row_stride;
                    let mut bi = 0usize;
                    let mut bv = qd[base];
                    for j in 1..action_dim {
                        let v = qd[base + j];
                        if v > bv {
                            bv = v;
                            bi = j;
                        }
                    }
                    next_actions.push(bi);
                }
                let q_targ_s2 = q_targ.forward(&s2);
                let q_targ_g = q_targ_s2.gather(1, &next_actions, &[batch_size, 1]);
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&q_targ_g))
            };

            // Q(s,a) for current actions
            // Zero grads first
            {
                let mut params = q_net.parameters();
                q_opt.zero_grad(&mut params);
            }

            let q_all = q_net.forward(&s);
            let q_sa = q_all.gather(1, &a_idx, &[batch_size, 1]);
            let diff = q_sa.sub_tensor(&target_q);
            let mut loss = pseudo_huber_mean(&diff);
            loss.backward(None);

            // Step (filter only params with grads)
            {
                let params = q_net.parameters();
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    let gn = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    q_opt.step(&mut with_grads);
                    q_opt.zero_grad(&mut with_grads);
                    if t % 100 == 0 {
                        let mut pn = q_net.parameters();
                        let pn_l2 = params_l2_norm(&mut pn);
                        let q_mean = q_all.mean().value();
                        println!(
                            "t={:5} | loss={:.4} | q_mean={:.3} | grad_norm={:.3} | param_norm={:.3} | eps={:.3}",
                            t, loss.value(), q_mean, gn, pn_l2, epsilon
                        );
                    }
                }
            }

            // Target hard update
            if t % target_update_interval == 0 {
                q_targ.net.copy_from(&q_net.net);
            }

            // Clear graphs
            clear_all_graphs_known();
        }
    }

    println!("=== DQN training finished ===");
    Ok(())
}

examples/RL_training/ppo_discrete.rs (line 266)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

// log-softmax for selected actions: given logits [B,A] and actions Vec<usize> -> log_prob [B,1]
fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

// probability ratio = exp(new_logp - old_logp)
fn ratio_from_logps(new_logp: &Tensor, old_logp: &Tensor) -> Tensor {
    new_logp.sub_tensor(old_logp).exp()
}

// Clamp ratio to [1-clip, 1+clip] using ReLU-based clamp (no custom ops)
fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/supervised_training/supervised_bce.rs (line 37)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn accuracy(pred: &Tensor, targets: &Tensor) -> f32 {
    // pred: [B,1] with sigmoid; threshold at 0.5
    let p = pred.data();
    let t = targets.data();
    let mut correct = 0usize;
    for i in 0..p.len() {
        let yhat = if p[i] >= 0.5 { 1.0 } else { 0.0 };
        if (yhat - t[i]).abs() < 1e-6 {
            correct += 1;
        }
    }
    correct as f32 / (p.len() as f32)
}

// Numerically stable BCE with logits:
// L = mean( relu(z) - z*y + log(1 + exp(-|z|)) )
fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

examples/supervised_training/supervised_classification.rs (line 37)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

// Cross-entropy over logits: CE = -mean(log_softmax(logits)[range, labels])
fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

examples/supervised_training/supervised_regression.rs (line 36)

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

Additional examples can be found in:

• examples/RL_training/ppo_continuous.rs
• examples/getting_started/tensor_basics.rs
• examples/RL_training/td3.rs
• examples/iterators/element_iteration.rs
• examples/iterators/performance_optimization.rs
• examples/iterators/advanced_patterns.rs

impl Tensor

pub fn pow_scalar(&self, exponent: f32) -> Tensor

Raises each element to a scalar power.

Computes element-wise power: output[i] = self[i]^exponent

Arguments

• exponent - The scalar exponent to raise each element to

Returns

A new tensor with each element raised to the given power

Examples

Basic Scalar Power

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let b = a.pow_scalar(2.0);
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 1.0); // 1.0^2 = 1.0
assert_eq!(b.get(&[1]), 4.0); // 2.0^2 = 4.0
assert_eq!(b.get(&[2]), 9.0); // 3.0^2 = 9.0

Square Root (Power 0.5)

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 4.0, 9.0], vec![3]).unwrap();
let b = a.pow_scalar(0.5);
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 1.0); // sqrt(1.0) = 1.0
assert_eq!(b.get(&[1]), 2.0); // sqrt(4.0) = 2.0
assert_eq!(b.get(&[2]), 3.0); // sqrt(9.0) = 3.0

Examples found in repository

examples/supervised_training/supervised_regression.rs (line 43)

fn mse(pred: &Tensor, target: &Tensor) -> Tensor {
    pred.sub_tensor(target).pow_scalar(2.0).mean()
}

fn rmse(pred: &Tensor, target: &Tensor) -> f32 {
    mse(pred, target).sqrt().value()
}

fn r2_score(pred: &Tensor, target: &Tensor) -> f32 {
    // R^2 = 1 - SS_res / SS_tot
    let y = target;
    let y_mean = y.mean();
    let ss_res = pred.sub_tensor(y).pow_scalar(2.0).sum();
    let ss_tot = y.sub_tensor(&y_mean).pow_scalar(2.0).sum();
    let ss_res_v = ss_res.value();
    let ss_tot_v = ss_tot.value().max(1e-12); // avoid divide by zero
    1.0 - (ss_res_v / ss_tot_v)
}

examples/RL_training/dqn.rs (line 322)

fn pseudo_huber_mean(diff: &Tensor) -> Tensor {
    diff.pow_scalar(2.0)
        .add_scalar(1.0)
        .sqrt()
        .sub_scalar(1.0)
        .mean()
}

examples/RL_training/ppo_continuous.rs (line 237)

fn gaussian_log_prob(action: &Tensor, mean: &Tensor, log_std: &Tensor) -> Tensor {
    // All tensors shaped [B, A] (log_std is broadcastable)
    let std = log_std.exp();
    let var = std.pow_scalar(2.0);
    let log_scale = log_std;
    let diff = action.sub_tensor(mean);
    let log_prob = diff
        .pow_scalar(2.0)
        .div_tensor(&var)
        .add_scalar(std::f32::consts::LN_2 + std::f32::consts::PI)
        .add_tensor(&log_scale.mul_scalar(2.0))
        .mul_scalar(0.5)
        .mul_scalar(-1.0);
    // Sum across action dim (dim=1) -> [B,1]
    log_prob.sum_dims(&[1], true)
}

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

examples/neural_networks/basic_transformer.rs (line 167)

    pub fn train_non_autoregressive_steps(
        &mut self,
        src: &Tensor,
        tgt: &Tensor,
        steps: usize,
        lr: f32,
    ) {
        let mut opt = Adam::with_learning_rate(lr);
        {
            let params_once = self.parameters();
            for p in &params_once {
                opt.add_parameter(p);
            }
        }
        for step in 0..steps {
            // forward + backward scope (immutable borrow)
            {
                let pred = self.forward(src, tgt);
                let diff = pred.sub_tensor(tgt);
                let mut loss = diff.pow_scalar(2.0).mean();
                if step == 0 || step + 1 == steps {
                    println!("NAR train step {}: loss={:.6}", step, loss.value());
                }
                loss.backward(None);
            }
            // step + zero_grad scope (mutable borrow)
            let mut params_step = self.parameters();
            opt.step(&mut params_step);
            opt.zero_grad(&mut params_step);
        }
    }

    /// Auto-regressive training (teacher forcing): predict next token with causal mask
    pub fn train_autoregressive_steps(
        &mut self,
        src: &Tensor,
        tgt: &Tensor,
        steps: usize,
        lr: f32,
    ) {
        let mut opt = Adam::with_learning_rate(lr);
        {
            let params_once = self.parameters();
            for p in &params_once {
                opt.add_parameter(p);
            }
        }

        // Build encoder memory once (static dataset demo)
        let mut memory = src.clone();
        for enc in &self.encoders {
            memory = enc.forward(&memory, None);
        }

        let (b, t, _e) = Self::triple(tgt);
        // Predict y[t] from y[:t] using causal mask; here we simply predict full seq with mask
        let causal = Self::build_causal_mask_static(b, self.num_heads, t);
        for step in 0..steps {
            // forward + backward scope
            {
                let mut out = tgt.clone();
                for dec in &self.decoders {
                    out = dec.forward(&out, &memory, Some(&causal), None);
                }
                let diff = out.sub_tensor(tgt);
                let mut loss = diff.pow_scalar(2.0).mean();
                if step == 0 || step + 1 == steps {
                    println!("AR  train step {}: loss={:.6}", step, loss.value());
                }
                loss.backward(None);
            }
            let mut params_step = self.parameters();
            opt.step(&mut params_step);
            opt.zero_grad(&mut params_step);
        }
    }

examples/iterators/element_iteration.rs (line 138)

fn demonstrate_standard_methods() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Standard Iterator Methods ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;

    // Using map for transformations
    println!("\nMap transformation (square each element):");
    let squared: Tensor = tensor.iter().map(|elem| elem.pow_scalar(2.0)).collect();
    println!("  Squared: {:?}", squared.data());

    // Using enumerate for indexed operations
    println!("\nEnumerate with indexed operations:");
    let indexed: Tensor = tensor
        .iter()
        .enumerate()
        .map(|(i, elem)| elem.add_scalar(i as f32))
        .collect();
    println!("  Indexed: {:?}", indexed.data());

    // Using fold for reduction
    println!("\nFold for sum calculation:");
    let sum: f32 = tensor.iter().fold(0.0, |acc, elem| acc + elem.value());
    println!("  Sum: {:.1}", sum);

    // Using find for element search
    println!("\nFind specific element:");
    if let Some(found) = tensor.iter().find(|elem| elem.value() == 3.0) {
        println!("  Found element with value 3.0: {:.1}", found.value());
    }

    // Using any/all for condition checking
    println!("\nCondition checking:");
    let all_positive = tensor.iter().all(|elem| elem.value() > 0.0);
    let any_large = tensor.iter().any(|elem| elem.value() > 4.0);
    println!("  All positive: {}", all_positive);
    println!("  Any > 4.0: {}", any_large);

    Ok(())
}

/// Demonstrate gradient tracking through element operations
///
/// Shows how gradient tracking works seamlessly through iterator
/// operations, maintaining the computational graph for backpropagation.
fn demonstrate_gradient_tracking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Gradient Tracking ---");

    // Create a tensor with gradient tracking enabled
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?.with_requires_grad();
    println!("Input tensor (requires_grad): {:?}", tensor.data());

    // Perform element-wise operations through iteration
    let result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Apply a complex transformation: (x^2 + 1) * 2
            elem.pow_scalar(2.0).add_scalar(1.0).mul_scalar(2.0)
        })
        .collect();

    println!("Result tensor: {:?}", result.data());
    println!("Result requires_grad: {}", result.requires_grad());

    // Compute gradients
    let mut loss = result.sum();
    loss.backward(None);

    println!("Loss: {:.6}", loss.value());
    println!("Input gradients: {:?}", tensor.grad().map(|g| g.data()));

    Ok(())
}

examples/optimizers/adam_configurations.rs (line 112)

fn demonstrate_default_adam() -> Result<(), Box<dyn std::error::Error>> {
    println!("--- Default Adam Configuration ---");

    // Create a simple regression problem: y = 2*x + 1
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(42)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create Adam optimizer with default configuration
    let mut optimizer = Adam::new();
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    println!("Default Adam configuration:");
    println!("  Learning rate: {}", optimizer.learning_rate());
    println!("  Initial weight: {:.6}", weight.value());
    println!("  Initial bias: {:.6}", bias.value());

    // Training loop
    let num_epochs = 50;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        losses.push(loss.value());

        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!("Epoch {:3}: Loss = {:.6}", epoch, loss.value());
        }
    }

    // Evaluate final model
    let _final_predictions = x_data.matmul(&weight) + &bias;
    println!("\nFinal model:");
    println!("  Learned weight: {:.6} (target: 2.0)", weight.value());
    println!("  Learned bias: {:.6} (target: 1.0)", bias.value());
    println!("  Final loss: {:.6}", losses[losses.len() - 1]);

    Ok(())
}

/// Demonstrate learning rate comparison
fn demonstrate_learning_rate_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Learning Rate Comparison ---");

    let learning_rates = [0.001, 0.01, 0.1];
    let mut results = Vec::new();

    for &lr in &learning_rates {
        println!("\nTesting learning rate: {}", lr);

        let stats = train_with_config(TrainingConfig {
            learning_rate: lr,
            ..Default::default()
        })?;

        results.push((lr, stats.clone()));

        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Convergence epoch: {}", stats.convergence_epoch);
    }

    // Compare results
    println!("\nLearning Rate Comparison Summary:");
    for (lr, stats) in &results {
        println!(
            "  LR={:6}: Loss={:.6}, Converged@{}",
            lr, stats.final_loss, stats.convergence_epoch
        );
    }

    Ok(())
}

/// Demonstrate weight decay comparison
fn demonstrate_weight_decay_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Weight Decay Comparison ---");

    let weight_decays = [0.0, 0.001, 0.01];
    let mut results = Vec::new();

    for &wd in &weight_decays {
        println!("\nTesting weight decay: {}", wd);

        let stats = train_with_config(TrainingConfig {
            weight_decay: wd,
            ..Default::default()
        })?;

        results.push((wd, stats.clone()));

        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Final weight norm: {:.6}", stats.weight_norm);
    }

    // Compare results
    println!("\nWeight Decay Comparison Summary:");
    for (wd, stats) in &results {
        println!(
            "  WD={:6}: Loss={:.6}, Weight Norm={:.6}",
            wd, stats.final_loss, stats.weight_norm
        );
    }

    Ok(())
}

/// Demonstrate beta parameter tuning
fn demonstrate_beta_parameter_tuning() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Beta Parameter Tuning ---");

    let beta_configs = [
        (0.9, 0.999),  // Default
        (0.8, 0.999),  // More aggressive momentum
        (0.95, 0.999), // Less aggressive momentum
        (0.9, 0.99),   // Faster second moment decay
    ];

    let mut results = Vec::new();

    for (i, (beta1, beta2)) in beta_configs.iter().enumerate() {
        println!(
            "\nTesting beta configuration {}: beta1={}, beta2={}",
            i + 1,
            beta1,
            beta2
        );

        let config = TrainingConfig {
            beta1: *beta1,
            beta2: *beta2,
            ..Default::default()
        };

        let stats = train_with_config(config)?;
        results.push(((*beta1, *beta2), stats.clone()));

        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Convergence epoch: {}", stats.convergence_epoch);
    }

    // Compare results
    println!("\nBeta Parameter Comparison Summary:");
    for ((beta1, beta2), stats) in &results {
        println!(
            "  B1={:4}, B2={:5}: Loss={:.6}, Converged@{}",
            beta1, beta2, stats.final_loss, stats.convergence_epoch
        );
    }

    Ok(())
}

/// Demonstrate configuration benchmarking
fn demonstrate_configuration_benchmarking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Configuration Benchmarking ---");

    // Define configurations to benchmark
    let configs = vec![
        (
            "Conservative",
            TrainingConfig {
                learning_rate: 0.001,
                weight_decay: 0.001,
                beta1: 0.95,
                ..Default::default()
            },
        ),
        (
            "Balanced",
            TrainingConfig {
                learning_rate: 0.01,
                weight_decay: 0.0,
                beta1: 0.9,
                ..Default::default()
            },
        ),
        (
            "Aggressive",
            TrainingConfig {
                learning_rate: 0.1,
                weight_decay: 0.0,
                beta1: 0.8,
                ..Default::default()
            },
        ),
    ];

    let mut benchmark_results = Vec::new();

    for (name, config) in configs {
        println!("\nBenchmarking {} configuration:", name);

        let start_time = std::time::Instant::now();
        let stats = train_with_config(config.clone())?;
        let elapsed = start_time.elapsed();

        println!("  Training time: {:.2}ms", elapsed.as_millis());
        println!("  Final loss: {:.6}", stats.final_loss);
        println!("  Convergence: {} epochs", stats.convergence_epoch);

        benchmark_results.push((name.to_string(), stats, elapsed));
    }

    // Summary
    println!("\nBenchmarking Summary:");
    for (name, stats, elapsed) in &benchmark_results {
        println!(
            "  {:12}: Loss={:.6}, Time={:4}ms, Converged@{}",
            name,
            stats.final_loss,
            elapsed.as_millis(),
            stats.convergence_epoch
        );
    }

    Ok(())
}

/// Helper function to train with specific configuration
fn train_with_config(config: TrainingConfig) -> Result<TrainingStats, Box<dyn std::error::Error>> {
    // Create training data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(123)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer with custom configuration
    let adam_config = AdamConfig {
        learning_rate: config.learning_rate,
        beta1: config.beta1,
        beta2: config.beta2,
        eps: 1e-8,
        weight_decay: config.weight_decay,
        amsgrad: false,
    };

    let mut optimizer = Adam::with_config(adam_config);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Training loop
    let mut losses = Vec::new();
    let mut convergence_epoch = config.epochs;

    for epoch in 0..config.epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        let loss_value = loss.value();
        losses.push(loss_value);

        // Check for convergence (loss < 0.01)
        if loss_value < 0.01 && convergence_epoch == config.epochs {
            convergence_epoch = epoch;
        }
    }

    Ok(TrainingStats {
        config,
        final_loss: losses[losses.len() - 1],
        loss_history: losses,
        convergence_epoch,
        weight_norm: weight.norm().value(),
    })
}

Additional examples can be found in:

• examples/optimizers/learning_rate_scheduling.rs
• examples/getting_started/optimizer_basics.rs
• examples/iterators/performance_optimization.rs
• examples/neural_networks/feedforward_network.rs
• examples/iterators/advanced_patterns.rs
• examples/RL_training/../neural_networks/basic_linear_layer.rs
• examples/RL_training/ppo_discrete.rs
• examples/RL_training/td3.rs

pub fn pow_tensor(&self, exponent: &Tensor) -> Tensor

Element-wise power with tensor exponents.

Computes element-wise power: output[i] = self[i]^exponent[i]

Arguments

• exponent - Tensor of exponents, must have the same shape as self

Returns

A new tensor with each element raised to the corresponding power

Examples

Basic Tensor Power

use train_station::Tensor;

let base = Tensor::from_slice(&[2.0, 3.0, 4.0], vec![3]).unwrap();
let exp = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let result = base.pow_tensor(&exp);
assert_eq!(result.shape().dims(), vec![3]);
assert_eq!(result.get(&[0]), 2.0); // 2.0^1.0 = 2.0
assert_eq!(result.get(&[1]), 9.0); // 3.0^2.0 = 9.0
assert_eq!(result.get(&[2]), 64.0); // 4.0^3.0 = 64.0

Mixed Exponents

use train_station::Tensor;

let base = Tensor::from_slice(&[4.0, 9.0, 16.0], vec![3]).unwrap();
let exp = Tensor::from_slice(&[0.5, 1.0, 2.0], vec![3]).unwrap();
let result = base.pow_tensor(&exp);
assert_eq!(result.shape().dims(), vec![3]);
assert_eq!(result.get(&[0]), 2.0); // sqrt(4.0) = 2.0
assert_eq!(result.get(&[1]), 9.0); // 9.0^1.0 = 9.0
assert_eq!(result.get(&[2]), 256.0); // 16.0^2.0 = 256.0

Panics

Panics if tensor shapes don’t match

impl Tensor

pub fn relu(&self) -> Tensor

Element-wise ReLU (Rectified Linear Unit) activation.

Applies ReLU to each element: output[i] = max(0, self[i])

Returns

A new tensor with ReLU applied to each element

Examples

Basic ReLU Activation

use train_station::Tensor;

let a = Tensor::from_slice(&[-1.0, 0.0, 2.5], vec![3]).unwrap();
let b = a.relu();
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 0.0); // max(0, -1.0) = 0.0
assert_eq!(b.get(&[1]), 0.0); // max(0, 0.0) = 0.0
assert_eq!(b.get(&[2]), 2.5); // max(0, 2.5) = 2.5

Mixed Positive and Negative Values

use train_station::Tensor;

let a = Tensor::from_slice(&[-5.0, -0.1, 0.0, 0.1, 5.0], vec![5]).unwrap();
let b = a.relu();
assert_eq!(b.shape().dims(), vec![5]);
assert_eq!(b.get(&[0]), 0.0); // max(0, -5.0) = 0.0
assert_eq!(b.get(&[1]), 0.0); // max(0, -0.1) = 0.0
assert_eq!(b.get(&[2]), 0.0); // max(0, 0.0) = 0.0
assert_eq!(b.get(&[3]), 0.1); // max(0, 0.1) = 0.1
assert_eq!(b.get(&[4]), 5.0); // max(0, 5.0) = 5.0

Examples found in repository

examples/neural_networks/feedforward_network.rs (line 55)

    pub fn forward(input: &Tensor) -> Tensor {
        input.relu()
    }

examples/supervised_training/supervised_bce.rs (line 60)

fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

examples/RL_training/ppo_continuous.rs (line 77)

    fn forward(&self, input: &Tensor) -> Tensor {
        let mut current: Option<Tensor> = None;
        for (i, layer) in self.layers.iter().enumerate() {
            let out = if i == 0 {
                layer.forward(input)
            } else {
                layer.forward(current.as_ref().unwrap())
            };
            let is_last = i + 1 == self.layers.len();
            let out = if !is_last { out.relu() } else { out };
            current = Some(out);
        }
        current.expect("MLP has at least one layer")
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.layers
            .iter_mut()
            .flat_map(|l| l.parameters())
            .collect()
    }
}

// -------------------------------
// Actor: mean = MLP(state); log_std is a learnable parameter tensor
// -------------------------------

struct Actor {
    net: Mlp,
    log_std: Tensor, // shape [action_dim]
}
impl Actor {
    fn new(state_dim: usize, action_dim: usize, seed: Option<u64>) -> Self {
        let net = Mlp::new(&[state_dim, 64, 64, action_dim], seed);
        let log_std = Tensor::from_slice(&vec![0.0; action_dim], vec![action_dim])
            .unwrap()
            .with_requires_grad();
        Self { net, log_std }
    }
    fn forward(&self, state: &Tensor) -> (Tensor, Tensor) {
        // Returns (mean [B, A], log_std [A])
        let mean = self.net.forward(state);
        (
            mean,
            self.log_std
                .view(vec![1, self.log_std.shape().dims()[0] as i32]),
        )
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        let mut ps = self.net.parameters();
        ps.push(&mut self.log_std);
        ps
    }
}

// -------------------------------
// Critic: value function V(s)
// -------------------------------

struct Critic {
    net: Mlp,
}
impl Critic {
    fn new(state_dim: usize, seed: Option<u64>) -> Self {
        Self {
            net: Mlp::new(&[state_dim, 64, 64, 1], seed),
        }
    }
    fn forward(&self, state: &Tensor) -> Tensor {
        self.net.forward(state)
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.net.parameters()
    }
}

// -------------------------------
// Continuous YardEnv (same dynamics as TD3 env)
// -------------------------------

struct YardEnv {
    pos: f32,
    vel: f32,
    steps: usize,
    max_steps: usize,
    rng: SmallRng,
}
impl YardEnv {
    fn new(seed: u64) -> Self {
        let mut e = Self {
            pos: 0.0,
            vel: 0.0,
            steps: 0,
            max_steps: 200,
            rng: SmallRng::new(seed),
        };
        e.reset();
        e
    }
    fn reset(&mut self) -> Tensor {
        self.pos = (self.rng.next_f32() * 1.0) - 0.5;
        self.vel = (self.rng.next_f32() * 0.2) - 0.1;
        self.steps = 0;
        self.state_tensor()
    }
    fn state_tensor(&self) -> Tensor {
        Tensor::from_slice(&[self.pos, self.vel, 0.0], vec![1, 3]).unwrap()
    }
    fn step(&mut self, action_value: f32) -> (Tensor, f32, bool) {
        let a = action_value.clamp(-1.0, 1.0);
        self.vel += 0.1 * a - 0.01 * self.pos;
        self.pos += self.vel;
        self.steps += 1;
        let reward = -(self.pos * self.pos) - 0.1 * (a * a);
        let done = self.pos.abs() > 3.0 || self.steps >= self.max_steps;
        (self.state_tensor(), reward, done)
    }
}

// -------------------------------
// Trajectory storage
// -------------------------------

struct RolloutBatch {
    states: Vec<f32>,
    actions: Vec<f32>,
    log_probs: Vec<f32>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    values: Vec<f32>,
    next_states: Vec<f32>,
    _state_dim: usize,
}
impl RolloutBatch {
    fn new(capacity: usize, state_dim: usize) -> Self {
        Self {
            states: Vec::with_capacity(capacity * state_dim),
            actions: Vec::with_capacity(capacity),
            log_probs: Vec::with_capacity(capacity),
            rewards: Vec::with_capacity(capacity),
            dones: Vec::with_capacity(capacity),
            values: Vec::with_capacity(capacity),
            next_states: Vec::with_capacity(capacity * state_dim),
            _state_dim: state_dim,
        }
    }

    #[allow(clippy::too_many_arguments)]
    fn push(&mut self, s: &[f32], a: f32, lp: f32, r: f32, d: f32, v: f32, s2: &[f32]) {
        self.states.extend_from_slice(s);
        self.actions.push(a);
        self.log_probs.push(lp);
        self.rewards.push(r);
        self.dones.push(d);
        self.values.push(v);
        self.next_states.extend_from_slice(s2);
    }

    fn len(&self) -> usize {
        self.actions.len()
    }
}

// -------------------------------
// Math helpers
// -------------------------------

fn gaussian_log_prob(action: &Tensor, mean: &Tensor, log_std: &Tensor) -> Tensor {
    // All tensors shaped [B, A] (log_std is broadcastable)
    let std = log_std.exp();
    let var = std.pow_scalar(2.0);
    let log_scale = log_std;
    let diff = action.sub_tensor(mean);
    let log_prob = diff
        .pow_scalar(2.0)
        .div_tensor(&var)
        .add_scalar(std::f32::consts::LN_2 + std::f32::consts::PI)
        .add_tensor(&log_scale.mul_scalar(2.0))
        .mul_scalar(0.5)
        .mul_scalar(-1.0);
    // Sum across action dim (dim=1) -> [B,1]
    log_prob.sum_dims(&[1], true)
}

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

examples/RL_training/ppo_discrete.rs (line 69)

    fn forward(&self, input: &Tensor) -> Tensor {
        let mut current: Option<Tensor> = None;
        for (i, layer) in self.layers.iter().enumerate() {
            let out = if i == 0 {
                layer.forward(input)
            } else {
                layer.forward(current.as_ref().unwrap())
            };
            let is_last = i + 1 == self.layers.len();
            let out = if !is_last { out.relu() } else { out };
            current = Some(out);
        }
        current.expect("MLP has at least one layer")
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.layers
            .iter_mut()
            .flat_map(|l| l.parameters())
            .collect()
    }
}

// -------------------------------
// Actor (logits) + Critic
// -------------------------------

struct Actor {
    net: Mlp,
}
impl Actor {
    fn new(state_dim: usize, action_dim: usize, seed: Option<u64>) -> Self {
        Self {
            net: Mlp::new(&[state_dim, 64, 64, action_dim], seed),
        }
    }
    fn forward(&self, state: &Tensor) -> Tensor {
        self.net.forward(state)
    } // logits [B, A]
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.net.parameters()
    }
}

struct Critic {
    net: Mlp,
}
impl Critic {
    fn new(state_dim: usize, seed: Option<u64>) -> Self {
        Self {
            net: Mlp::new(&[state_dim, 64, 64, 1], seed),
        }
    }
    fn forward(&self, state: &Tensor) -> Tensor {
        self.net.forward(state)
    }
    fn parameters(&mut self) -> Vec<&mut Tensor> {
        self.net.parameters()
    }
}

// -------------------------------
// Discrete YardEnv (3 actions -> -1, 0, +1)
// -------------------------------

struct YardEnv {
    pos: f32,
    vel: f32,
    steps: usize,
    max_steps: usize,
    rng: SmallRng,
}
impl YardEnv {
    const ACTIONS: [f32; 3] = [-1.0, 0.0, 1.0];
    fn new(seed: u64) -> Self {
        let mut e = Self {
            pos: 0.0,
            vel: 0.0,
            steps: 0,
            max_steps: 200,
            rng: SmallRng::new(seed),
        };
        e.reset();
        e
    }
    fn reset(&mut self) -> Tensor {
        self.pos = (self.rng.next_f32() * 1.0) - 0.5;
        self.vel = (self.rng.next_f32() * 0.2) - 0.1;
        self.steps = 0;
        self.state_tensor()
    }
    fn state_tensor(&self) -> Tensor {
        Tensor::from_slice(&[self.pos, self.vel, 0.0], vec![1, 3]).unwrap()
    }
    fn step(&mut self, action_idx: usize) -> (Tensor, f32, bool) {
        let a = Self::ACTIONS[action_idx.min(2)];
        self.vel += 0.1 * a - 0.01 * self.pos;
        self.pos += self.vel;
        self.steps += 1;
        let reward = -(self.pos * self.pos) - 0.05 * (a * a);
        let done = self.pos.abs() > 3.0 || self.steps >= self.max_steps;
        (self.state_tensor(), reward, done)
    }
}

// -------------------------------
// Rollout storage
// -------------------------------

struct RolloutBatch {
    states: Vec<f32>,
    actions: Vec<usize>,
    old_logps: Vec<f32>,
    rewards: Vec<f32>,
    dones: Vec<f32>,
    values: Vec<f32>,
    next_states: Vec<f32>,
    _state_dim: usize,
}
impl RolloutBatch {
    fn new(cap: usize, sd: usize) -> Self {
        Self {
            states: Vec::with_capacity(cap * sd),
            actions: Vec::with_capacity(cap),
            old_logps: Vec::with_capacity(cap),
            rewards: Vec::with_capacity(cap),
            dones: Vec::with_capacity(cap),
            values: Vec::with_capacity(cap),
            next_states: Vec::with_capacity(cap * sd),
            _state_dim: sd,
        }
    }
    #[allow(clippy::too_many_arguments)]
    fn push(&mut self, s: &[f32], a: usize, lp: f32, r: f32, d: f32, v: f32, s2: &[f32]) {
        self.states.extend_from_slice(s);
        self.actions.push(a);
        self.old_logps.push(lp);
        self.rewards.push(r);
        self.dones.push(d);
        self.values.push(v);
        self.next_states.extend_from_slice(s2);
    }
    fn len(&self) -> usize {
        self.actions.len()
    }
}

// -------------------------------
// Helpers
// -------------------------------

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

// log-softmax for selected actions: given logits [B,A] and actions Vec<usize> -> log_prob [B,1]
fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

// probability ratio = exp(new_logp - old_logp)
fn ratio_from_logps(new_logp: &Tensor, old_logp: &Tensor) -> Tensor {
    new_logp.sub_tensor(old_logp).exp()
}

// Clamp ratio to [1-clip, 1+clip] using ReLU-based clamp (no custom ops)
fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/neural_networks/basic_encoder.rs (line 60)

    pub fn forward(&self, input: &Tensor, attn_mask: Option<&Tensor>) -> Tensor {
        let attn = self.mha.forward(input, input, input, attn_mask);
        let res1 = attn.add_tensor(input);

        // Feed-forward network with ReLU and residual
        let (b, t, e) = Self::triple(input);
        let x2d = res1.contiguous().view(vec![(b * t) as i32, e as i32]);
        let hidden = self.ffn_in.forward(&x2d).relu();
        let out2d = self.ffn_out.forward(&hidden);
        let out = out2d.view(vec![b as i32, t as i32, e as i32]);
        out.add_tensor(&res1)
    }

examples/RL_training/dqn.rs (line 81)

    fn forward(&self, input: &Tensor, final_activation: Option<fn(&Tensor) -> Tensor>) -> Tensor {
        let mut current: Option<Tensor> = None;
        for (i, layer) in self.layers.iter().enumerate() {
            let out = if i == 0 {
                layer.forward(input)
            } else {
                layer.forward(current.as_ref().unwrap())
            };
            let is_last = i + 1 == self.layers.len();
            let out = if !is_last {
                out.relu()
            } else if let Some(act) = final_activation {
                act(&out)
            } else {
                out
            };
            current = Some(out);
        }
        current.expect("MLP has at least one layer")
    }

Additional examples can be found in:

• examples/RL_training/td3.rs
• examples/neural_networks/basic_decoder.rs

impl Tensor

pub fn sigmoid(&self) -> Tensor

Element-wise sigmoid activation function

Computes the sigmoid function for each element: output[i] = 1 / (1 + e^(-self[i]))

Uses a numerically stable implementation that avoids overflow for large positive/negative values by using different computation paths for positive and negative inputs.

Returns

A new tensor with sigmoid applied to each element, values in range (0, 1)

Performance Characteristics

• Numerical Stability: Avoids overflow using stable implementation
• Scalar Implementation: Optimized scalar computation for mathematical accuracy
• Cache-friendly: Linear memory access patterns
• Mathematical Accuracy: High-precision exponential and division operations
• Gradient Tracking: Full gradtrack support with efficient gradient computation

Implementation Details

Uses a numerically stable implementation:

• For x ≥ 0: computes 1 / (1 + e^(-x)) to avoid overflow in e^x for large positive x
• For x < 0: computes e^x / (1 + e^x) to avoid overflow in e^(-x) for large negative x This ensures the result is always in the range (0, 1) without numerical overflow.

Examples

Basic Sigmoid Activation

use train_station::Tensor;

let a = Tensor::from_slice(&[-1.0, 0.0, 1.0], vec![3]).unwrap();
let b = a.sigmoid();
assert_eq!(b.shape().dims(), vec![3]);
assert!((b.get(&[0]) - 0.26894143).abs() < 1e-6); // sigmoid(-1.0)
assert!((b.get(&[1]) - 0.5).abs() < 1e-6); // sigmoid(0.0)
assert!((b.get(&[2]) - 0.7310586).abs() < 1e-6); // sigmoid(1.0)

Extreme Values

use train_station::Tensor;

let a = Tensor::from_slice(&[-10.0, 10.0], vec![2]).unwrap();
let b = a.sigmoid();
assert_eq!(b.shape().dims(), vec![2]);
assert!(b.get(&[0]) < 1e-4); // sigmoid(-10.0) ≈ 0
assert!(b.get(&[1]) > 0.9999); // sigmoid(10.0) ≈ 1

Examples found in repository

examples/supervised_training/supervised_bce.rs (line 142)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Supervised FFN Example (XOR) ===");

    // Dataset: XOR (repeat to form a small batch)
    let inputs: Vec<f32> = vec![
        0.0, 0.0, // -> 0
        0.0, 1.0, // -> 1
        1.0, 0.0, // -> 1
        1.0, 1.0, // -> 0
    ];
    let targets: Vec<f32> = vec![0.0, 1.0, 1.0, 0.0];

    // Repeat the base patterns to stabilize training
    let repeats = 64usize; // effective batch = 4 * repeats = 256
    let mut xs = Vec::with_capacity(repeats * inputs.len());
    let mut ys = Vec::with_capacity(repeats * targets.len());
    for _ in 0..repeats {
        xs.extend_from_slice(&inputs);
        ys.extend_from_slice(&targets);
    }

    let batch = xs.len() / 2; // two features
    let x_t = Tensor::from_slice(&xs, vec![batch, 2]).unwrap();
    let y_t = Tensor::from_slice(&ys, vec![batch, 1]).unwrap();

    // Model config: 2 -> 32 -> 32 -> 1, final sigmoid via loss path
    let cfg = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![32, 32],
        output_size: 1,
        use_bias: true,
    };
    let mut net = FeedForwardNetwork::new(cfg, Some(777));

    // Optimizer and parameter linking
    let mut opt = Adam::with_learning_rate(1e-3);
    for p in net.parameters() {
        opt.add_parameter(p);
    }

    let epochs = 1000usize;
    let max_grad_norm = 1.0f32;
    let mut best_loss = f32::INFINITY;
    let mut best_acc = 0.0f32;

    for e in 0..epochs {
        // Zero grads each iteration
        {
            let mut params = net.parameters();
            opt.zero_grad(&mut params);
        }

        // Forward -> logits; use numerically stable BCE-with-logits for loss
        let logits = net.forward(&x_t);
        let mut loss = bce_with_logits(&logits, &y_t);
        loss.backward(None);

        // Step only params with grads
        {
            let params = net.parameters();
            let mut with_grads: Vec<&mut Tensor> = Vec::new();
            for p in params {
                if p.grad_owned().is_some() {
                    with_grads.push(p);
                }
            }
            if !with_grads.is_empty() {
                clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                opt.step(&mut with_grads);
                opt.zero_grad(&mut with_grads);
            }
        }

        // Metrics (use sigmoid only for reporting accuracy)
        let preds = logits.sigmoid();
        let acc = accuracy(&preds, &y_t);
        if loss.value() < best_loss {
            best_loss = loss.value();
        }
        if acc > best_acc {
            best_acc = acc;
        }
        if e % 10 == 0 || e + 1 == epochs {
            println!(
                "epoch {:4} | loss={:.5} acc={:.3} | best_loss={:.5} best_acc={:.3}",
                e,
                loss.value(),
                acc,
                best_loss,
                best_acc
            );
        }

        // Clear graphs to avoid stale accumulation across epochs
        clear_all_graphs_known();
    }

    // Quick sanity check predictions
    let test = Tensor::from_slice(&inputs, vec![4, 2]).unwrap();
    let out = net.forward(&test).sigmoid();
    println!("predictions (approx): {:?}", out.data());

    println!("=== Supervised training finished ===");
    Ok(())
}

impl Tensor

pub fn softmax(&self, dim: usize) -> Tensor

Computes softmax activation along the specified dimension

Applies the softmax function along dimension dim, transforming values into probabilities that sum to 1 along that dimension. Uses numerically stable computation to avoid overflow: softmax(x_i) = exp(x_i - max(x)) / sum(exp(x_j - max(x)))

Arguments

• dim - Dimension along which to compute softmax (0-based indexing)

Returns

A new tensor with softmax applied along the specified dimension. Values are in range (0, 1) and sum to 1 along dim.

Performance Characteristics

• Numerical Stability: Avoids overflow using max subtraction technique
• Scalar Implementation: Optimized scalar computation for mathematical accuracy
• Cache-friendly: Optimized memory access patterns for dimension operations
• Mathematical Accuracy: High-precision exponential and division operations
• GradTrack Support: Full automatic differentiation with efficient gradient computation

Implementation Details

Uses a numerically stable three-pass algorithm:

1. Max Computation: Find the maximum value along the specified dimension
2. Exponential Sum: Compute exp(x - max) and sum the results
3. Normalization: Divide each exp(x - max) by the sum to get probabilities

This approach prevents overflow by subtracting the maximum value before computing exponentials, ensuring numerical stability for any input range.

Examples

Basic Softmax Activation

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let b = a.softmax(0);
assert_eq!(b.shape().dims(), vec![3]);

// Verify probabilities sum to 1
let sum = b.get(&[0]) + b.get(&[1]) + b.get(&[2]);
assert!((sum - 1.0).abs() < 1e-6);

// Verify relative ordering is preserved
assert!(b.get(&[0]) < b.get(&[1]));
assert!(b.get(&[1]) < b.get(&[2]));

2D Softmax Along Different Dimensions

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let b = a.softmax(0); // Softmax along first dimension
assert_eq!(b.shape().dims(), vec![2, 2]);

// Each column should sum to 1
let col1_sum = b.get(&[0, 0]) + b.get(&[1, 0]);
let col2_sum = b.get(&[0, 1]) + b.get(&[1, 1]);
assert!((col1_sum - 1.0).abs() < 1e-6);
assert!((col2_sum - 1.0).abs() < 1e-6);

Panics

• Panics if dim is out of bounds for the tensor’s rank
• Panics if the dimension size is 0

Examples found in repository

examples/neural_networks/multi_head_attention.rs (line 108)

    pub fn forward(
        &self,
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        attn_mask: Option<&Tensor>,
    ) -> Tensor {
        let qkv = Self::project_qkv(query, key, value, &self.q_proj, &self.k_proj, &self.v_proj);
        let (q, k, v) = qkv;

        // Split heads: [b, t, e] -> [b, h, t, d]
        let (b, tq, _e) = Self::triple(query);
        let (_b2, tk, _e2) = Self::triple(key);
        let q = Self::split_heads(&q, b, tq, self.num_heads, self.head_dim);
        let k = Self::split_heads(&k, b, tk, self.num_heads, self.head_dim);
        let v = Self::split_heads(&v, b, tk, self.num_heads, self.head_dim);

        // Scaled dot-product attention
        // logits: [b, h, tq, tk]
        let k_t = k.transpose(2, 3);
        let mut logits = q.matmul(&k_t).div_scalar((self.head_dim as f32).sqrt());
        if let Some(mask) = attn_mask {
            let dims = mask.shape().dims().to_vec();
            // If boolean-like mask matching [b,h,tq,tk], apply masked_fill
            if dims.len() == 4 && dims[0] == b && dims[1] == self.num_heads && dims[2] == tq {
                // Interpret mask > 0.5 as keep; we invert to build masked positions
                let cond: Vec<bool> = mask.data().iter().map(|&v| v < 0.5).collect();
                // Apply masked fill on a flattened view, then reshape back
                let flat_logits = logits.view(vec![(b * self.num_heads * tq * tk) as i32]);
                let filled = flat_logits.masked_fill(&cond, f32::NEG_INFINITY);
                logits = filled.view(vec![b as i32, self.num_heads as i32, tq as i32, tk as i32]);
            } else {
                // Fallback: additive mask
                logits = logits.add_tensor(mask);
            }
        }
        let attn = logits.softmax(3);

        // context: [b, h, tq, d]
        let context = attn.matmul(&v);
        let context = context.permute(vec![0, 2, 1, 3]); // [b, tq, h, d]
        let context = context.contiguous().view(vec![
            b as i32,
            tq as i32,
            (self.num_heads * self.head_dim) as i32,
        ]);

        // Output projection (flatten to 2D, project, then restore 3D)
        let flat = context.view(vec![(b * tq) as i32, self.embed_dim as i32]);
        let out2d = self.out_proj.forward(&flat);
        out2d.view(vec![b as i32, tq as i32, self.embed_dim as i32])
    }

examples/supervised_training/supervised_classification.rs (line 231)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Supervised Classification Example (Cross-Entropy) ===");

    // Synthetic 2D inputs, 3 classes with linear-ish separations
    let n = 1200usize;
    let classes = 3usize;
    let mut xs: Vec<f32> = Vec::with_capacity(n * 2);
    let mut ys: Vec<usize> = Vec::with_capacity(n);

    // Simple RNG
    let mut state: u64 = 424242;
    let mut rand_f32 = || {
        state = state.wrapping_mul(1664525).wrapping_add(1013904223);
        (state >> 16) as f32 / (u32::MAX as f32)
    };

    for _ in 0..n {
        let x1 = rand_f32() * 4.0 - 2.0;
        let x2 = rand_f32() * 4.0 - 2.0;
        // Class by quadrant-ish rule with noise
        let mut c = if x1 + 0.5 * x2 > 0.5 {
            0
        } else if x1 - x2 < -0.5 {
            1
        } else {
            2
        };
        if rand_f32() < 0.05 {
            c = (c + 1) % classes;
        }
        xs.push(x1);
        xs.push(x2);
        ys.push(c);
    }

    // Normalize inputs per-feature to [-1, 1]
    let mut min1 = f32::INFINITY;
    let mut max1 = f32::NEG_INFINITY;
    let mut min2 = f32::INFINITY;
    let mut max2 = f32::NEG_INFINITY;
    for i in (0..xs.len()).step_by(2) {
        let a = xs[i];
        let b = xs[i + 1];
        if a < min1 {
            min1 = a;
        }
        if a > max1 {
            max1 = a;
        }
        if b < min2 {
            min2 = b;
        }
        if b > max2 {
            max2 = b;
        }
    }
    let rng1 = (max1 - min1).max(1e-8);
    let rng2 = (max2 - min2).max(1e-8);
    for i in (0..xs.len()).step_by(2) {
        let a = xs[i];
        let b = xs[i + 1];
        xs[i] = 2.0 * (a - min1) / rng1 - 1.0;
        xs[i + 1] = 2.0 * (b - min2) / rng2 - 1.0;
    }

    // Train/Val split (80/20)
    let n_train = (n as f32 * 0.8) as usize;
    let x_train = Tensor::from_slice(&xs[..n_train * 2], vec![n_train, 2]).unwrap();
    let y_train = ys[..n_train].to_vec();
    let x_val = Tensor::from_slice(&xs[n_train * 2..], vec![n - n_train, 2]).unwrap();
    let y_val = ys[n_train..].to_vec();

    // Model: 2 -> 64 -> 64 -> 3 (logits)
    let cfg = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![64, 64],
        output_size: classes,
        use_bias: true,
    };
    let mut net = FeedForwardNetwork::new(cfg, Some(303));

    // Optimizer
    let mut opt = Adam::with_learning_rate(1e-3);
    for p in net.parameters() {
        opt.add_parameter(p);
    }

    let epochs = 300usize;
    let max_grad_norm = 1.0f32;
    let mut best_val_acc = 0.0f32;
    let mut best_val_loss = f32::INFINITY;

    for e in 0..epochs {
        // Zero grads
        {
            let mut params = net.parameters();
            opt.zero_grad(&mut params);
        }

        // Forward logits
        let logits = net.forward(&x_train);
        let mut loss = cross_entropy_logits(&logits, &y_train, n_train, classes);
        loss.backward(None);

        // Step clipped
        {
            let params = net.parameters();
            let mut with_grads: Vec<&mut Tensor> = Vec::new();
            for p in params {
                if p.grad_owned().is_some() {
                    with_grads.push(p);
                }
            }
            if !with_grads.is_empty() {
                clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                opt.step(&mut with_grads);
                opt.zero_grad(&mut with_grads);
            }
        }

        // Metrics
        let train_acc = accuracy_from_logits(&logits, &y_train, n_train, classes);
        let val_logits = net.forward(&x_val);
        let val_loss = cross_entropy_logits(&val_logits, &y_val, n - n_train, classes).value();
        let val_acc = accuracy_from_logits(&val_logits, &y_val, n - n_train, classes);
        if val_acc > best_val_acc {
            best_val_acc = val_acc;
        }
        if val_loss < best_val_loss {
            best_val_loss = val_loss;
        }

        if e % 10 == 0 || e + 1 == epochs {
            println!(
                "epoch {:4} | loss={:.4} acc={:.3} | val_loss={:.4} val_acc={:.3} | best_val_acc={:.3}",
                e, loss.value(), train_acc, val_loss, val_acc, best_val_acc
            );
        }

        clear_all_graphs_known();
    }

    // Quick sample preds via softmax
    let samples = Tensor::from_slice(&[-1.0, -1.0, 0.0, 0.0, 1.0, 1.0], vec![3, 2]).unwrap();
    let sm = net.forward(&samples).softmax(1);
    println!("sample class probs: {:?}", sm.data());

    println!("=== Supervised classification finished ===");
    Ok(())
}

examples/RL_training/ppo_discrete.rs (line 364)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

impl Tensor

pub fn sqrt(&self) -> Tensor

Element-wise square root

Computes the square root for each element: output[i] = sqrt(self[i])

Uses SIMD optimization when available for maximum performance, with automatic fallback to optimized scalar computation for non-SIMD hardware.

Returns

A new tensor with the square root of each element

Performance Characteristics

• SIMD Optimization: AVX2-optimized with 32-element blocks and 4x unrolling
• Scalar Fallback: 4x unrolled scalar implementation for non-SIMD hardware
• Cache-friendly: Linear memory access patterns
• Mathematical Accuracy: High-precision square root computation
• GradTrack Support: Full automatic differentiation with efficient gradient computation

Implementation Details

Automatically selects between SIMD and scalar implementations based on hardware capabilities. SIMD implementation uses AVX2 vector square root operations for optimal performance. Scalar implementation uses 4x unrolling for better instruction-level parallelism.

Examples

Basic Square Root

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 4.0, 9.0], vec![3]).unwrap();
let b = a.sqrt();
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 1.0); // sqrt(1.0) = 1.0
assert_eq!(b.get(&[1]), 2.0); // sqrt(4.0) = 2.0
assert_eq!(b.get(&[2]), 3.0); // sqrt(9.0) = 3.0

Zero and Special Values

use train_station::Tensor;

let a = Tensor::from_slice(&[0.0, 1.0, 16.0], vec![3]).unwrap();
let b = a.sqrt();
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 0.0); // sqrt(0.0) = 0.0
assert_eq!(b.get(&[1]), 1.0); // sqrt(1.0) = 1.0
assert_eq!(b.get(&[2]), 4.0); // sqrt(16.0) = 4.0

Note

Results are undefined for negative values (may produce NaN)

Examples found in repository

examples/supervised_training/supervised_regression.rs (line 47)

fn rmse(pred: &Tensor, target: &Tensor) -> f32 {
    mse(pred, target).sqrt().value()
}

examples/RL_training/dqn.rs (line 324)

fn pseudo_huber_mean(diff: &Tensor) -> Tensor {
    diff.pow_scalar(2.0)
        .add_scalar(1.0)
        .sqrt()
        .sub_scalar(1.0)
        .mean()
}

examples/iterators/performance_optimization.rs (line 179)

fn demonstrate_memory_optimization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Memory Optimization ---");

    // Create a large tensor for memory testing
    let size = 10000;
    let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
    let tensor = Tensor::from_slice(&data, vec![size])?;

    println!("Processing tensor of size: {}", size);

    // Pattern 1: Streaming processing with iterator chunks (process in blocks, collect with shape)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();
    let flattened = tensor.view(vec![size as i32]);
    let _streamed_result: Tensor = flattened
        .chunks(chunk_size)
        .map(|c| c.pow_scalar(2.0).sqrt())
        .collect_shape(vec![size]);
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter_elements()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect_shape(vec![size]);
    let full_time = start.elapsed();

    println!("  Streaming time: {:?}", streamed_time);
    println!("  Full processing time: {:?}", full_time);
    println!(
        "  Memory efficiency ratio: {:.2}x",
        full_time.as_nanos() as f64 / streamed_time.as_nanos() as f64
    );

    // Pattern 3: Lazy evaluation with take
    println!("\nPattern 2: Lazy Evaluation");
    let start = Instant::now();
    let lazy_result: Tensor = tensor
        .iter_elements()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect_shape(vec![1000]);
    let lazy_time = start.elapsed();

    println!("  Lazy processing (1000 elements): {:?}", lazy_time);
    println!("  Lazy result size: {}", lazy_result.size());

    // Pattern 4: Memory-efficient filtering
    println!("\nPattern 3: Memory-Efficient Filtering");
    let start = Instant::now();
    let filtered_result: Tensor = tensor
        .iter_elements()
        .filter(|elem| elem.value() > size as f32 / 2.0) // Keep only large values
        .map(|elem| elem.mul_scalar(2.0))
        .collect();
    let filtered_time = start.elapsed();

    println!("  Filtered processing: {:?}", filtered_time);
    println!(
        "  Filtered result size: {} (reduced from {})",
        filtered_result.size(),
        size
    );

    Ok(())
}

/// Demonstrate large-scale processing techniques
///
/// Shows how to efficiently process very large datasets using
/// iterator patterns and optimization strategies.
fn demonstrate_large_scale_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Large-Scale Processing ---");

    // Simulate large dataset processing
    let sizes = vec![10000, 50000, 100000];

    for size in sizes {
        println!("\nProcessing dataset of size: {}", size);

        // Generate large dataset
        let data: Vec<f32> = (0..size)
            .map(|i| {
                let x = i as f32 / size as f32;
                x * x + 0.1 * (i % 10) as f32 // Quadratic with noise
            })
            .collect();

        let tensor = Tensor::from_slice(&data, vec![size])?;

        // Technique 1: Batch processing
        let batch_size = 1000;
        let start = Instant::now();

        let mut batch_results = Vec::new();
        for batch_start in (0..size).step_by(batch_size) {
            let batch_end = (batch_start + batch_size).min(size);
            let batch: Tensor = tensor
                .iter_range(batch_start, batch_end)
                .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
                .collect();
            batch_results.push(batch);
        }
        let batch_time = start.elapsed();

        // Technique 2: Parallel-like processing with stride
        let start = Instant::now();
        let stride = 4;
        let strided_result: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % stride == 0)
            .map(|(_, elem)| elem.pow_scalar(2.0).add_scalar(1.0))
            .collect();
        let strided_time = start.elapsed();

        // Technique 3: Hierarchical processing
        let start = Instant::now();
        let coarse: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % 10 == 0) // Every 10th element
            .map(|(_, elem)| elem.pow_scalar(2.0).add_scalar(1.0))
            .collect();
        let fine: Tensor = tensor
            .iter()
            .enumerate()
            .filter(|(i, _)| i % 10 != 0) // Rest of elements
            .map(|(_, elem)| elem.pow_scalar(1.5).add_scalar(0.5))
            .collect();
        let hierarchical_time = start.elapsed();

        // Report performance
        println!("  Batch processing: {:?}", batch_time);
        println!("  Strided processing: {:?}", strided_time);
        println!("  Hierarchical processing: {:?}", hierarchical_time);

        // Memory usage analysis
        let total_batches = size.div_ceil(batch_size);
        println!("  Batch count: {}", total_batches);
        println!("  Strided result size: {}", strided_result.size());
        println!(
            "  Hierarchical: coarse={}, fine={}",
            coarse.size(),
            fine.size()
        );
    }

    Ok(())
}

/// Demonstrate advanced optimization techniques
///
/// Shows sophisticated optimization strategies and techniques
/// for maximizing performance in tensor iterator operations.
fn demonstrate_optimization_techniques() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Optimization Techniques ---");

    let size = 50000;
    let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
    let tensor = Tensor::from_slice(&data, vec![size])?;

    println!("Optimizing processing for size: {}", size);

    // Technique 1: Operation fusion
    println!("\nTechnique 1: Operation Fusion");
    let start = Instant::now();
    let fused_result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Fuse multiple operations into single chain
            elem.mul_scalar(2.0).add_scalar(1.0).pow_scalar(2.0).sqrt()
        })
        .collect();
    let fused_time = start.elapsed();

    // Technique 2: Conditional optimization
    println!("\nTechnique 2: Conditional Optimization");
    let start = Instant::now();
    let conditional_result: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val < size as f32 / 2.0 {
                elem.mul_scalar(2.0) // Simple operation for small values
            } else {
                elem.pow_scalar(2.0).sqrt() // Complex operation for large values
            }
        })
        .collect();
    let conditional_time = start.elapsed();

    // Technique 3: Cache-friendly processing
    println!("\nTechnique 3: Cache-Friendly Processing");
    let start = Instant::now();
    let cache_friendly_result: Tensor = tensor
        .iter()
        .take(1000) // Process in cache-friendly chunks
        .map(|elem| elem.mul_scalar(2.0))
        .collect();
    let cache_friendly_time = start.elapsed();

    // Technique 4: Memory pooling simulation
    println!("\nTechnique 4: Memory Pooling Simulation");
    let start = Instant::now();
    let pooled_result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 100 == 0) // Process every 100th element
        .map(|(_, elem)| elem.pow_scalar(2.0))
        .collect();
    let pooled_time = start.elapsed();

    // Report optimization results
    println!("  Fused operations: {:?}", fused_time);
    println!("  Conditional optimization: {:?}", conditional_time);
    println!("  Cache-friendly processing: {:?}", cache_friendly_time);
    println!("  Memory pooling simulation: {:?}", pooled_time);

    // Performance analysis
    let fastest = fused_time
        .min(conditional_time)
        .min(cache_friendly_time)
        .min(pooled_time);
    println!("  Fastest technique: {:?}", fastest);

    // Memory efficiency analysis
    println!("  Fused result size: {}", fused_result.size());
    println!("  Conditional result size: {}", conditional_result.size());
    println!(
        "  Cache-friendly result size: {}",
        cache_friendly_result.size()
    );
    println!("  Pooled result size: {}", pooled_result.size());

    // Technique 5: Gradient optimization
    println!("\nTechnique 5: Gradient Optimization");
    let grad_tensor = tensor.with_requires_grad();
    let start = Instant::now();

    let grad_result: Tensor = grad_tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
        .collect();

    let mut loss = grad_result.sum();
    loss.backward(None);
    let grad_time = start.elapsed();

    println!("  Gradient computation: {:?}", grad_time);
    println!(
        "  Gradient tracking enabled: {}",
        grad_result.requires_grad()
    );

    Ok(())
}

examples/iterators/advanced_patterns.rs (line 191)

fn demonstrate_conditional_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Conditional Processing ---");

    // Create data with mixed characteristics
    let data = vec![1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0, 9.0, -10.0];
    let tensor = Tensor::from_slice(&data, vec![10])?;
    println!("Input data: {:?}", tensor.data());

    // Conditional transformation based on sign
    println!("\nConditional transformation (positive/negative handling):");
    let processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val > 0.0 {
                elem.pow_scalar(2.0) // Square positive values
            } else {
                elem.mul_scalar(-1.0).sqrt() // Square root of absolute negative values
            }
        })
        .collect();
    println!("  Processed: {:?}", processed.data());

    // Adaptive filtering based on local statistics
    println!("\nAdaptive filtering (remove values > 2 std from local mean):");
    let window_size = 3;
    let adaptive_filtered: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, elem)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(tensor.size());

            // Calculate local mean and std
            let local_values: Vec<f32> = (start..end)
                .map(|j| tensor.element_view(j).value())
                .collect();

            let local_mean = local_values.iter().sum::<f32>() / local_values.len() as f32;
            let local_variance = local_values
                .iter()
                .map(|v| (v - local_mean).powi(2))
                .sum::<f32>()
                / local_values.len() as f32;
            let local_std = local_variance.sqrt();

            let threshold = local_mean + 2.0 * local_std;
            elem.value() <= threshold
        })
        .map(|(_, elem)| elem)
        .collect();
    println!("  Adaptive filtered: {:?}", adaptive_filtered.data());

    // Multi-condition processing
    println!("\nMulti-condition processing:");
    let multi_processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            match () {
                _ if val > 5.0 => elem.mul_scalar(2.0), // Double large values
                _ if val < -5.0 => elem.div_scalar(2.0), // Halve small values
                _ if val.abs() < 2.0 => elem.add_scalar(1.0), // Add 1 to small values
                _ => elem.clone(),                      // Keep others unchanged
            }
        })
        .collect();
    println!("  Multi-condition: {:?}", multi_processed.data());

    Ok(())
}

impl Tensor

pub fn sub_tensor(&self, other: &Tensor) -> Tensor

Element-wise subtraction with another tensor with broadcasting support

Performs element-wise subtraction with automatic broadcasting: output[i] = self[i] - other[i]

Broadcasting enables subtraction between tensors of different but compatible shapes. Compatible shapes follow NumPy broadcasting rules:

• Dimensions are aligned from the rightmost dimension
• Dimensions are compatible if they are equal, or one of them is 1
• Missing dimensions are treated as 1

Arguments

• other - Tensor to subtract. Shapes must be broadcast-compatible.

Returns

A new tensor containing the element-wise difference with broadcast result shape

Performance Characteristics

• Fast Path: Optimized for identical shapes to avoid broadcasting overhead
• SIMD Optimization: AVX2-optimized with 32-element blocks and 4x unrolling
• Broadcasting: Efficient broadcasting for compatible shapes
• Cache-friendly: Linear memory access patterns
• GradTrack Support: Full automatic differentiation with efficient gradient computation

Implementation Details

Uses a fast path for identical shapes to avoid broadcasting overhead. For different shapes, performs broadcasting followed by optimized element-wise subtraction. Automatically selects between SIMD and scalar implementations based on hardware capabilities.

Examples

Same Shape Subtraction

use train_station::Tensor;

let a = Tensor::from_slice(&[5.0, 7.0, 9.0], vec![3]).unwrap();
let b = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let c = a.sub_tensor(&b);
assert_eq!(c.shape().dims(), vec![3]);
assert_eq!(c.get(&[0]), 4.0); // 5.0 - 1.0
assert_eq!(c.get(&[1]), 5.0); // 7.0 - 2.0
assert_eq!(c.get(&[2]), 6.0); // 9.0 - 3.0

Broadcasting Subtraction

use train_station::Tensor;

let a = Tensor::from_slice(&[5.0, 10.0], vec![2, 1]).unwrap();
let b = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3]).unwrap();
let c = a.sub_tensor(&b);
assert_eq!(c.shape().dims(), vec![2, 3]);
// Result: [[4.0, 3.0, 2.0], [9.0, 8.0, 7.0]]
assert_eq!(c.get(&[0, 0]), 4.0); // 5.0 - 1.0
assert_eq!(c.get(&[0, 1]), 3.0); // 5.0 - 2.0
assert_eq!(c.get(&[1, 0]), 9.0); // 10.0 - 1.0

Scalar Subtraction

use train_station::Tensor;

let a = Tensor::ones(vec![2, 3]);
let b = Tensor::from_slice(&[0.5], vec![1]).unwrap();
let c = a.sub_tensor(&b);
assert_eq!(c.shape().dims(), vec![2, 3]);
assert_eq!(c.get(&[0, 0]), 0.5); // 1.0 - 0.5

Panics

Panics if tensor shapes are not broadcast-compatible

Examples found in repository

examples/supervised_training/supervised_regression.rs (line 43)

fn mse(pred: &Tensor, target: &Tensor) -> Tensor {
    pred.sub_tensor(target).pow_scalar(2.0).mean()
}

fn rmse(pred: &Tensor, target: &Tensor) -> f32 {
    mse(pred, target).sqrt().value()
}

fn r2_score(pred: &Tensor, target: &Tensor) -> f32 {
    // R^2 = 1 - SS_res / SS_tot
    let y = target;
    let y_mean = y.mean();
    let ss_res = pred.sub_tensor(y).pow_scalar(2.0).sum();
    let ss_tot = y.sub_tensor(&y_mean).pow_scalar(2.0).sum();
    let ss_res_v = ss_res.value();
    let ss_tot_v = ss_tot.value().max(1e-12); // avoid divide by zero
    1.0 - (ss_res_v / ss_tot_v)
}

examples/supervised_training/supervised_bce.rs (line 65)

fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

examples/supervised_training/supervised_classification.rs (line 52)

fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

examples/RL_training/ppo_discrete.rs (line 280)

fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

// probability ratio = exp(new_logp - old_logp)
fn ratio_from_logps(new_logp: &Tensor, old_logp: &Tensor) -> Tensor {
    new_logp.sub_tensor(old_logp).exp()
}

// Clamp ratio to [1-clip, 1+clip] using ReLU-based clamp (no custom ops)
fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/RL_training/ppo_continuous.rs (line 239)

fn gaussian_log_prob(action: &Tensor, mean: &Tensor, log_std: &Tensor) -> Tensor {
    // All tensors shaped [B, A] (log_std is broadcastable)
    let std = log_std.exp();
    let var = std.pow_scalar(2.0);
    let log_scale = log_std;
    let diff = action.sub_tensor(mean);
    let log_prob = diff
        .pow_scalar(2.0)
        .div_tensor(&var)
        .add_scalar(std::f32::consts::LN_2 + std::f32::consts::PI)
        .add_tensor(&log_scale.mul_scalar(2.0))
        .mul_scalar(0.5)
        .mul_scalar(-1.0);
    // Sum across action dim (dim=1) -> [B,1]
    log_prob.sum_dims(&[1], true)
}

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

examples/getting_started/tensor_basics.rs (line 100)

fn demonstrate_basic_operations() {
    println!("\n--- Basic Operations ---");

    let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
    let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();

    // Addition
    let sum = a.add_tensor(&b);
    println!("A + B: {:?}", sum.data());

    // Subtraction
    let diff = a.sub_tensor(&b);
    println!("A - B: {:?}", diff.data());

    // Multiplication
    let product = a.mul_tensor(&b);
    println!("A * B: {:?}", product.data());

    // Division
    let quotient = a.div_tensor(&b);
    println!("A / B: {:?}", quotient.data());

    // Scalar operations
    let scalar_add = a.add_scalar(5.0);
    println!("A + 5.0: {:?}", scalar_add.data());

    let scalar_mul = a.mul_scalar(2.0);
    println!("A * 2.0: {:?}", scalar_mul.data());
}

Additional examples can be found in:

• examples/neural_networks/basic_transformer.rs
• examples/neural_networks/feedforward_network.rs
• examples/RL_training/../neural_networks/basic_linear_layer.rs
• examples/RL_training/dqn.rs
• examples/RL_training/td3.rs

pub fn sub_scalar(&self, scalar: f32) -> Tensor

Element-wise subtraction of a scalar from this tensor

Performs element-wise subtraction of a scalar value: output[i] = self[i] - scalar

Arguments

• scalar - The scalar value to subtract from each element

Returns

A new tensor with the scalar subtracted from each element

Performance Characteristics

• SIMD Optimization: AVX2-optimized with 32-element blocks and 4x unrolling
• Scalar Fallback: 4x unrolled scalar implementation for non-SIMD hardware
• Cache-friendly: Linear memory access patterns
• Mathematical Accuracy: High-precision subtraction computation
• GradTrack Support: Full automatic differentiation with efficient gradient computation

Examples

Basic Scalar Subtraction

use train_station::Tensor;

let a = Tensor::from_slice(&[5.0, 7.0, 9.0], vec![3]).unwrap();
let b = a.sub_scalar(2.0);
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 3.0); // 5.0 - 2.0
assert_eq!(b.get(&[1]), 5.0); // 7.0 - 2.0
assert_eq!(b.get(&[2]), 7.0); // 9.0 - 2.0

Negative Scalar Subtraction

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let b = a.sub_scalar(-2.0); // Subtracting negative = adding
assert_eq!(b.shape().dims(), vec![3]);
assert_eq!(b.get(&[0]), 3.0); // 1.0 - (-2.0) = 3.0
assert_eq!(b.get(&[1]), 4.0); // 2.0 - (-2.0) = 4.0
assert_eq!(b.get(&[2]), 5.0); // 3.0 - (-2.0) = 5.0

Examples found in repository

examples/RL_training/dqn.rs (line 325)

fn pseudo_huber_mean(diff: &Tensor) -> Tensor {
    diff.pow_scalar(2.0)
        .add_scalar(1.0)
        .sqrt()
        .sub_scalar(1.0)
        .mean()
}

examples/iterators/advanced_patterns.rs (line 114)

fn demonstrate_data_pipeline() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Data Processing Pipeline ---");

    // Simulate raw sensor data with noise
    let raw_data: Vec<f32> = (0..20)
        .map(|i| {
            let base = i as f32 * 0.5;
            let noise = (i % 3) as f32 * 0.1;
            base + noise
        })
        .collect();

    let tensor = Tensor::from_slice(&raw_data, vec![20])?;
    println!("Raw sensor data: {:?}", tensor.data());

    // Multi-stage processing pipeline
    println!("\nProcessing pipeline:");
    println!("1. Normalize data (z-score)");
    println!("2. Apply smoothing filter");
    println!("3. Detect outliers");
    println!("4. Apply feature scaling");

    // Stage 1: Normalization
    let mean = tensor.mean().value();
    let std = tensor.std().value();
    let normalized: Tensor = tensor
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();
    println!(
        "  Normalized (mean={:.3}, std={:.3}): {:?}",
        mean,
        std,
        normalized.data()
    );

    // Stage 2: Smoothing (simple moving average)
    let smoothed: Tensor = normalized
        .iter()
        .enumerate()
        .map(|(i, elem)| {
            if i == 0 || i == normalized.size() - 1 {
                elem.clone()
            } else {
                // Simple 3-point average
                let prev = normalized.element_view(i - 1);
                let next = normalized.element_view(i + 1);
                elem.add_tensor(&prev).add_tensor(&next).div_scalar(3.0)
            }
        })
        .collect();
    println!("  Smoothed: {:?}", smoothed.data());

    // Stage 3: Outlier detection and removal
    let outlier_threshold = 2.0;
    let cleaned: Tensor = smoothed
        .iter()
        .filter(|elem| elem.value().abs() < outlier_threshold)
        .collect();
    println!(
        "  Outliers removed (threshold={}): {:?}",
        outlier_threshold,
        cleaned.data()
    );

    // Stage 4: Feature scaling to [0, 1] range
    let min_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);
    let scaled: Tensor = cleaned
        .iter()
        .map(|elem| elem.sub_scalar(min_val).div_scalar(max_val - min_val))
        .collect();
    println!("  Scaled to [0,1]: {:?}", scaled.data());

    Ok(())
}

/// Demonstrate conditional processing patterns
///
/// Shows how to implement dynamic filtering and transformation
/// based on data characteristics and conditions.
fn demonstrate_conditional_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Conditional Processing ---");

    // Create data with mixed characteristics
    let data = vec![1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0, 9.0, -10.0];
    let tensor = Tensor::from_slice(&data, vec![10])?;
    println!("Input data: {:?}", tensor.data());

    // Conditional transformation based on sign
    println!("\nConditional transformation (positive/negative handling):");
    let processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val > 0.0 {
                elem.pow_scalar(2.0) // Square positive values
            } else {
                elem.mul_scalar(-1.0).sqrt() // Square root of absolute negative values
            }
        })
        .collect();
    println!("  Processed: {:?}", processed.data());

    // Adaptive filtering based on local statistics
    println!("\nAdaptive filtering (remove values > 2 std from local mean):");
    let window_size = 3;
    let adaptive_filtered: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, elem)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(tensor.size());

            // Calculate local mean and std
            let local_values: Vec<f32> = (start..end)
                .map(|j| tensor.element_view(j).value())
                .collect();

            let local_mean = local_values.iter().sum::<f32>() / local_values.len() as f32;
            let local_variance = local_values
                .iter()
                .map(|v| (v - local_mean).powi(2))
                .sum::<f32>()
                / local_values.len() as f32;
            let local_std = local_variance.sqrt();

            let threshold = local_mean + 2.0 * local_std;
            elem.value() <= threshold
        })
        .map(|(_, elem)| elem)
        .collect();
    println!("  Adaptive filtered: {:?}", adaptive_filtered.data());

    // Multi-condition processing
    println!("\nMulti-condition processing:");
    let multi_processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            match () {
                _ if val > 5.0 => elem.mul_scalar(2.0), // Double large values
                _ if val < -5.0 => elem.div_scalar(2.0), // Halve small values
                _ if val.abs() < 2.0 => elem.add_scalar(1.0), // Add 1 to small values
                _ => elem.clone(),                      // Keep others unchanged
            }
        })
        .collect();
    println!("  Multi-condition: {:?}", multi_processed.data());

    Ok(())
}

/// Demonstrate batch processing operations
///
/// Shows efficient processing of large datasets using iterator
/// patterns and batch operations for performance optimization.
fn demonstrate_batch_operations() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Batch Operations ---");

    // Create a larger dataset for batch processing
    let size = 100;
    let data: Vec<f32> = (0..size)
        .map(|i| {
            let x = i as f32 / size as f32;
            x * x + 0.1 * (i % 7) as f32 // Quadratic with some noise
        })
        .collect();

    let tensor = Tensor::from_slice(&data, vec![size])?;
    println!("Dataset size: {}", tensor.size());

    // Batch processing with windowing (iterator views)
    println!("\nBatch processing with sliding windows:");
    let batch_size = 10;
    let batches: Vec<Tensor> = tensor
        .iter()
        .collect::<Vec<_>>()
        .chunks(batch_size)
        .map(|chunk| {
            // Process each batch independently
            chunk
                .iter()
                .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
                .collect()
        })
        .collect();

    println!(
        "  Processed {} batches of size {}",
        batches.len(),
        batch_size
    );
    for (i, batch) in batches.iter().enumerate() {
        println!(
            "    Batch {}: mean={:.3}, std={:.3}",
            i,
            batch.mean().value(),
            batch.std().value()
        );
    }

    // Parallel-like processing with stride
    println!("\nStrided processing (every nth element):");
    let stride = 5;
    let strided: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % stride == 0)
        .map(|(_, elem)| elem)
        .collect();
    println!("  Strided (every {}th): {:?}", stride, strided.data());

    // Hierarchical processing
    println!("\nHierarchical processing (coarse to fine):");
    let coarse: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 == 0) // Take every 4th element
        .map(|(_, elem)| elem)
        .collect();

    let fine: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 != 0) // Take the rest
        .map(|(_, elem)| elem)
        .collect();

    println!("  Coarse (every 4th): {:?}", coarse.data());
    println!("  Fine (rest): {:?}", fine.data());

    // Combine coarse and fine with different processing
    let combined: Tensor = coarse
        .iter()
        .map(|elem| elem.mul_scalar(2.0)) // Scale coarse
        .chain(fine.iter().map(|elem| elem.div_scalar(2.0))) // Scale fine
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

/// Demonstrate real-world processing scenarios
///
/// Shows practical applications of iterator patterns for
/// common data processing tasks in machine learning and analytics.
fn demonstrate_real_world_scenarios() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Real-world Scenarios ---");

    // Scenario 1: Time series analysis
    println!("\nScenario 1: Time Series Analysis");
    let time_series: Vec<f32> = (0..24)
        .map(|hour| {
            let base = 20.0 + 10.0 * (hour as f32 * std::f32::consts::PI / 12.0).sin();
            base + (hour % 3) as f32 * 2.0 // Add some noise
        })
        .collect();

    let series = Tensor::from_slice(&time_series, vec![24])?;
    println!("  Time series (24 hours): {:?}", series.data());

    // Calculate moving average with view-based iteration
    let window_size = 3;
    let moving_avg: Tensor = series
        .iter()
        .enumerate()
        .map(|(i, _)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(series.size());
            let window = series.iter_range(start, end);
            window.fold(0.0, |acc, elem| acc + elem.value()) / (end - start) as f32
        })
        .map(|val| Tensor::from_slice(&[val], vec![1]).unwrap())
        .collect();
    println!(
        "  Moving average (window={}): {:?}",
        window_size,
        moving_avg.data()
    );

    // Inference pipeline with NoGrad + streaming
    println!("\nInference pipeline (NoGrad + streaming)");
    let features = Tensor::from_slice(
        &(0..48).map(|i| i as f32 * 0.125).collect::<Vec<_>>(),
        vec![6, 8],
    )?;
    let fast = with_no_grad(|| {
        // Stream values directly, apply light affine, and collect back to same shape
        features
            .data()
            .iter()
            .copied()
            .map(|x| 0.75 * x + 0.1)
            .collect_shape(vec![6, 8])
    });
    println!(
        "  NoGrad streamed transform shape: {:?}",
        fast.shape().dims()
    );

    // Row-wise iteration with shape-preserving collection (GradTrack-friendly)
    let per_row: Tensor = features
        .iter()
        .map(|row| row.mul_scalar(0.5).add_scalar(2.0))
        .collect_shape(vec![6, 8]);
    println!("  Row-wise mapped shape: {:?}", per_row.shape().dims());

    // Scenario 2: Feature engineering
    println!("\nScenario 2: Feature Engineering");
    let features = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;
    println!("  Original features: {:?}", features.data());

    // Create polynomial features
    let poly_features: Tensor = features
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // x^1
                elem.pow_scalar(2.0), // x^2
                elem.pow_scalar(3.0), // x^3
            ]
        })
        .collect();
    println!(
        "  Polynomial features (x, x^2, x^3): {:?}",
        poly_features.data()
    );

    // Scenario 3: Data augmentation
    println!("\nScenario 3: Data Augmentation");
    let original = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?;
    println!("  Original data: {:?}", original.data());

    // Augment with noise and scaling
    let augmented: Tensor = original
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // Original
                elem.add_scalar(0.1), // Add noise
                elem.sub_scalar(0.1), // Subtract noise
                elem.mul_scalar(1.1), // Scale up
                elem.mul_scalar(0.9), // Scale down
            ]
        })
        .collect();
    println!("  Augmented data: {:?}", augmented.data());

    // Scenario 4: Statistical analysis
    println!("\nScenario 4: Statistical Analysis");
    let sample_data = Tensor::from_slice(&[1.1, 2.3, 1.8, 2.1, 1.9, 2.0, 1.7, 2.2], vec![8])?;
    println!("  Sample data: {:?}", sample_data.data());

    // Calculate various statistics
    let mean = sample_data.mean().value();
    let std = sample_data.std().value();
    let min = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);

    // Z-score normalization
    let z_scores: Tensor = sample_data
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();

    println!(
        "  Statistics: mean={:.3}, std={:.3}, min={:.3}, max={:.3}",
        mean, std, min, max
    );
    println!("  Z-scores: {:?}", z_scores.data());

    Ok(())
}

impl Tensor

pub fn tanh(&self) -> Tensor

Element-wise hyperbolic tangent activation

Computes hyperbolic tangent for each element: output[i] = tanh(self[i])

The hyperbolic tangent function maps any real number to the range (-1, 1), making it useful as an activation function in neural networks.

Returns

A new tensor with tanh applied to each element, values in range (-1, 1)

Performance Characteristics

• High Precision: Accurate scalar implementation for mathematical validation
• 4x Unrolling: Optimized scalar operations with instruction-level parallelism
• Cache-friendly: Linear memory access patterns
• Numerical Stability: Robust handling of extreme input values
• GradTrack Support: Full automatic differentiation with efficient gradient computation

Mathematical Properties

• Range: Output values are in the range (-1, 1)
• Symmetry: tanh(-x) = -tanh(x) (odd function)
• Zero: tanh(0) = 0
• Gradient: ∂tanh(x)/∂x = 1 - tanh²(x) = sech²(x)

Examples

Basic Hyperbolic Tangent

use train_station::Tensor;

let a = Tensor::from_slice(&[-1.0, 0.0, 1.0], vec![3]).unwrap();
let b = a.tanh();
assert_eq!(b.shape().dims(), vec![3]);
assert!((b.get(&[0]) - (-0.7615942)).abs() < 1e-6); // tanh(-1.0)
assert!((b.get(&[1]) - 0.0).abs() < 1e-6); // tanh(0.0)
assert!((b.get(&[2]) - 0.7615942).abs() < 1e-6); // tanh(1.0)

Extreme Values

use train_station::Tensor;

let a = Tensor::from_slice(&[-10.0, 10.0], vec![2]).unwrap();
let b = a.tanh();
assert_eq!(b.shape().dims(), vec![2]);
assert!((b.get(&[0]) - (-1.0)).abs() < 1e-6); // tanh(-10.0) ≈ -1
assert!((b.get(&[1]) - 1.0).abs() < 1e-6); // tanh(10.0) ≈ 1

Examples found in repository

examples/RL_training/td3.rs (line 55)

fn tanh_bounded(x: &Tensor) -> Tensor {
    x.tanh()
}

impl Tensor

pub fn argmax(&self) -> Tensor

Returns the index of the maximum value across all elements in the tensor

This operation finds the flat index (0-based) of the element with the highest value. If multiple elements have the same maximum value, the index of the first occurrence is returned. The output is a scalar tensor with shape [1] containing the index as a float.

This operation is non-differentiable and the output never requires gradients.

Returns

A tensor with shape [1] containing the flat index of the maximum value

Examples

use train_station::Tensor;

// 1D tensor
let tensor = Tensor::from_slice(&[1.0, 5.0, 3.0, 2.0], vec![4]).unwrap();
let max_idx = tensor.argmax();
assert_eq!(max_idx.shape().dims(), vec![1]);
assert_eq!(max_idx.get(&[0]), 1.0); // Index 1 has value 5.0

use train_station::Tensor;

// 2D tensor
let tensor = Tensor::from_slice(&[1.0, 3.0, 2.0, 4.0, 0.0, 5.0], vec![2, 3]).unwrap();
let max_idx = tensor.argmax();
assert_eq!(max_idx.get(&[0]), 5.0); // Flat index 5 has value 5.0

use train_station::Tensor;

// Tied values return first occurrence
let tensor = Tensor::from_slice(&[3.0, 5.0, 5.0, 2.0], vec![4]).unwrap();
let max_idx = tensor.argmax();
assert_eq!(max_idx.get(&[0]), 1.0); // First occurrence of 5.0 at index 1

pub fn argmax_dim(&self, dim: usize, keepdim: bool) -> Tensor

Returns the indices of maximum values along a specified dimension

This operation finds the indices of maximum values along the specified dimension. For each slice along the dimension, it returns the index of the maximum value. If multiple elements have the same maximum value, the index of the first occurrence is returned.

The output shape depends on the keepdim parameter:

• If keepdim is true, the reduced dimension is kept with size 1
• If keepdim is false, the reduced dimension is removed

This operation is non-differentiable and the output never requires gradients.

Arguments

• dim - The dimension along which to find argmax indices (0-based)
• keepdim - Whether to keep the reduced dimension with size 1

Returns

A tensor containing the indices of maximum values along the specified dimension

Panics

Panics if dim is out of bounds for the tensor’s rank or if the dimension size is 0.

Examples

use train_station::Tensor;

// 2D tensor: [[1.0, 3.0, 2.0],
//             [4.0, 0.0, 5.0]]
let tensor = Tensor::from_slice(&[1.0, 3.0, 2.0, 4.0, 0.0, 5.0], vec![2, 3]).unwrap();

// argmax along columns (dim=1)
let col_max_idx = tensor.argmax_dim(1, false);
assert_eq!(col_max_idx.shape().dims(), vec![2]);
assert_eq!(col_max_idx.get(&[0]), 1.0); // Row 0: max at index 1 (value 3.0)
assert_eq!(col_max_idx.get(&[1]), 2.0); // Row 1: max at index 2 (value 5.0)

use train_station::Tensor;

// argmax along rows (dim=0) with keepdim
let tensor = Tensor::from_slice(&[1.0, 3.0, 2.0, 4.0, 0.0, 5.0], vec![2, 3]).unwrap();
let row_max_idx = tensor.argmax_dim(0, true);
assert_eq!(row_max_idx.shape().dims(), vec![1, 3]);
assert_eq!(row_max_idx.get(&[0, 0]), 1.0); // Col 0: max at index 1 (value 4.0)
assert_eq!(row_max_idx.get(&[0, 1]), 0.0); // Col 1: max at index 0 (value 3.0)
assert_eq!(row_max_idx.get(&[0, 2]), 1.0); // Col 2: max at index 1 (value 5.0)

use train_station::Tensor;

// 1D tensor edge case
let tensor = Tensor::from_slice(&[5.0, 1.0, 8.0, 3.0], vec![4]).unwrap();
let max_idx = tensor.argmax_dim(0, false);
assert_eq!(max_idx.shape().dims(), vec![1]); // Special case: becomes [1] not []
assert_eq!(max_idx.get(&[0]), 2.0); // Index 2 has maximum value 8.0

impl Tensor

pub fn argmin(&self) -> Tensor

Returns the index of the minimum value in the tensor

This method finds the flat index of the minimum value across all elements in the tensor. The result is a scalar tensor containing the index as a floating-point value. This operation is non-differentiable and the output never requires gradient tracking.

Returns

A tensor with shape [1] containing the flat index of the minimum value as a f32. If the input tensor is empty, returns 0.0.

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[3.0, -2.0, 5.0, -1.0], vec![4]).unwrap();
let min_index = tensor.argmin();
assert_eq!(min_index.get(&[0]), 1.0); // -2.0 is at index 1

use train_station::Tensor;

// Empty tensor case
let empty_tensor = Tensor::new(vec![0]);
let min_index = empty_tensor.argmin();
assert_eq!(min_index.get(&[0]), 0.0);

pub fn argmin_dim(&self, dim: usize, keepdim: bool) -> Tensor

Returns the indices of minimum values along a specified dimension

This method finds the indices of minimum values along the specified dimension. The result contains the indices where the minimum values occur in that dimension. This operation is non-differentiable and the output never requires gradient tracking.

Arguments

• dim - The dimension along which to find minimum indices (0-based)
• keepdim - Whether to keep the reduced dimension in the output shape
  • If true, the reduced dimension is kept with size 1
  • If false, the reduced dimension is removed from the output shape

Returns

A tensor containing the indices of minimum values along the specified dimension. The output shape depends on keepdim:

• If keepdim is true, the reduced dimension has size 1
• If keepdim is false, the reduced dimension is removed

Panics

• If dim is out of bounds for the tensor’s rank
• If the dimension to reduce has size 0

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[3.0, -2.0, 5.0, -1.0, 0.0, -3.0], vec![2, 3]).unwrap();

// Find minimum indices along dimension 1 (columns), keeping the dimension
let indices = tensor.argmin_dim(1, true);
assert_eq!(indices.shape().dims(), vec![2, 1]);
assert_eq!(indices.get(&[0, 0]), 1.0); // -2.0 is at index 1 in first row
assert_eq!(indices.get(&[1, 0]), 2.0); // -3.0 is at index 2 in second row

use train_station::Tensor;

let tensor = Tensor::from_slice(&[3.0, -2.0, 5.0, -1.0, 0.0, -3.0], vec![2, 3]).unwrap();

// Find minimum indices along dimension 1 (columns), removing the dimension
let indices = tensor.argmin_dim(1, false);
assert_eq!(indices.shape().dims(), vec![2]);
assert_eq!(indices.get(&[0]), 1.0); // -2.0 is at index 1 in first row
assert_eq!(indices.get(&[1]), 2.0); // -3.0 is at index 2 in second row

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();

// Find minimum index in a 1D tensor
let index = tensor.argmin_dim(0, false);
assert_eq!(index.shape().dims(), vec![1]);
assert_eq!(index.get(&[0]), 0.0); // 1.0 is at index 0

impl Tensor

pub fn max(&self) -> Tensor

Computes the maximum value over all elements in the tensor

Returns a scalar tensor containing the maximum value. For empty tensors, returns negative infinity. This operation supports gradient tracking through the GradTrack system.

Returns

A tensor with shape [1] containing the maximum value

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 5.0, 3.0, 2.0], vec![2, 2]).unwrap();
let max_val = tensor.max();
assert_eq!(max_val.get(&[0]), 5.0);

GradTrack Support

When requires_grad is true, this operation is tracked for automatic differentiation. The gradient computation uses the saved input and output for efficient backward pass.

pub fn max_dims(&self, dims: &[usize], keepdim: bool) -> Tensor

Computes the maximum value over specified dimensions

Reduces the tensor along the specified dimensions by computing the maximum value in each reduction group. The keepdim parameter determines whether reduced dimensions are kept with size 1 or removed entirely.

Arguments

• dims - Dimensions to reduce over (must be valid for the tensor’s rank)
• keepdim - If true, reduced dimensions are kept with size 1; if false, they are removed

Returns

A tensor with the specified dimensions reduced

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();

// Max over columns (dim 1), keeping dimensions
let max_cols = tensor.max_dims(&[1], true);
assert_eq!(max_cols.shape().dims(), vec![2, 1]);
assert_eq!(max_cols.get(&[0, 0]), 3.0);
assert_eq!(max_cols.get(&[1, 0]), 6.0);

// Max over rows (dim 0), removing dimensions
let max_rows = tensor.max_dims(&[0], false);
assert_eq!(max_rows.shape().dims(), vec![3]);
assert_eq!(max_rows.get(&[0]), 4.0);
assert_eq!(max_rows.get(&[1]), 5.0);
assert_eq!(max_rows.get(&[2]), 6.0);

Panics

Panics if:

• dims is empty
• Any dimension in dims is out of bounds for the tensor’s rank

GradTrack Support

When requires_grad is true, this operation is tracked for automatic differentiation. The gradient computation preserves the original input shape and handles broadcasting correctly.

Examples found in repository

examples/supervised_training/supervised_classification.rs (line 51)

fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

examples/RL_training/ppo_discrete.rs (line 279)

fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

impl Tensor

pub fn mean(&self) -> Tensor

Computes the arithmetic mean of all elements in the tensor

This method calculates the average value across all tensor elements by summing all values and dividing by the total number of elements. The result is a scalar tensor containing the mean value. This operation supports gradient tracking through the GradTrack system.

Returns

A tensor with shape [1] containing the arithmetic mean of all elements. For empty tensors, returns 0.0 as a safe default.

Performance Characteristics

• Linear Time: O(n) complexity for computing the sum
• Memory Efficient: Single pass through tensor data with SIMD-optimized accumulation
• Numerical Stability: Uses direct accumulation for typical tensor sizes
• Edge Case Handling: Returns 0.0 for empty tensors

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let mean_val = tensor.mean();
assert_eq!(mean_val.get(&[0]), 2.5); // (1+2+3+4)/4 = 2.5

use train_station::Tensor;

// Empty tensor case
let empty_tensor = Tensor::new(vec![0]);
let mean_val = empty_tensor.mean();
assert_eq!(mean_val.get(&[0]), 0.0);

GradTrack Support

When requires_grad is true, this operation is tracked for automatic differentiation. The gradient computation distributes the gradient equally across all input elements.

Examples found in repository

examples/supervised_training/supervised_regression.rs (line 43)

fn mse(pred: &Tensor, target: &Tensor) -> Tensor {
    pred.sub_tensor(target).pow_scalar(2.0).mean()
}

fn rmse(pred: &Tensor, target: &Tensor) -> f32 {
    mse(pred, target).sqrt().value()
}

fn r2_score(pred: &Tensor, target: &Tensor) -> f32 {
    // R^2 = 1 - SS_res / SS_tot
    let y = target;
    let y_mean = y.mean();
    let ss_res = pred.sub_tensor(y).pow_scalar(2.0).sum();
    let ss_tot = y.sub_tensor(&y_mean).pow_scalar(2.0).sum();
    let ss_res_v = ss_res.value();
    let ss_tot_v = ss_tot.value().max(1e-12); // avoid divide by zero
    1.0 - (ss_res_v / ss_tot_v)
}

examples/RL_training/dqn.rs (line 326)

fn pseudo_huber_mean(diff: &Tensor) -> Tensor {
    diff.pow_scalar(2.0)
        .add_scalar(1.0)
        .sqrt()
        .sub_scalar(1.0)
        .mean()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== DQN Example (YardEnv discrete) ===");

    // Dims
    let state_dim = 3usize;
    let action_dim = 3usize;

    // Hparams
    let gamma = 0.99f32;
    let batch_size = 64usize;
    let start_steps = 200usize;
    let target_update_interval = 200usize; // hard update cadence
    let max_grad_norm = 1.0f32;
    let mut epsilon = 1.0f32;
    let eps_min = 0.05f32;
    let eps_decay_steps = 2_000usize; // linear decay
    let total_steps = std::env::var("DQN_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3000usize);

    // Models
    let mut q_net = QNet::new(state_dim, action_dim, Some(7));
    let mut q_targ = QNet::new(state_dim, action_dim, Some(8));
    q_targ.net.copy_from(&q_net.net);
    q_targ.set_requires_grad_all(false);

    // Optimizer
    let mut q_opt = Adam::with_learning_rate(3e-4);
    for p in q_net.parameters() {
        q_opt.add_parameter(p);
    }

    // Replay + env
    let mut rb = ReplayBuffer::new(100_000, state_dim);
    let mut env = YardEnv::new(12345);
    let mut rng = SmallRng::new(999_111);

    // Metrics
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    for t in 0..total_steps {
        // Epsilon-greedy action
        let action_index = if t < start_steps || rng.next_f32() < epsilon {
            rng.sample_index(action_dim)
        } else {
            let _ng = NoGradTrack::new();
            let q_vals = q_net.forward(&state);
            let row = q_vals.data();
            let mut best_i = 0usize;
            let mut best_v = row[0];
            for (i, &r) in row.iter().enumerate().take(action_dim).skip(1) {
                if r > best_v {
                    best_v = r;
                    best_i = i;
                }
            }
            best_i
        };

        // Env step
        let (next_state, reward, done) = env.step(action_index);
        episode_return += reward;

        // Store
        let s_slice = state.data().to_vec();
        let s2_slice = next_state.data().to_vec();
        rb.push(
            &s_slice,
            action_index,
            reward,
            if done { 1.0 } else { 0.0 },
            &s2_slice,
        );

        // Reset on done
        state = if done {
            let st = env.reset();
            ema_return = Some(match ema_return {
                None => episode_return,
                Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
            });
            if episode_return > best_return {
                best_return = episode_return;
            }
            println!(
                "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3} | rb_size={}",
                t,
                episode,
                episode_return,
                ema_return.unwrap_or(episode_return),
                best_return,
                rb.size
            );
            episode_return = 0.0;
            episode += 1;
            st
        } else {
            next_state
        };

        // Epsilon linear decay
        if t < eps_decay_steps {
            epsilon = (1.0 - (t as f32) / (eps_decay_steps as f32)) * (1.0 - eps_min) + eps_min;
        }

        // Train
        if rb.can_sample(batch_size) {
            let (s, a_idx, r, d, s2) = rb.sample(batch_size, &mut rng);

            // Double DQN target: a* = argmax_a Q_online(s2,a); y = r + (1-d)*gamma*Q_target(s2, a*)
            let target_q = {
                let _ng = NoGradTrack::new();
                let q_online_s2 = q_net.forward(&s2);
                // argmax per row (manual on CPU)
                let row_stride = action_dim;
                let qd = q_online_s2.data();
                let mut next_actions: Vec<usize> = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let base = i * row_stride;
                    let mut bi = 0usize;
                    let mut bv = qd[base];
                    for j in 1..action_dim {
                        let v = qd[base + j];
                        if v > bv {
                            bv = v;
                            bi = j;
                        }
                    }
                    next_actions.push(bi);
                }
                let q_targ_s2 = q_targ.forward(&s2);
                let q_targ_g = q_targ_s2.gather(1, &next_actions, &[batch_size, 1]);
                let not_done = Tensor::ones(vec![batch_size, 1]).sub_tensor(&d);
                r.add_tensor(&not_done.mul_scalar(gamma).mul_tensor(&q_targ_g))
            };

            // Q(s,a) for current actions
            // Zero grads first
            {
                let mut params = q_net.parameters();
                q_opt.zero_grad(&mut params);
            }

            let q_all = q_net.forward(&s);
            let q_sa = q_all.gather(1, &a_idx, &[batch_size, 1]);
            let diff = q_sa.sub_tensor(&target_q);
            let mut loss = pseudo_huber_mean(&diff);
            loss.backward(None);

            // Step (filter only params with grads)
            {
                let params = q_net.parameters();
                let mut with_grads: Vec<&mut Tensor> = Vec::new();
                for p in params {
                    if p.grad_owned().is_some() {
                        with_grads.push(p);
                    }
                }
                if !with_grads.is_empty() {
                    let gn = grad_global_norm(&mut with_grads);
                    clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                    q_opt.step(&mut with_grads);
                    q_opt.zero_grad(&mut with_grads);
                    if t % 100 == 0 {
                        let mut pn = q_net.parameters();
                        let pn_l2 = params_l2_norm(&mut pn);
                        let q_mean = q_all.mean().value();
                        println!(
                            "t={:5} | loss={:.4} | q_mean={:.3} | grad_norm={:.3} | param_norm={:.3} | eps={:.3}",
                            t, loss.value(), q_mean, gn, pn_l2, epsilon
                        );
                    }
                }
            }

            // Target hard update
            if t % target_update_interval == 0 {
                q_targ.net.copy_from(&q_net.net);
            }

            // Clear graphs
            clear_all_graphs_known();
        }
    }

    println!("=== DQN training finished ===");
    Ok(())
}

examples/supervised_training/supervised_bce.rs (line 65)

fn bce_with_logits(logits: &Tensor, targets: &Tensor) -> Tensor {
    let relu_z = logits.relu();
    let zy = logits.mul_tensor(targets);
    // |z| = relu(z) + relu(-z)
    let abs_z = relu_z.add_tensor(&logits.mul_scalar(-1.0).relu());
    let log_term = abs_z.mul_scalar(-1.0).exp().add_scalar(1.0).log();
    relu_z.sub_tensor(&zy).add_tensor(&log_term).mean()
}

examples/supervised_training/supervised_classification.rs (line 58)

fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

examples/neural_networks/basic_encoder.rs (line 94)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Basic Encoder Example ===");

    let batch = 2usize;
    let seq = 6usize;
    let embed = 32usize;
    let heads = 4usize;

    let input = Tensor::randn(vec![batch, seq, embed], Some(11));
    let mut enc = EncoderBlock::new(embed, heads, Some(123));

    // Example: no mask (set Some(mask) to use masking)
    let out = enc.forward(&input, None);
    println!("Output shape: {:?}", out.shape().dims());

    // Verify gradients/optimization
    let mut opt = Adam::with_learning_rate(0.01);
    let mut params = enc.parameters();
    for p in &params {
        opt.add_parameter(p);
    }
    let mut loss = out.mean();
    loss.backward(None);
    opt.step(&mut params);
    opt.zero_grad(&mut params);
    println!("Loss: {:.6}", loss.value());
    println!("=== Done ===");
    Ok(())
}

examples/getting_started/tensor_basics.rs (line 194)

fn demonstrate_utility_functions() {
    println!("\n--- Utility Functions ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

    // Basic properties
    println!("Shape: {:?}", tensor.shape().dims());
    println!("Size: {}", tensor.size());
    println!("Is contiguous: {}", tensor.is_contiguous());
    println!("Device: {:?}", tensor.device());

    // Mathematical operations
    let sum = tensor.sum();
    println!("Sum: {}", sum.value());

    let mean = tensor.mean();
    println!("Mean: {}", mean.value());

    let norm = tensor.norm();
    println!("Norm: {}", norm.value());

    // Device placement
    let cpu_tensor = Tensor::zeros_on_device(vec![3, 3], train_station::Device::cpu());
    println!(
        "CPU tensor: shape {:?}, device: {:?}",
        cpu_tensor.shape().dims(),
        cpu_tensor.device()
    );
}

Additional examples can be found in:

• examples/neural_networks/basic_decoder.rs
• examples/neural_networks/basic_transformer.rs
• examples/neural_networks/multi_head_attention.rs
• examples/optimizers/adam_configurations.rs
• examples/optimizers/learning_rate_scheduling.rs
• examples/getting_started/optimizer_basics.rs
• examples/neural_networks/feedforward_network.rs
• examples/iterators/advanced_patterns.rs
• examples/RL_training/../neural_networks/basic_linear_layer.rs
• examples/RL_training/ppo_discrete.rs
• examples/RL_training/ppo_continuous.rs
• examples/RL_training/td3.rs

pub fn mean_dims(&self, dims: &[usize], keepdim: bool) -> Tensor

Computes the arithmetic mean over specified dimensions

This method calculates the mean value along the specified dimensions by first computing the sum over those dimensions and then dividing by the product of the reduced dimension sizes. The keepdim parameter determines whether reduced dimensions are kept with size 1 or removed entirely.

Arguments

• dims - Dimensions to reduce over (must be valid for the tensor’s rank)
• keepdim - If true, reduced dimensions are kept with size 1; if false, they are removed

Returns

A tensor with the specified dimensions reduced by computing the mean. The output shape depends on keepdim:

• If keepdim is true, reduced dimensions have size 1
• If keepdim is false, reduced dimensions are removed

Performance Characteristics

• Efficient Implementation: Uses sum_dims followed by scalar multiplication
• Memory Optimized: Leverages existing sum reduction for optimal performance
• Shape Computation: Fast output shape calculation with dimension preservation
• Numerical Stability: Maintains precision through direct computation

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();

// Mean over columns (dim 1), keeping dimensions
let mean_cols = tensor.mean_dims(&[1], true);
assert_eq!(mean_cols.shape().dims(), vec![2, 1]);
assert_eq!(mean_cols.get(&[0, 0]), 2.0); // (1+2+3)/3 = 2.0
assert_eq!(mean_cols.get(&[1, 0]), 5.0); // (4+5+6)/3 = 5.0

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();

// Mean over rows (dim 0), removing dimensions
let mean_rows = tensor.mean_dims(&[0], false);
assert_eq!(mean_rows.shape().dims(), vec![3]);
assert_eq!(mean_rows.get(&[0]), 2.5); // (1+4)/2 = 2.5
assert_eq!(mean_rows.get(&[1]), 3.5); // (2+5)/2 = 3.5
assert_eq!(mean_rows.get(&[2]), 4.5); // (3+6)/2 = 4.5

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

// Mean over multiple dimensions
let mean_all = tensor.mean_dims(&[0, 1], false);
assert_eq!(mean_all.shape().dims(), vec![1]);
assert_eq!(mean_all.get(&[0]), 2.5); // (1+2+3+4)/4 = 2.5

Panics

Panics if:

• dims is empty
• Any dimension in dims is out of bounds for the tensor’s rank

GradTrack Support

When requires_grad is true, this operation is tracked for automatic differentiation. The gradient computation preserves the original input shape and handles broadcasting correctly through the ReduceMeanDims gradient function.

impl Tensor

pub fn min(&self) -> Tensor

Computes the minimum value over all elements in the tensor

Returns a scalar tensor containing the minimum value. For empty tensors, returns positive infinity. This operation supports gradient tracking through the GradTrack system.

Returns

A tensor with shape [1] containing the minimum value

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 5.0, 3.0, 2.0], vec![2, 2]).unwrap();
let min_val = tensor.min();
assert_eq!(min_val.get(&[0]), 1.0);

use train_station::Tensor;

// Empty tensor case
let empty_tensor = Tensor::new(vec![0]);
let min_val = empty_tensor.min();
assert_eq!(min_val.get(&[0]), f32::INFINITY);

GradTrack Support

When requires_grad is true, this operation is tracked for automatic differentiation. The gradient computation uses the saved input and output for efficient backward pass.

pub fn min_dims(&self, dims: &[usize], keepdim: bool) -> Tensor

Computes the minimum value over specified dimensions

Reduces the tensor along the specified dimensions by computing the minimum value in each reduction group. The keepdim parameter determines whether reduced dimensions are kept with size 1 or removed entirely.

Arguments

• dims - Dimensions to reduce over (must be valid for the tensor’s rank)
• keepdim - If true, reduced dimensions are kept with size 1; if false, they are removed

Returns

A tensor with the specified dimensions reduced

Examples

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();

// Min over columns (dim 1), keeping dimensions
let min_cols = tensor.min_dims(&[1], true);
assert_eq!(min_cols.shape().dims(), vec![2, 1]);
assert_eq!(min_cols.get(&[0, 0]), 1.0);
assert_eq!(min_cols.get(&[1, 0]), 4.0);

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();

// Min over rows (dim 0), removing dimensions
let min_rows = tensor.min_dims(&[0], false);
assert_eq!(min_rows.shape().dims(), vec![3]);
assert_eq!(min_rows.get(&[0]), 1.0);
assert_eq!(min_rows.get(&[1]), 2.0);
assert_eq!(min_rows.get(&[2]), 3.0);

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

// Min over multiple dimensions
let min_all = tensor.min_dims(&[0, 1], false);
assert_eq!(min_all.shape().dims(), vec![1]);
assert_eq!(min_all.get(&[0]), 1.0);

Panics

Panics if:

• dims is empty
• Any dimension in dims is out of bounds for the tensor’s rank

GradTrack Support

When requires_grad is true, this operation is tracked for automatic differentiation. The gradient computation preserves the original input shape and handles broadcasting correctly.

impl Tensor

pub fn norm(&self) -> Tensor

Computes the L2 norm (Euclidean norm) over all elements

The L2 norm is calculated as sqrt(sum(x²)) where x represents each element in the tensor. This operation reduces the tensor to a scalar value [1].

Returns

A scalar tensor containing the L2 norm value

Examples

use train_station::Tensor;

// Basic L2 norm calculation
let tensor = Tensor::from_slice(&[3.0, 4.0], vec![2]).unwrap();
let norm = tensor.norm();
assert!((norm.get(&[0]) - 5.0).abs() < 1e-6); // sqrt(3² + 4²) = 5

use train_station::Tensor;

// L2 norm of a larger tensor
let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
let tensor = Tensor::from_slice(&data, vec![2, 2, 2]).unwrap();
let norm = tensor.norm();
// sqrt(1² + 2² + 3² + 4² + 5² + 6² + 7² + 8²) = sqrt(204) ≈ 14.283
let expected = 204.0_f32.sqrt();
assert!((norm.get(&[0]) - expected).abs() < 1e-5);

Performance

Uses optimized contiguous tensor path with 4x loop unrolling for better performance. Non-contiguous tensors use stride-aware iteration.

Examples found in repository

examples/getting_started/tensor_basics.rs (line 197)

fn demonstrate_utility_functions() {
    println!("\n--- Utility Functions ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

    // Basic properties
    println!("Shape: {:?}", tensor.shape().dims());
    println!("Size: {}", tensor.size());
    println!("Is contiguous: {}", tensor.is_contiguous());
    println!("Device: {:?}", tensor.device());

    // Mathematical operations
    let sum = tensor.sum();
    println!("Sum: {}", sum.value());

    let mean = tensor.mean();
    println!("Mean: {}", mean.value());

    let norm = tensor.norm();
    println!("Norm: {}", norm.value());

    // Device placement
    let cpu_tensor = Tensor::zeros_on_device(vec![3, 3], train_station::Device::cpu());
    println!(
        "CPU tensor: shape {:?}, device: {:?}",
        cpu_tensor.shape().dims(),
        cpu_tensor.device()
    );
}

examples/optimizers/adam_configurations.rs (line 370)

fn train_with_config(config: TrainingConfig) -> Result<TrainingStats, Box<dyn std::error::Error>> {
    // Create training data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0, 11.0], vec![5, 1]).unwrap();

    // Create model parameters
    let mut weight = Tensor::randn(vec![1, 1], Some(123)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer with custom configuration
    let adam_config = AdamConfig {
        learning_rate: config.learning_rate,
        beta1: config.beta1,
        beta2: config.beta2,
        eps: 1e-8,
        weight_decay: config.weight_decay,
        amsgrad: false,
    };

    let mut optimizer = Adam::with_config(adam_config);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Training loop
    let mut losses = Vec::new();
    let mut convergence_epoch = config.epochs;

    for epoch in 0..config.epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        let loss_value = loss.value();
        losses.push(loss_value);

        // Check for convergence (loss < 0.01)
        if loss_value < 0.01 && convergence_epoch == config.epochs {
            convergence_epoch = epoch;
        }
    }

    Ok(TrainingStats {
        config,
        final_loss: losses[losses.len() - 1],
        loss_history: losses,
        convergence_epoch,
        weight_norm: weight.norm().value(),
    })
}

examples/getting_started/optimizer_basics.rs (line 214)

fn demonstrate_advanced_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Advanced Training Patterns ---");

    // Create a more complex model
    let mut weight = Tensor::randn(vec![1, 2], Some(44)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![2]).with_requires_grad();

    // Create optimizer with different learning rate
    let mut optimizer = Adam::with_learning_rate(0.005);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Create training data: y = 2*x + [1, 3]
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5, 1]).unwrap();
    let y_true = Tensor::from_slice(
        &[3.0, 5.0, 7.0, 9.0, 11.0, 6.0, 8.0, 10.0, 12.0, 14.0],
        vec![5, 2],
    )
    .unwrap();

    println!("Advanced training with monitoring:");
    println!("  Initial learning rate: {}", optimizer.learning_rate());

    // Training loop with monitoring
    let num_epochs = 50;
    let mut losses = Vec::new();
    let mut weight_norms = Vec::new();
    let mut gradient_norms = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Compute gradient norm before optimizer step
        let gradient_norm = weight.grad_owned().unwrap().norm();

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Learning rate scheduling: reduce every 10 epochs
        if epoch > 0 && epoch % 10 == 0 {
            let current_lr = optimizer.learning_rate();
            let new_lr = current_lr * 0.5;
            optimizer.set_learning_rate(new_lr);
            println!(
                "Epoch {:2}: Reduced learning rate from {:.3} to {:.3}",
                epoch, current_lr, new_lr
            );
        }

        // Record metrics
        losses.push(loss.value());
        weight_norms.push(weight.norm().value());
        gradient_norms.push(gradient_norm.value());

        // Print detailed progress
        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, Weight Norm = {:.6}, Gradient Norm = {:.6}",
                epoch,
                loss.value(),
                weight.norm().value(),
                gradient_norm.value()
            );
        }
    }

    println!("Final learning rate: {}", optimizer.learning_rate());

    // Analyze training progression
    let initial_loss = losses[0];
    let final_loss = losses[losses.len() - 1];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final weight norm: {:.6}", weight.norm().value());
    println!("  Final bias: {:?}", bias.data());

    Ok(())
}

/// Demonstrate learning rate scheduling
fn demonstrate_learning_rate_scheduling() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Learning Rate Scheduling ---");

    // Create simple model
    let mut weight = Tensor::randn(vec![1, 1], Some(45)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer with high initial learning rate
    let mut optimizer = Adam::with_learning_rate(0.1);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Simple data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3, 1]).unwrap();
    let y_true = Tensor::from_slice(&[2.0, 4.0, 6.0], vec![3, 1]).unwrap();

    println!("Initial learning rate: {}", optimizer.learning_rate());

    // Training loop with learning rate scheduling
    let num_epochs = 50;
    let mut losses = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Learning rate scheduling: reduce every 10 epochs
        if epoch > 0 && epoch % 10 == 0 {
            let current_lr = optimizer.learning_rate();
            let new_lr = current_lr * 0.5;
            optimizer.set_learning_rate(new_lr);
            println!(
                "Epoch {:2}: Reduced learning rate from {:.3} to {:.3}",
                epoch, current_lr, new_lr
            );
        }

        losses.push(loss.value());

        // Print progress
        if epoch % 10 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, LR = {:.3}",
                epoch,
                loss.value(),
                optimizer.learning_rate()
            );
        }
    }

    println!("Final learning rate: {}", optimizer.learning_rate());

    Ok(())
}

/// Demonstrate training monitoring and analysis
fn demonstrate_training_monitoring() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Monitoring ---");

    // Create model
    let mut weight = Tensor::randn(vec![1, 1], Some(46)).with_requires_grad();
    let mut bias = Tensor::zeros(vec![1]).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    // Training data
    let x_data = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
    let y_true = Tensor::from_slice(&[3.0, 5.0, 7.0, 9.0], vec![4, 1]).unwrap();

    // Training loop with comprehensive monitoring
    let num_epochs = 30;
    let mut losses = Vec::new();
    let mut weight_history = Vec::new();
    let mut bias_history = Vec::new();

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = x_data.matmul(&weight) + &bias;
        let mut loss = (&y_pred - &y_true).pow_scalar(2.0).mean();

        // Backward pass
        loss.backward(None);

        // Optimizer step
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);

        // Record history
        losses.push(loss.value());
        weight_history.push(weight.value());
        bias_history.push(bias.value());

        // Print detailed monitoring
        if epoch % 5 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:2}: Loss = {:.6}, Weight = {:.6}, Bias = {:.6}",
                epoch,
                loss.value(),
                weight.value(),
                bias.value()
            );
        }
    }

    // Analyze training progression
    println!("\nTraining Analysis:");
    println!("  Initial loss: {:.6}", losses[0]);
    println!("  Final loss: {:.6}", losses[losses.len() - 1]);
    println!(
        "  Loss reduction: {:.1}%",
        (losses[0] - losses[losses.len() - 1]) / losses[0] * 100.0
    );

    // Compute statistics
    let loss_mean = compute_mean(&losses);
    let loss_std = compute_std(&losses);
    let weight_change = (weight_history[weight_history.len() - 1] - weight_history[0]).abs();
    let bias_change = (bias_history[bias_history.len() - 1] - bias_history[0]).abs();

    println!("  Average loss: {:.6} ± {:.6}", loss_mean, loss_std);
    println!("  Weight change: {:.6}", weight_change);
    println!("  Bias change: {:.6}", bias_change);
    println!("  Final weight norm: {:.6}", weight.norm().value());
    println!("  Final bias: {:.6}", bias.value());

    Ok(())
}

pub fn norm_dims(&self, dims: &[usize], keepdim: bool) -> Tensor

Computes the L2 norm over specified dimensions

Reduces the tensor along the specified dimensions by computing the L2 norm of each slice. The result maintains the original tensor structure with reduced dimensions optionally preserved as size-1 dimensions.

Arguments

• dims - Vector of dimension indices to reduce over (must be valid for tensor rank)
• keepdim - Whether to keep reduced dimensions as size-1 dimensions

Returns

A tensor with L2 norm computed over the specified dimensions

Examples

use train_station::Tensor;

// Norm along rows (dimension 1) with keepdim=true
let matrix = Tensor::from_slice(&[3.0, 4.0, 0.0, 5.0], vec![2, 2]).unwrap();
let row_norms = matrix.norm_dims(&[1], true);
assert_eq!(row_norms.shape().dims(), vec![2, 1]);
assert!((row_norms.get(&[0, 0]) - 5.0).abs() < 1e-6); // sqrt(3² + 4²)
assert!((row_norms.get(&[1, 0]) - 5.0).abs() < 1e-6); // sqrt(0² + 5²)

use train_station::Tensor;

// Norm along columns (dimension 0) with keepdim=false
let matrix = Tensor::from_slice(&[3.0, 4.0, 0.0, 5.0], vec![2, 2]).unwrap();
let col_norms = matrix.norm_dims(&[0], false);
assert_eq!(col_norms.shape().dims(), vec![2]);
assert!((col_norms.get(&[0]) - 3.0).abs() < 1e-6); // sqrt(3² + 0²)
assert!((col_norms.get(&[1]) - 6.403).abs() < 1e-3); // sqrt(4² + 5²)

use train_station::Tensor;

// Norm over multiple dimensions
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let norm_all = tensor.norm_dims(&[0, 1], false);
assert_eq!(norm_all.shape().dims(), vec![1]);
// sqrt(1² + 2² + 3² + 4²) = sqrt(30) ≈ 5.477
assert!((norm_all.get(&[0]) - 30.0_f32.sqrt()).abs() < 1e-5);

Panics

• If dims is empty
• If any dimension index is out of bounds for the tensor rank

Performance

Uses efficient coordinate-based iteration that works correctly with both contiguous and non-contiguous tensor layouts.

impl Tensor

pub fn std(&self) -> Tensor

Computes the standard deviation over all elements

The standard deviation is calculated as sqrt(variance) where variance is the mean of squared differences from the mean. This operation reduces the tensor to a scalar value [1].

The implementation uses population standard deviation (divides by n rather than n-1) to match PyTorch’s default behavior.

Returns

A scalar tensor containing the standard deviation value

Examples

use train_station::Tensor;

// Basic standard deviation calculation
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
let std_dev = tensor.std();
assert!((std_dev.get(&[0]) - 1.118_034).abs() < 1e-5);

use train_station::Tensor;

// Standard deviation of a larger dataset
let data = vec![1.0, 3.0, 5.0, 7.0, 2.0, 4.0, 6.0, 8.0];
let tensor = Tensor::from_slice(&data, vec![2, 2, 2]).unwrap();
let std_dev = tensor.std();
// mean=4.5, var=5.25, std=sqrt(5.25)≈2.291
let expected = 5.25_f32.sqrt();
assert!((std_dev.get(&[0]) - expected).abs() < 1e-5);

use train_station::Tensor;

// Standard deviation of constant values (should be 0)
let tensor = Tensor::from_slice(&[5.0, 5.0, 5.0, 5.0], vec![4]).unwrap();
let std_dev = tensor.std();
assert!((std_dev.get(&[0]) - 0.0).abs() < 1e-6);

Performance

Uses optimized contiguous tensor path with 4x loop unrolling for better performance. Non-contiguous tensors use stride-aware iteration. The algorithm performs two passes: first to compute the mean, then to compute the variance.

Examples found in repository

examples/iterators/advanced_patterns.rs (line 111)

fn demonstrate_data_pipeline() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Data Processing Pipeline ---");

    // Simulate raw sensor data with noise
    let raw_data: Vec<f32> = (0..20)
        .map(|i| {
            let base = i as f32 * 0.5;
            let noise = (i % 3) as f32 * 0.1;
            base + noise
        })
        .collect();

    let tensor = Tensor::from_slice(&raw_data, vec![20])?;
    println!("Raw sensor data: {:?}", tensor.data());

    // Multi-stage processing pipeline
    println!("\nProcessing pipeline:");
    println!("1. Normalize data (z-score)");
    println!("2. Apply smoothing filter");
    println!("3. Detect outliers");
    println!("4. Apply feature scaling");

    // Stage 1: Normalization
    let mean = tensor.mean().value();
    let std = tensor.std().value();
    let normalized: Tensor = tensor
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();
    println!(
        "  Normalized (mean={:.3}, std={:.3}): {:?}",
        mean,
        std,
        normalized.data()
    );

    // Stage 2: Smoothing (simple moving average)
    let smoothed: Tensor = normalized
        .iter()
        .enumerate()
        .map(|(i, elem)| {
            if i == 0 || i == normalized.size() - 1 {
                elem.clone()
            } else {
                // Simple 3-point average
                let prev = normalized.element_view(i - 1);
                let next = normalized.element_view(i + 1);
                elem.add_tensor(&prev).add_tensor(&next).div_scalar(3.0)
            }
        })
        .collect();
    println!("  Smoothed: {:?}", smoothed.data());

    // Stage 3: Outlier detection and removal
    let outlier_threshold = 2.0;
    let cleaned: Tensor = smoothed
        .iter()
        .filter(|elem| elem.value().abs() < outlier_threshold)
        .collect();
    println!(
        "  Outliers removed (threshold={}): {:?}",
        outlier_threshold,
        cleaned.data()
    );

    // Stage 4: Feature scaling to [0, 1] range
    let min_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max_val = cleaned
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);
    let scaled: Tensor = cleaned
        .iter()
        .map(|elem| elem.sub_scalar(min_val).div_scalar(max_val - min_val))
        .collect();
    println!("  Scaled to [0,1]: {:?}", scaled.data());

    Ok(())
}

/// Demonstrate conditional processing patterns
///
/// Shows how to implement dynamic filtering and transformation
/// based on data characteristics and conditions.
fn demonstrate_conditional_processing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Conditional Processing ---");

    // Create data with mixed characteristics
    let data = vec![1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0, 9.0, -10.0];
    let tensor = Tensor::from_slice(&data, vec![10])?;
    println!("Input data: {:?}", tensor.data());

    // Conditional transformation based on sign
    println!("\nConditional transformation (positive/negative handling):");
    let processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val > 0.0 {
                elem.pow_scalar(2.0) // Square positive values
            } else {
                elem.mul_scalar(-1.0).sqrt() // Square root of absolute negative values
            }
        })
        .collect();
    println!("  Processed: {:?}", processed.data());

    // Adaptive filtering based on local statistics
    println!("\nAdaptive filtering (remove values > 2 std from local mean):");
    let window_size = 3;
    let adaptive_filtered: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, elem)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(tensor.size());

            // Calculate local mean and std
            let local_values: Vec<f32> = (start..end)
                .map(|j| tensor.element_view(j).value())
                .collect();

            let local_mean = local_values.iter().sum::<f32>() / local_values.len() as f32;
            let local_variance = local_values
                .iter()
                .map(|v| (v - local_mean).powi(2))
                .sum::<f32>()
                / local_values.len() as f32;
            let local_std = local_variance.sqrt();

            let threshold = local_mean + 2.0 * local_std;
            elem.value() <= threshold
        })
        .map(|(_, elem)| elem)
        .collect();
    println!("  Adaptive filtered: {:?}", adaptive_filtered.data());

    // Multi-condition processing
    println!("\nMulti-condition processing:");
    let multi_processed: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            match () {
                _ if val > 5.0 => elem.mul_scalar(2.0), // Double large values
                _ if val < -5.0 => elem.div_scalar(2.0), // Halve small values
                _ if val.abs() < 2.0 => elem.add_scalar(1.0), // Add 1 to small values
                _ => elem.clone(),                      // Keep others unchanged
            }
        })
        .collect();
    println!("  Multi-condition: {:?}", multi_processed.data());

    Ok(())
}

/// Demonstrate batch processing operations
///
/// Shows efficient processing of large datasets using iterator
/// patterns and batch operations for performance optimization.
fn demonstrate_batch_operations() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Batch Operations ---");

    // Create a larger dataset for batch processing
    let size = 100;
    let data: Vec<f32> = (0..size)
        .map(|i| {
            let x = i as f32 / size as f32;
            x * x + 0.1 * (i % 7) as f32 // Quadratic with some noise
        })
        .collect();

    let tensor = Tensor::from_slice(&data, vec![size])?;
    println!("Dataset size: {}", tensor.size());

    // Batch processing with windowing (iterator views)
    println!("\nBatch processing with sliding windows:");
    let batch_size = 10;
    let batches: Vec<Tensor> = tensor
        .iter()
        .collect::<Vec<_>>()
        .chunks(batch_size)
        .map(|chunk| {
            // Process each batch independently
            chunk
                .iter()
                .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
                .collect()
        })
        .collect();

    println!(
        "  Processed {} batches of size {}",
        batches.len(),
        batch_size
    );
    for (i, batch) in batches.iter().enumerate() {
        println!(
            "    Batch {}: mean={:.3}, std={:.3}",
            i,
            batch.mean().value(),
            batch.std().value()
        );
    }

    // Parallel-like processing with stride
    println!("\nStrided processing (every nth element):");
    let stride = 5;
    let strided: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % stride == 0)
        .map(|(_, elem)| elem)
        .collect();
    println!("  Strided (every {}th): {:?}", stride, strided.data());

    // Hierarchical processing
    println!("\nHierarchical processing (coarse to fine):");
    let coarse: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 == 0) // Take every 4th element
        .map(|(_, elem)| elem)
        .collect();

    let fine: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 4 != 0) // Take the rest
        .map(|(_, elem)| elem)
        .collect();

    println!("  Coarse (every 4th): {:?}", coarse.data());
    println!("  Fine (rest): {:?}", fine.data());

    // Combine coarse and fine with different processing
    let combined: Tensor = coarse
        .iter()
        .map(|elem| elem.mul_scalar(2.0)) // Scale coarse
        .chain(fine.iter().map(|elem| elem.div_scalar(2.0))) // Scale fine
        .collect();
    println!("  Combined: {:?}", combined.data());

    Ok(())
}

/// Demonstrate real-world processing scenarios
///
/// Shows practical applications of iterator patterns for
/// common data processing tasks in machine learning and analytics.
fn demonstrate_real_world_scenarios() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Real-world Scenarios ---");

    // Scenario 1: Time series analysis
    println!("\nScenario 1: Time Series Analysis");
    let time_series: Vec<f32> = (0..24)
        .map(|hour| {
            let base = 20.0 + 10.0 * (hour as f32 * std::f32::consts::PI / 12.0).sin();
            base + (hour % 3) as f32 * 2.0 // Add some noise
        })
        .collect();

    let series = Tensor::from_slice(&time_series, vec![24])?;
    println!("  Time series (24 hours): {:?}", series.data());

    // Calculate moving average with view-based iteration
    let window_size = 3;
    let moving_avg: Tensor = series
        .iter()
        .enumerate()
        .map(|(i, _)| {
            let start = i.saturating_sub(window_size / 2);
            let end = (i + window_size / 2 + 1).min(series.size());
            let window = series.iter_range(start, end);
            window.fold(0.0, |acc, elem| acc + elem.value()) / (end - start) as f32
        })
        .map(|val| Tensor::from_slice(&[val], vec![1]).unwrap())
        .collect();
    println!(
        "  Moving average (window={}): {:?}",
        window_size,
        moving_avg.data()
    );

    // Inference pipeline with NoGrad + streaming
    println!("\nInference pipeline (NoGrad + streaming)");
    let features = Tensor::from_slice(
        &(0..48).map(|i| i as f32 * 0.125).collect::<Vec<_>>(),
        vec![6, 8],
    )?;
    let fast = with_no_grad(|| {
        // Stream values directly, apply light affine, and collect back to same shape
        features
            .data()
            .iter()
            .copied()
            .map(|x| 0.75 * x + 0.1)
            .collect_shape(vec![6, 8])
    });
    println!(
        "  NoGrad streamed transform shape: {:?}",
        fast.shape().dims()
    );

    // Row-wise iteration with shape-preserving collection (GradTrack-friendly)
    let per_row: Tensor = features
        .iter()
        .map(|row| row.mul_scalar(0.5).add_scalar(2.0))
        .collect_shape(vec![6, 8]);
    println!("  Row-wise mapped shape: {:?}", per_row.shape().dims());

    // Scenario 2: Feature engineering
    println!("\nScenario 2: Feature Engineering");
    let features = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5])?;
    println!("  Original features: {:?}", features.data());

    // Create polynomial features
    let poly_features: Tensor = features
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // x^1
                elem.pow_scalar(2.0), // x^2
                elem.pow_scalar(3.0), // x^3
            ]
        })
        .collect();
    println!(
        "  Polynomial features (x, x^2, x^3): {:?}",
        poly_features.data()
    );

    // Scenario 3: Data augmentation
    println!("\nScenario 3: Data Augmentation");
    let original = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?;
    println!("  Original data: {:?}", original.data());

    // Augment with noise and scaling
    let augmented: Tensor = original
        .iter()
        .flat_map(|elem| {
            vec![
                elem.clone(),         // Original
                elem.add_scalar(0.1), // Add noise
                elem.sub_scalar(0.1), // Subtract noise
                elem.mul_scalar(1.1), // Scale up
                elem.mul_scalar(0.9), // Scale down
            ]
        })
        .collect();
    println!("  Augmented data: {:?}", augmented.data());

    // Scenario 4: Statistical analysis
    println!("\nScenario 4: Statistical Analysis");
    let sample_data = Tensor::from_slice(&[1.1, 2.3, 1.8, 2.1, 1.9, 2.0, 1.7, 2.2], vec![8])?;
    println!("  Sample data: {:?}", sample_data.data());

    // Calculate various statistics
    let mean = sample_data.mean().value();
    let std = sample_data.std().value();
    let min = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::INFINITY, f32::min);
    let max = sample_data
        .iter()
        .map(|e| e.value())
        .fold(f32::NEG_INFINITY, f32::max);

    // Z-score normalization
    let z_scores: Tensor = sample_data
        .iter()
        .map(|elem| elem.sub_scalar(mean).div_scalar(std))
        .collect();

    println!(
        "  Statistics: mean={:.3}, std={:.3}, min={:.3}, max={:.3}",
        mean, std, min, max
    );
    println!("  Z-scores: {:?}", z_scores.data());

    Ok(())
}

pub fn std_dims(&self, dims: &[usize], keepdim: bool) -> Tensor

Computes the standard deviation over specified dimensions

Reduces the tensor along the specified dimensions by computing the standard deviation of each slice. The result maintains the original tensor structure with reduced dimensions optionally preserved as size-1 dimensions.

Uses population standard deviation (divides by n rather than n-1) to match PyTorch’s default behavior.

Arguments

• dims - Vector of dimension indices to reduce over (must be valid for tensor rank)
• keepdim - Whether to keep reduced dimensions as size-1 dimensions

Returns

A tensor with standard deviation computed over the specified dimensions

Examples

use train_station::Tensor;

// Standard deviation along rows (dimension 1) with keepdim=true
let matrix = Tensor::from_slice(&[1.0, 3.0, 2.0, 2.0], vec![2, 2]).unwrap();
let row_stds = matrix.std_dims(&[1], true);
assert_eq!(row_stds.shape().dims(), vec![2, 1]);
assert!((row_stds.get(&[0, 0]) - 1.0).abs() < 1e-6); // std([1, 3]) = 1.0
assert!((row_stds.get(&[1, 0]) - 0.0).abs() < 1e-6); // std([2, 2]) = 0.0

use train_station::Tensor;

// Standard deviation along columns (dimension 0) with keepdim=false
let matrix = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let col_stds = matrix.std_dims(&[0], false);
assert_eq!(col_stds.shape().dims(), vec![2]);
// std([1, 3]) = 1.0, std([2, 4]) = 1.0
assert!((col_stds.get(&[0]) - 1.0).abs() < 1e-6);
assert!((col_stds.get(&[1]) - 1.0).abs() < 1e-6);

use train_station::Tensor;

// Standard deviation over multiple dimensions
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let std_all = tensor.std_dims(&[0, 1], false);
assert_eq!(std_all.shape().dims(), vec![1]);
// std([1, 2, 3, 4]) = sqrt(1.25) ≈ 1.118
assert!((std_all.get(&[0]) - 1.25_f32.sqrt()).abs() < 1e-5);

Panics

• If dims is empty
• If any dimension index is out of bounds for the tensor rank
• If the reduced size is 0 (invalid for standard deviation calculation)

Performance

Uses efficient coordinate-based iteration that works correctly with both contiguous and non-contiguous tensor layouts. The algorithm performs two passes: first to compute means, then to compute variances.

impl Tensor

pub fn sum(&self) -> Tensor

Returns the sum of all elements in the tensor

This operation computes the sum of all elements across all dimensions, reducing the tensor to a scalar value. The output is a tensor with shape [1] containing the sum as a float.

When requires_grad is enabled, this operation supports automatic gradient tracking through the GradTrack system.

Returns

A tensor with shape [1] containing the sum of all elements

Examples

use train_station::Tensor;

// Basic sum calculation
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let total = tensor.sum();
assert_eq!(total.shape().dims(), vec![1]);
assert_eq!(total.get(&[0]), 10.0); // 1 + 2 + 3 + 4 = 10

use train_station::Tensor;

// Sum with gradient tracking
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])
    .unwrap()
    .with_requires_grad();
let mut total = tensor.sum();
total.backward(None);
let grad = tensor.grad_owned().expect("gradient should exist");
// Gradient should be [1.0, 1.0, 1.0] for each element
assert_eq!(grad.get(&[0]), 1.0);
assert_eq!(grad.get(&[1]), 1.0);
assert_eq!(grad.get(&[2]), 1.0);

use train_station::Tensor;

// Sum of empty tensor
let tensor = Tensor::new(vec![0]);
let total = tensor.sum();
assert_eq!(total.get(&[0]), 0.0); // Sum of empty tensor is 0

Performance

Uses optimized contiguous tensor path with 4x loop unrolling for better performance. Non-contiguous tensors use stride-aware iteration.

Examples found in repository

examples/supervised_training/supervised_regression.rs (line 54)

fn r2_score(pred: &Tensor, target: &Tensor) -> f32 {
    // R^2 = 1 - SS_res / SS_tot
    let y = target;
    let y_mean = y.mean();
    let ss_res = pred.sub_tensor(y).pow_scalar(2.0).sum();
    let ss_tot = y.sub_tensor(&y_mean).pow_scalar(2.0).sum();
    let ss_res_v = ss_res.value();
    let ss_tot_v = ss_tot.value().max(1e-12); // avoid divide by zero
    1.0 - (ss_res_v / ss_tot_v)
}

examples/getting_started/tensor_basics.rs (line 191)

fn demonstrate_utility_functions() {
    println!("\n--- Utility Functions ---");

    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();

    // Basic properties
    println!("Shape: {:?}", tensor.shape().dims());
    println!("Size: {}", tensor.size());
    println!("Is contiguous: {}", tensor.is_contiguous());
    println!("Device: {:?}", tensor.device());

    // Mathematical operations
    let sum = tensor.sum();
    println!("Sum: {}", sum.value());

    let mean = tensor.mean();
    println!("Mean: {}", mean.value());

    let norm = tensor.norm();
    println!("Norm: {}", norm.value());

    // Device placement
    let cpu_tensor = Tensor::zeros_on_device(vec![3, 3], train_station::Device::cpu());
    println!(
        "CPU tensor: shape {:?}, device: {:?}",
        cpu_tensor.shape().dims(),
        cpu_tensor.device()
    );
}

examples/iterators/element_iteration.rs (line 195)

fn demonstrate_gradient_tracking() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Gradient Tracking ---");

    // Create a tensor with gradient tracking enabled
    let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3])?.with_requires_grad();
    println!("Input tensor (requires_grad): {:?}", tensor.data());

    // Perform element-wise operations through iteration
    let result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Apply a complex transformation: (x^2 + 1) * 2
            elem.pow_scalar(2.0).add_scalar(1.0).mul_scalar(2.0)
        })
        .collect();

    println!("Result tensor: {:?}", result.data());
    println!("Result requires_grad: {}", result.requires_grad());

    // Compute gradients
    let mut loss = result.sum();
    loss.backward(None);

    println!("Loss: {:.6}", loss.value());
    println!("Input gradients: {:?}", tensor.grad().map(|g| g.data()));

    Ok(())
}

examples/getting_started/serialization_basics.rs (line 137)

fn demonstrate_optimizer_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Optimizer Serialization ---");

    // Create an optimizer with some parameters
    let mut weight = Tensor::randn(vec![2, 2], Some(42)).with_requires_grad();
    let mut bias = Tensor::randn(vec![2], Some(43)).with_requires_grad();

    let config = AdamConfig {
        learning_rate: 0.001,
        beta1: 0.9,
        beta2: 0.999,
        eps: 1e-8,
        weight_decay: 0.0,
        amsgrad: false,
    };

    let mut optimizer = Adam::with_config(config);
    optimizer.add_parameter(&weight);
    optimizer.add_parameter(&bias);

    println!(
        "Created optimizer with {} parameters",
        optimizer.parameter_count()
    );
    println!("Learning rate: {}", optimizer.learning_rate());

    // Simulate some training steps
    for _ in 0..3 {
        let mut loss = weight.sum() + bias.sum();
        loss.backward(None);
        optimizer.step(&mut [&mut weight, &mut bias]);
        optimizer.zero_grad(&mut [&mut weight, &mut bias]);
    }

    // Save optimizer state
    let optimizer_path = "temp_optimizer.json";
    optimizer.save_json(optimizer_path)?;
    println!("Saved optimizer to: {}", optimizer_path);

    // Load optimizer state
    let loaded_optimizer = Adam::load_json(optimizer_path)?;
    println!(
        "Loaded optimizer with {} parameters",
        loaded_optimizer.parameter_count()
    );
    println!("Learning rate: {}", loaded_optimizer.learning_rate());

    // Verify optimizer state
    assert_eq!(
        optimizer.parameter_count(),
        loaded_optimizer.parameter_count()
    );
    assert_eq!(optimizer.learning_rate(), loaded_optimizer.learning_rate());
    println!("Optimizer serialization verification: PASSED");

    Ok(())
}

/// Demonstrate format comparison and performance characteristics
fn demonstrate_format_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Format Comparison ---");

    // Create a larger tensor for comparison
    let tensor = Tensor::randn(vec![10, 10], Some(44));

    // Save in both formats
    tensor.save_json("temp_comparison.json")?;
    tensor.save_binary("temp_comparison.bin")?;

    // Compare file sizes
    let json_size = fs::metadata("temp_comparison.json")?.len();
    let binary_size = fs::metadata("temp_comparison.bin")?.len();

    println!("JSON file size: {} bytes", json_size);
    println!("Binary file size: {} bytes", binary_size);
    println!(
        "Compression ratio: {:.2}x",
        json_size as f64 / binary_size as f64
    );

    // Load and verify both formats
    let json_tensor = Tensor::load_json("temp_comparison.json")?;
    let binary_tensor = Tensor::load_binary("temp_comparison.bin")?;

    assert_eq!(tensor.shape().dims(), json_tensor.shape().dims());
    assert_eq!(tensor.shape().dims(), binary_tensor.shape().dims());
    assert_eq!(tensor.data(), json_tensor.data());
    assert_eq!(tensor.data(), binary_tensor.data());

    println!("Format comparison verification: PASSED");

    Ok(())
}

/// Demonstrate a basic model checkpointing workflow
fn demonstrate_model_checkpointing() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Model Checkpointing ---");

    // Create a simple model (weights and bias)
    let mut weights = Tensor::randn(vec![2, 1], Some(45)).with_requires_grad();
    let mut bias = Tensor::randn(vec![1], Some(46)).with_requires_grad();

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.01);
    optimizer.add_parameter(&weights);
    optimizer.add_parameter(&bias);

    println!("Initial weights: {:?}", weights.data());
    println!("Initial bias: {:?}", bias.data());

    // Simulate training
    for epoch in 0..5 {
        let mut loss = weights.sum() + bias.sum();
        loss.backward(None);
        optimizer.step(&mut [&mut weights, &mut bias]);
        optimizer.zero_grad(&mut [&mut weights, &mut bias]);

        if epoch % 2 == 0 {
            // Save checkpoint
            let checkpoint_dir = format!("checkpoint_epoch_{}", epoch);
            fs::create_dir_all(&checkpoint_dir)?;

            weights.save_json(format!("{}/weights.json", checkpoint_dir))?;
            bias.save_json(format!("{}/bias.json", checkpoint_dir))?;
            optimizer.save_json(format!("{}/optimizer.json", checkpoint_dir))?;

            println!("Saved checkpoint for epoch {}", epoch);
        }
    }

    // Load from checkpoint
    let loaded_weights = Tensor::load_json("checkpoint_epoch_4/weights.json")?;
    let loaded_bias = Tensor::load_json("checkpoint_epoch_4/bias.json")?;
    let loaded_optimizer = Adam::load_json("checkpoint_epoch_4/optimizer.json")?;

    println!("Loaded weights: {:?}", loaded_weights.data());
    println!("Loaded bias: {:?}", loaded_bias.data());
    println!(
        "Loaded optimizer learning rate: {}",
        loaded_optimizer.learning_rate()
    );

    // Verify checkpoint integrity
    assert_eq!(weights.shape().dims(), loaded_weights.shape().dims());
    assert_eq!(bias.shape().dims(), loaded_bias.shape().dims());
    assert_eq!(optimizer.learning_rate(), loaded_optimizer.learning_rate());

    println!("Checkpointing verification: PASSED");

    Ok(())
}

examples/iterators/performance_optimization.rs (line 409)

fn demonstrate_optimization_techniques() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Optimization Techniques ---");

    let size = 50000;
    let data: Vec<f32> = (0..size).map(|i| i as f32).collect();
    let tensor = Tensor::from_slice(&data, vec![size])?;

    println!("Optimizing processing for size: {}", size);

    // Technique 1: Operation fusion
    println!("\nTechnique 1: Operation Fusion");
    let start = Instant::now();
    let fused_result: Tensor = tensor
        .iter()
        .map(|elem| {
            // Fuse multiple operations into single chain
            elem.mul_scalar(2.0).add_scalar(1.0).pow_scalar(2.0).sqrt()
        })
        .collect();
    let fused_time = start.elapsed();

    // Technique 2: Conditional optimization
    println!("\nTechnique 2: Conditional Optimization");
    let start = Instant::now();
    let conditional_result: Tensor = tensor
        .iter()
        .map(|elem| {
            let val = elem.value();
            if val < size as f32 / 2.0 {
                elem.mul_scalar(2.0) // Simple operation for small values
            } else {
                elem.pow_scalar(2.0).sqrt() // Complex operation for large values
            }
        })
        .collect();
    let conditional_time = start.elapsed();

    // Technique 3: Cache-friendly processing
    println!("\nTechnique 3: Cache-Friendly Processing");
    let start = Instant::now();
    let cache_friendly_result: Tensor = tensor
        .iter()
        .take(1000) // Process in cache-friendly chunks
        .map(|elem| elem.mul_scalar(2.0))
        .collect();
    let cache_friendly_time = start.elapsed();

    // Technique 4: Memory pooling simulation
    println!("\nTechnique 4: Memory Pooling Simulation");
    let start = Instant::now();
    let pooled_result: Tensor = tensor
        .iter()
        .enumerate()
        .filter(|(i, _)| i % 100 == 0) // Process every 100th element
        .map(|(_, elem)| elem.pow_scalar(2.0))
        .collect();
    let pooled_time = start.elapsed();

    // Report optimization results
    println!("  Fused operations: {:?}", fused_time);
    println!("  Conditional optimization: {:?}", conditional_time);
    println!("  Cache-friendly processing: {:?}", cache_friendly_time);
    println!("  Memory pooling simulation: {:?}", pooled_time);

    // Performance analysis
    let fastest = fused_time
        .min(conditional_time)
        .min(cache_friendly_time)
        .min(pooled_time);
    println!("  Fastest technique: {:?}", fastest);

    // Memory efficiency analysis
    println!("  Fused result size: {}", fused_result.size());
    println!("  Conditional result size: {}", conditional_result.size());
    println!(
        "  Cache-friendly result size: {}",
        cache_friendly_result.size()
    );
    println!("  Pooled result size: {}", pooled_result.size());

    // Technique 5: Gradient optimization
    println!("\nTechnique 5: Gradient Optimization");
    let grad_tensor = tensor.with_requires_grad();
    let start = Instant::now();

    let grad_result: Tensor = grad_tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).add_scalar(1.0))
        .collect();

    let mut loss = grad_result.sum();
    loss.backward(None);
    let grad_time = start.elapsed();

    println!("  Gradient computation: {:?}", grad_time);
    println!(
        "  Gradient tracking enabled: {}",
        grad_result.requires_grad()
    );

    Ok(())
}

pub fn sum_dims(&self, dims: &[usize], keepdim: bool) -> Tensor

Returns the sum of elements along specified dimensions

This operation computes the sum of elements along the specified dimensions, reducing the tensor while optionally preserving the reduced dimensions as size-1 dimensions.

The output shape depends on the keepdim parameter:

• If keepdim is true, the reduced dimensions are kept with size 1
• If keepdim is false, the reduced dimensions are removed

When requires_grad is enabled, this operation supports automatic gradient tracking through the GradTrack system.

Arguments

• dims - Vector of dimension indices to sum over (must be valid for tensor rank)
• keepdim - Whether to keep reduced dimensions as size-1 dimensions

Returns

A tensor with sum computed over the specified dimensions

Panics

• If dims is empty
• If any dimension index is out of bounds for the tensor rank

Examples

use train_station::Tensor;

// Sum along rows (dimension 0) with keepdim=false
let matrix = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let row_sums = matrix.sum_dims(&[0], false);
assert_eq!(row_sums.shape().dims(), vec![2]);
assert_eq!(row_sums.get(&[0]), 4.0); // 1 + 3 = 4
assert_eq!(row_sums.get(&[1]), 6.0); // 2 + 4 = 6

use train_station::Tensor;

// Sum along columns (dimension 1) with keepdim=true
let matrix = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let col_sums = matrix.sum_dims(&[1], true);
assert_eq!(col_sums.shape().dims(), vec![2, 1]);
assert_eq!(col_sums.get(&[0, 0]), 3.0); // 1 + 2 = 3
assert_eq!(col_sums.get(&[1, 0]), 7.0); // 3 + 4 = 7

use train_station::Tensor;

// Sum over multiple dimensions
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let total = tensor.sum_dims(&[0, 1], false);
assert_eq!(total.shape().dims(), vec![1]);
assert_eq!(total.get(&[0]), 10.0); // 1 + 2 + 3 + 4 = 10

use train_station::Tensor;

// Sum with gradient tracking
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2])
    .unwrap()
    .with_requires_grad();
let mut row_sums = tensor.sum_dims(&[0], false);
row_sums.backward(None);
let grad = tensor.grad_owned().expect("gradient should exist");
// Gradient should be [1.0, 1.0, 1.0, 1.0] for each element
assert_eq!(grad.get(&[0, 0]), 1.0);
assert_eq!(grad.get(&[0, 1]), 1.0);
assert_eq!(grad.get(&[1, 0]), 1.0);
assert_eq!(grad.get(&[1, 1]), 1.0);

Performance

Uses efficient coordinate-based iteration that works correctly with both contiguous and non-contiguous tensor layouts.

Examples found in repository

examples/supervised_training/supervised_classification.rs (line 54)

fn cross_entropy_logits(
    logits: &Tensor,
    labels: &[usize],
    batch: usize,
    _num_classes: usize,
) -> Tensor {
    // log_softmax = logits - logsumexp(logits, dim=1)
    let max_logits = logits.max_dims(&[1], true);
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true);
    let log_sum_exp = sum_exp.log();
    let log_softmax = shifted.sub_tensor(&log_sum_exp);
    let ll = log_softmax.gather(1, labels, &[batch, 1]); // selected log-probs
    ll.mul_scalar(-1.0).mean()
}

examples/RL_training/ppo_discrete.rs (line 282)

fn log_prob_actions(
    logits: &Tensor,
    actions: &[usize],
    batch: usize,
    _action_dim: usize,
) -> Tensor {
    let max_logits = logits.max_dims(&[1], true); // [B,1]
    let shifted = logits.sub_tensor(&max_logits);
    let exp = shifted.exp();
    let sum_exp = exp.sum_dims(&[1], true); // [B,1]
    let log_sum_exp = sum_exp.log(); // [B,1]
    let log_softmax = shifted.sub_tensor(&log_sum_exp); // [B,A]
                                                        // gather selected action log-probs
    log_softmax.gather(1, actions, &[batch, 1])
}

// probability ratio = exp(new_logp - old_logp)
fn ratio_from_logps(new_logp: &Tensor, old_logp: &Tensor) -> Tensor {
    new_logp.sub_tensor(old_logp).exp()
}

// Clamp ratio to [1-clip, 1+clip] using ReLU-based clamp (no custom ops)
fn clamp_ratio(ratio: &Tensor, clip_eps: f32) -> Tensor {
    let b = ratio.shape().dims()[0];
    let low = Tensor::from_slice(&vec![1.0 - clip_eps; b], vec![b, 1]).unwrap();
    let high = Tensor::from_slice(&vec![1.0 + clip_eps; b], vec![b, 1]).unwrap();
    let ge_low = ratio.sub_tensor(&low).relu().add_tensor(&low);
    high.sub_tensor(&ge_low.sub_tensor(&high).relu())
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/RL_training/ppo_continuous.rs (line 248)

fn gaussian_log_prob(action: &Tensor, mean: &Tensor, log_std: &Tensor) -> Tensor {
    // All tensors shaped [B, A] (log_std is broadcastable)
    let std = log_std.exp();
    let var = std.pow_scalar(2.0);
    let log_scale = log_std;
    let diff = action.sub_tensor(mean);
    let log_prob = diff
        .pow_scalar(2.0)
        .div_tensor(&var)
        .add_scalar(std::f32::consts::LN_2 + std::f32::consts::PI)
        .add_tensor(&log_scale.mul_scalar(2.0))
        .mul_scalar(0.5)
        .mul_scalar(-1.0);
    // Sum across action dim (dim=1) -> [B,1]
    log_prob.sum_dims(&[1], true)
}

#[allow(clippy::too_many_arguments)]
fn compute_gae(
    returns_out: &mut [f32],
    adv_out: &mut [f32],
    rewards: &[f32],
    dones: &[f32],
    values: &[f32],
    next_values: &[f32],
    gamma: f32,
    lam: f32,
) {
    let n = rewards.len();
    let mut gae = 0.0f32;
    for t in (0..n).rev() {
        let not_done = 1.0 - dones[t];
        let delta = rewards[t] + gamma * next_values[t] * not_done - values[t];
        gae = delta + gamma * lam * not_done * gae;
        adv_out[t] = gae;
        returns_out[t] = gae + values[t];
    }
}

fn normalize_in_place(x: &mut [f32], eps: f32) {
    let n = x.len() as f32;
    if n <= 1.0 {
        return;
    }
    let mean = x.iter().copied().sum::<f32>() / n;
    let var = x
        .iter()
        .map(|v| {
            let d = v - mean;
            d * d
        })
        .sum::<f32>()
        / n;
    let std = (var + eps).sqrt();
    for v in x.iter_mut() {
        *v = (*v - mean) / std;
    }
}

fn clip_gradients(parameters: &mut [&mut Tensor], max_norm: f32, eps: f32) {
    let mut total_sq = 0.0f32;
    for p in parameters.iter() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    let norm = total_sq.sqrt();
    if norm > max_norm {
        let scale = max_norm / (norm + eps);
        for p in parameters.iter_mut() {
            if let Some(g) = p.grad_owned() {
                p.set_grad(g.mul_scalar(scale));
            }
        }
    }
}

fn grad_global_norm(parameters: &mut [&mut Tensor]) -> f32 {
    let mut total_sq = 0.0f32;
    for p in parameters.iter_mut() {
        if let Some(g) = p.grad_owned() {
            for &v in g.data() {
                total_sq += v * v;
            }
        }
    }
    total_sq.sqrt()
}

// -------------------------------
// Main
// -------------------------------

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

impl Tensor

pub fn var(&self) -> Tensor

Computes the variance over all elements

The variance is calculated as the mean of squared differences from the mean. This operation reduces the tensor to a scalar value [1].

The implementation uses population variance (divides by n rather than n-1) to match PyTorch’s default behavior.

Returns

A scalar tensor containing the variance value

Examples

use train_station::Tensor;

// Basic variance calculation
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
let variance = tensor.var();
assert!((variance.get(&[0]) - 1.25).abs() < 1e-5);

use train_station::Tensor;

// Variance of a larger dataset
let data = vec![1.0, 3.0, 5.0, 7.0, 2.0, 4.0, 6.0, 8.0];
let tensor = Tensor::from_slice(&data, vec![2, 2, 2]).unwrap();
let variance = tensor.var();
// mean=4.5, var=mean([3.5², 1.5², 0.5², 2.5², 2.5², 0.5², 1.5², 3.5²]) = 5.25
assert!((variance.get(&[0]) - 5.25).abs() < 1e-5);

use train_station::Tensor;

// Variance of constant values (should be 0)
let tensor = Tensor::from_slice(&[5.0, 5.0, 5.0, 5.0], vec![4]).unwrap();
let variance = tensor.var();
assert!((variance.get(&[0]) - 0.0).abs() < 1e-6);

Performance

Uses optimized contiguous tensor path with manual loop unrolling for better performance. Non-contiguous tensors use stride-aware iteration. The algorithm performs two passes: first to compute the mean, then to compute the variance.

pub fn var_dims(&self, dims: &[usize], keepdim: bool) -> Tensor

Computes the variance over specified dimensions

Reduces the tensor along the specified dimensions by computing the variance of each slice. The result maintains the original tensor structure with reduced dimensions optionally preserved as size-1 dimensions.

Uses population variance (divides by n rather than n-1) to match PyTorch’s default behavior.

Arguments

• dims - Vector of dimension indices to reduce over (must be valid for tensor rank)
• keepdim - Whether to keep reduced dimensions as size-1 dimensions

Returns

A tensor with variance computed over the specified dimensions

Examples

use train_station::Tensor;

// Variance along rows (dimension 1) with keepdim=true
let matrix = Tensor::from_slice(&[1.0, 3.0, 2.0, 2.0], vec![2, 2]).unwrap();
let row_vars = matrix.var_dims(&[1], true);
assert_eq!(row_vars.shape().dims(), vec![2, 1]);
assert!((row_vars.get(&[0, 0]) - 1.0).abs() < 1e-6); // var([1, 3]) = 1.0
assert!((row_vars.get(&[1, 0]) - 0.0).abs() < 1e-6); // var([2, 2]) = 0.0

use train_station::Tensor;

// Variance along columns (dimension 0) with keepdim=false
let matrix = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let col_vars = matrix.var_dims(&[0], false);
assert_eq!(col_vars.shape().dims(), vec![2]);
// var([1, 3]) = 1.0, var([2, 4]) = 1.0
assert!((col_vars.get(&[0]) - 1.0).abs() < 1e-6);
assert!((col_vars.get(&[1]) - 1.0).abs() < 1e-6);

use train_station::Tensor;

// Variance over multiple dimensions
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let var_all = tensor.var_dims(&[0, 1], false);
assert_eq!(var_all.shape().dims(), vec![1]);
// var([1, 2, 3, 4]) = 1.25
assert!((var_all.get(&[0]) - 1.25).abs() < 1e-5);

Panics

• If dims is empty
• If any dimension index is out of bounds for the tensor rank
• If the reduced size is 0 (invalid for variance calculation)

Performance

Uses efficient coordinate-based iteration that works correctly with both contiguous and non-contiguous tensor layouts. The algorithm performs two passes: first to compute means, then to compute variances.

impl Tensor

pub fn cat(tensors: &[Tensor], dim: usize) -> Tensor

Concatenate tensors along a given dimension

Joins multiple tensors along the specified dimension, creating a new tensor with the combined data. All input tensors must have the same rank and matching dimensions except for the concatenation dimension.

Arguments

• tensors - Slice of tensors to concatenate (must not be empty)
• dim - Dimension along which to concatenate (must be < tensor rank)

Returns

A new tensor containing the concatenated data with shape where the concatenation dimension is the sum of all input tensor sizes along that dimension.

Panics

• If tensors is empty
• If dim is out of bounds for the tensor rank
• If tensors have different ranks
• If tensors have mismatched dimensions (except along concatenation dimension)

Examples

use train_station::Tensor;

// Concatenate 1D tensors
let a = Tensor::from_slice(&[1.0, 2.0], vec![2]).unwrap();
let b = Tensor::from_slice(&[3.0, 4.0], vec![2]).unwrap();
let result = Tensor::cat(&[a, b], 0);
assert_eq!(result.shape().dims(), vec![4]);
assert_eq!(result.get(&[0]), 1.0);
assert_eq!(result.get(&[1]), 2.0);
assert_eq!(result.get(&[2]), 3.0);
assert_eq!(result.get(&[3]), 4.0);

use train_station::Tensor;

// Concatenate 2D tensors along dimension 1
let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let b = Tensor::from_slice(&[5.0, 6.0], vec![2, 1]).unwrap();
let result = Tensor::cat(&[a, b], 1);
assert_eq!(result.shape().dims(), vec![2, 3]);
assert_eq!(result.get(&[0, 0]), 1.0);
assert_eq!(result.get(&[0, 1]), 2.0);
assert_eq!(result.get(&[0, 2]), 5.0);

use train_station::Tensor;

// Concatenate with gradient tracking
let mut a = Tensor::from_slice(&[1.0, 2.0], vec![2]).unwrap();
let mut b = Tensor::from_slice(&[3.0, 4.0], vec![2]).unwrap();
a.set_requires_grad(true);
b.set_requires_grad(true);

let result = Tensor::cat(&[a, b], 0);
assert!(result.requires_grad());

Examples found in repository

examples/RL_training/td3.rs (line 198)

    fn forward(&self, state: &Tensor, action: &Tensor) -> Tensor {
        // Concatenate along feature dim (dim=1) for batched inputs
        // IMPORTANT: use views to preserve gradient graph; cloning would detach autograd
        let s_view = state.view(state.shape().dims().iter().map(|&d| d as i32).collect());
        let a_view = action.view(action.shape().dims().iter().map(|&d| d as i32).collect());
        let sa = Tensor::cat(&[s_view, a_view], 1);
        self.net.forward(&sa, None)
    }

impl Tensor

pub fn contiguous(&self) -> Tensor

Creates a contiguous copy of the tensor

This operation ensures that the tensor data is stored in a linear, cache-friendly memory layout. If the tensor is already contiguous, this operation returns a clone. For non-contiguous tensors, it creates a new tensor with the same data but in contiguous memory layout.

The operation uses different optimization strategies based on tensor size:

• Small tensors (≤64 elements): Simple coordinate-based copy
• Medium tensors (65-1023 elements): Unrolled copy for better performance
• Large tensors (≥1024 elements): Blocked copy with cache optimization

Returns

A new tensor with contiguous memory layout containing the same data

Examples

use train_station::Tensor;

// Already contiguous tensor
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let contiguous = tensor.contiguous();
assert!(contiguous.is_contiguous());
assert_eq!(contiguous.shape().dims(), vec![2, 2]);

use train_station::Tensor;

// Non-contiguous tensor from transpose
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let transposed = tensor.transpose(0, 1);
assert!(!transposed.is_contiguous());

let contiguous = transposed.contiguous();
assert!(contiguous.is_contiguous());
assert_eq!(contiguous.get(&[0, 0]), 1.0);
assert_eq!(contiguous.get(&[0, 1]), 3.0);

use train_station::Tensor;

// Preserves gradient tracking
let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
tensor.set_requires_grad(true);

let contiguous = tensor.contiguous();
assert!(contiguous.requires_grad());

Performance

• Already contiguous: O(1) time complexity, returns a clone
• Non-contiguous: O(n) time complexity with size-dependent optimizations
• Memory usage: Creates a new tensor with the same size as the original

Examples found in repository

examples/neural_networks/basic_encoder.rs (line 59)

    pub fn forward(&self, input: &Tensor, attn_mask: Option<&Tensor>) -> Tensor {
        let attn = self.mha.forward(input, input, input, attn_mask);
        let res1 = attn.add_tensor(input);

        // Feed-forward network with ReLU and residual
        let (b, t, e) = Self::triple(input);
        let x2d = res1.contiguous().view(vec![(b * t) as i32, e as i32]);
        let hidden = self.ffn_in.forward(&x2d).relu();
        let out2d = self.ffn_out.forward(&hidden);
        let out = out2d.view(vec![b as i32, t as i32, e as i32]);
        out.add_tensor(&res1)
    }

examples/neural_networks/basic_decoder.rs (line 70)

    pub fn forward(
        &self,
        tgt: &Tensor,
        memory: &Tensor,
        causal_mask: Option<&Tensor>,
        cross_mask: Option<&Tensor>,
    ) -> Tensor {
        let self_attn = self.self_attn.forward(tgt, tgt, tgt, causal_mask);
        let res1 = self_attn.add_tensor(tgt);

        let cross = self.cross_attn.forward(&res1, memory, memory, cross_mask);
        let res2 = cross.add_tensor(&res1);

        let (b, t, e) = Self::triple(tgt);
        let x2d = res2.contiguous().view(vec![(b * t) as i32, e as i32]);
        let hidden = self.ffn_in.forward(&x2d).relu();
        let out2d = self.ffn_out.forward(&hidden);
        let out = out2d.view(vec![b as i32, t as i32, e as i32]);
        out.add_tensor(&res2)
    }

examples/neural_networks/multi_head_attention.rs (line 113)

    pub fn forward(
        &self,
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        attn_mask: Option<&Tensor>,
    ) -> Tensor {
        let qkv = Self::project_qkv(query, key, value, &self.q_proj, &self.k_proj, &self.v_proj);
        let (q, k, v) = qkv;

        // Split heads: [b, t, e] -> [b, h, t, d]
        let (b, tq, _e) = Self::triple(query);
        let (_b2, tk, _e2) = Self::triple(key);
        let q = Self::split_heads(&q, b, tq, self.num_heads, self.head_dim);
        let k = Self::split_heads(&k, b, tk, self.num_heads, self.head_dim);
        let v = Self::split_heads(&v, b, tk, self.num_heads, self.head_dim);

        // Scaled dot-product attention
        // logits: [b, h, tq, tk]
        let k_t = k.transpose(2, 3);
        let mut logits = q.matmul(&k_t).div_scalar((self.head_dim as f32).sqrt());
        if let Some(mask) = attn_mask {
            let dims = mask.shape().dims().to_vec();
            // If boolean-like mask matching [b,h,tq,tk], apply masked_fill
            if dims.len() == 4 && dims[0] == b && dims[1] == self.num_heads && dims[2] == tq {
                // Interpret mask > 0.5 as keep; we invert to build masked positions
                let cond: Vec<bool> = mask.data().iter().map(|&v| v < 0.5).collect();
                // Apply masked fill on a flattened view, then reshape back
                let flat_logits = logits.view(vec![(b * self.num_heads * tq * tk) as i32]);
                let filled = flat_logits.masked_fill(&cond, f32::NEG_INFINITY);
                logits = filled.view(vec![b as i32, self.num_heads as i32, tq as i32, tk as i32]);
            } else {
                // Fallback: additive mask
                logits = logits.add_tensor(mask);
            }
        }
        let attn = logits.softmax(3);

        // context: [b, h, tq, d]
        let context = attn.matmul(&v);
        let context = context.permute(vec![0, 2, 1, 3]); // [b, tq, h, d]
        let context = context.contiguous().view(vec![
            b as i32,
            tq as i32,
            (self.num_heads * self.head_dim) as i32,
        ]);

        // Output projection (flatten to 2D, project, then restore 3D)
        let flat = context.view(vec![(b * tq) as i32, self.embed_dim as i32]);
        let out2d = self.out_proj.forward(&flat);
        out2d.view(vec![b as i32, tq as i32, self.embed_dim as i32])
    }

impl Tensor

pub fn flatten(&self) -> Tensor

Flatten the tensor into a 1D representation

Transforms a multi-dimensional tensor into a 1D tensor by reshaping all dimensions into a single dimension. This is equivalent to reshape(vec![-1]) where -1 automatically calculates the size based on the total number of elements.

The flatten operation preserves the total number of elements while changing the tensor’s shape to have a single dimension. This is commonly used in neural networks to prepare tensor data for linear layers or feature extraction.

Returns

A 1D tensor containing the same data as the original tensor, with shape [total_elements] where total_elements is the product of all original dimensions.

Examples

use train_station::Tensor;

// Flatten a 2D tensor
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let flattened = tensor.flatten();
assert_eq!(flattened.shape().dims(), vec![4]);
assert_eq!(flattened.get(&[0]), 1.0);
assert_eq!(flattened.get(&[1]), 2.0);
assert_eq!(flattened.get(&[2]), 3.0);
assert_eq!(flattened.get(&[3]), 4.0);

use train_station::Tensor;

// Flatten a 3D tensor
let data: Vec<f32> = (0..12).map(|i| i as f32).collect();
let tensor = Tensor::from_slice(&data, vec![2, 2, 3]).unwrap();
let flattened = tensor.flatten();
assert_eq!(flattened.shape().dims(), vec![12]);
assert_eq!(flattened.size(), 12);

use train_station::Tensor;

// Flatten with gradient tracking
let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
tensor.set_requires_grad(true);

let flattened = tensor.flatten();
assert!(flattened.requires_grad());
assert_eq!(flattened.shape().dims(), vec![4]);

use train_station::Tensor;

// Flatten an already 1D tensor (no change)
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let flattened = tensor.flatten();
assert_eq!(flattened.shape().dims(), vec![3]);
assert_eq!(flattened.size(), 3);

Performance

• Time Complexity: O(1) - Returns a view when possible
• Memory Usage: No additional memory allocation for view operations
• Gradient Tracking: Preserves gradient requirements and tracking

Relationship to Other Operations

This operation is equivalent to:

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let flattened = tensor.reshape(vec![-1]);

Where -1 is a special value that automatically calculates the dimension size based on the total number of elements in the tensor.

impl Tensor

pub fn permute(&self, dims: Vec<usize>) -> Tensor

Permute tensor dimensions according to specified order

Rearranges the dimensions of the tensor according to the provided dimension order. This operation returns a view with reordered strides, avoiding data copying while changing the logical arrangement of the tensor’s dimensions.

The permutation is specified as a vector where each element represents the new position of the corresponding dimension from the original tensor. For example, permute(vec![1, 0]) swaps the first two dimensions.

Arguments

• dims - Vector specifying the new order of dimensions (must have length equal to tensor rank)

Returns

A new tensor view with rearranged dimensions and correspondingly adjusted strides. The total number of elements remains unchanged.

Panics

• If dims length does not equal the tensor rank
• If any dimension index is out of bounds for the tensor rank
• If dims contains duplicate dimension indices

Examples

use train_station::Tensor;

// Permute 2D tensor (swap dimensions)
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
let permuted = tensor.permute(vec![1, 0]);
assert_eq!(permuted.shape().dims(), vec![3, 2]);
assert_eq!(permuted.get(&[0, 0]), 1.0);
assert_eq!(permuted.get(&[1, 0]), 2.0);
assert_eq!(permuted.get(&[2, 1]), 6.0);

use train_station::Tensor;

// Permute 3D tensor (reorder dimensions)
let data: Vec<f32> = (0..24).map(|i| i as f32).collect();
let tensor = Tensor::from_slice(&data, vec![2, 3, 4]).unwrap();
let permuted = tensor.permute(vec![2, 0, 1]);
assert_eq!(permuted.shape().dims(), vec![4, 2, 3]);
assert_eq!(permuted.size(), 24); // Total elements unchanged

use train_station::Tensor;

// Permute with gradient tracking
let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
tensor.set_requires_grad(true);

let permuted = tensor.permute(vec![1, 0]);
assert!(permuted.requires_grad());
assert_eq!(permuted.shape().dims(), vec![2, 2]);

use train_station::Tensor;

// Identity permutation (no change)
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let permuted = tensor.permute(vec![0, 1]);
assert_eq!(permuted.shape().dims(), vec![2, 2]);
assert_eq!(permuted.get(&[0, 0]), 1.0);
assert_eq!(permuted.get(&[1, 1]), 4.0);

Performance

• Time Complexity: O(1) - Returns a view with reordered strides
• Memory Usage: No additional memory allocation (view operation)
• Gradient Tracking: Preserves gradient requirements and tracking

Relationship to Other Operations

This operation is similar to transpose() but more general:

• transpose(dim0, dim1) is equivalent to permute() with a swap of two dimensions
• permute() can handle arbitrary dimension reordering for tensors of any rank

Memory Layout

The permuted tensor maintains the same underlying data but with reordered strides. This means the tensor becomes non-contiguous unless the permutation is the identity permutation.

Examples found in repository

examples/neural_networks/multi_head_attention.rs (line 112)

    pub fn forward(
        &self,
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        attn_mask: Option<&Tensor>,
    ) -> Tensor {
        let qkv = Self::project_qkv(query, key, value, &self.q_proj, &self.k_proj, &self.v_proj);
        let (q, k, v) = qkv;

        // Split heads: [b, t, e] -> [b, h, t, d]
        let (b, tq, _e) = Self::triple(query);
        let (_b2, tk, _e2) = Self::triple(key);
        let q = Self::split_heads(&q, b, tq, self.num_heads, self.head_dim);
        let k = Self::split_heads(&k, b, tk, self.num_heads, self.head_dim);
        let v = Self::split_heads(&v, b, tk, self.num_heads, self.head_dim);

        // Scaled dot-product attention
        // logits: [b, h, tq, tk]
        let k_t = k.transpose(2, 3);
        let mut logits = q.matmul(&k_t).div_scalar((self.head_dim as f32).sqrt());
        if let Some(mask) = attn_mask {
            let dims = mask.shape().dims().to_vec();
            // If boolean-like mask matching [b,h,tq,tk], apply masked_fill
            if dims.len() == 4 && dims[0] == b && dims[1] == self.num_heads && dims[2] == tq {
                // Interpret mask > 0.5 as keep; we invert to build masked positions
                let cond: Vec<bool> = mask.data().iter().map(|&v| v < 0.5).collect();
                // Apply masked fill on a flattened view, then reshape back
                let flat_logits = logits.view(vec![(b * self.num_heads * tq * tk) as i32]);
                let filled = flat_logits.masked_fill(&cond, f32::NEG_INFINITY);
                logits = filled.view(vec![b as i32, self.num_heads as i32, tq as i32, tk as i32]);
            } else {
                // Fallback: additive mask
                logits = logits.add_tensor(mask);
            }
        }
        let attn = logits.softmax(3);

        // context: [b, h, tq, d]
        let context = attn.matmul(&v);
        let context = context.permute(vec![0, 2, 1, 3]); // [b, tq, h, d]
        let context = context.contiguous().view(vec![
            b as i32,
            tq as i32,
            (self.num_heads * self.head_dim) as i32,
        ]);

        // Output projection (flatten to 2D, project, then restore 3D)
        let flat = context.view(vec![(b * tq) as i32, self.embed_dim as i32]);
        let out2d = self.out_proj.forward(&flat);
        out2d.view(vec![b as i32, tq as i32, self.embed_dim as i32])
    }

    fn project_qkv(
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        q_proj: &LinearLayer,
        k_proj: &LinearLayer,
        v_proj: &LinearLayer,
    ) -> (Tensor, Tensor, Tensor) {
        let (bq, tq, eq) = Self::triple(query);
        let (bk, tk, ek) = Self::triple(key);
        let (_bv, tv, ev) = Self::triple(value);
        assert!(eq == ek && ek == ev, "Q,K,V embed dims must match");
        let q2d = query.view(vec![(bq * tq) as i32, eq as i32]);
        let k2d = key.view(vec![(bk * tk) as i32, ek as i32]);
        let v2d = value.view(vec![(_bv * tv) as i32, ev as i32]);
        let q = q_proj
            .forward(&q2d)
            .view(vec![bq as i32, tq as i32, eq as i32]);
        let k = k_proj
            .forward(&k2d)
            .view(vec![bk as i32, tk as i32, ek as i32]);
        let v = v_proj
            .forward(&v2d)
            .view(vec![bk as i32, tv as i32, ev as i32]);
        (q, k, v)
    }

    fn split_heads(x: &Tensor, b: usize, t: usize, h: usize, d: usize) -> Tensor {
        x.view(vec![b as i32, t as i32, h as i32, d as i32])
            .permute(vec![0, 2, 1, 3])
    }

impl Tensor

pub fn reshape(&self, new_shape: Vec<i32>) -> Tensor

Reshape the tensor to the specified dimensions

Changes the shape of the tensor while preserving the total number of elements. This operation returns a view when the tensor is contiguous, avoiding data copying. For non-contiguous tensors, data is copied to ensure the reshape is valid.

The reshape operation supports automatic dimension inference using -1, which allows one dimension to be automatically calculated based on the total number of elements and the other specified dimensions.

Arguments

• new_shape - Target shape for the tensor. Use -1 for one dimension to have it automatically inferred from the total size.

Returns

A new tensor with the specified shape containing the same data as the original tensor.

Panics

• If more than one dimension is -1
• If the total number of elements doesn’t match the original tensor
• If any dimension size is 0 or less than -1
• If the inferred dimension size is not a whole number

Examples

use train_station::Tensor;

// Basic reshape
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
let reshaped = tensor.reshape(vec![3, 2]);
assert_eq!(reshaped.shape().dims(), vec![3, 2]);
assert_eq!(reshaped.get(&[0, 0]), 1.0);
assert_eq!(reshaped.get(&[2, 1]), 6.0);

use train_station::Tensor;

// Using -1 for automatic dimension inference
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![4]).unwrap();
let reshaped = tensor.reshape(vec![2, -1]);
assert_eq!(reshaped.shape().dims(), vec![2, 2]);
assert_eq!(reshaped.get(&[0, 0]), 1.0);
assert_eq!(reshaped.get(&[1, 1]), 4.0);

use train_station::Tensor;

// Reshape with gradient tracking
let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
tensor.set_requires_grad(true);

let reshaped = tensor.reshape(vec![4]);
assert!(reshaped.requires_grad());
assert_eq!(reshaped.shape().dims(), vec![4]);

use train_station::Tensor;

// Reshape 3D tensor
let data: Vec<f32> = (0..24).map(|i| i as f32).collect();
let tensor = Tensor::from_slice(&data, vec![2, 3, 4]).unwrap();
let reshaped = tensor.reshape(vec![6, 4]);
assert_eq!(reshaped.shape().dims(), vec![6, 4]);
assert_eq!(reshaped.size(), 24);

Performance

• Contiguous tensors: O(1) time complexity, returns a view
• Non-contiguous tensors: O(n) time complexity with data copying
• Memory usage: No additional allocation for view operations
• Gradient tracking: Preserves gradient requirements and tracking

Automatic Dimension Inference

When using -1 for a dimension, the size is automatically calculated:

use train_station::Tensor;

// For a tensor with 12 elements
let data: Vec<f32> = (0..12).map(|i| i as f32).collect();
let tensor = Tensor::from_slice(&data, vec![3, 4]).unwrap();

let reshaped1 = tensor.reshape(vec![3, -1]);  // Results in shape [3, 4]
let reshaped2 = tensor.reshape(vec![-1, 6]);  // Results in shape [2, 6]
let reshaped3 = tensor.reshape(vec![-1]);     // Results in shape [12]

Examples found in repository

examples/RL_training/ppo_discrete.rs (line 476)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Discrete Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 3usize;
    let total_steps = std::env::var("PPOD_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(3500usize);
    let horizon = 128usize;
    let epochs = 4usize;
    let mini_batch_size = 64usize;
    let gamma = 0.99f32;
    let lam = 0.95f32;
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    let mut actor = Actor::new(state_dim, action_dim, Some(111));
    let mut critic = Critic::new(state_dim, Some(222));
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    let mut env = YardEnv::new(1234);
    let mut rng = SmallRng::new(98765);
    let mut state = env.reset();
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Actor logits and categorical sampling
            let logits = actor.forward(&state); // [1, A]
            let probs = logits.softmax(1); // [1, A]
                                           // sample action from probs (CPU sampling)
            let p = probs.data();
            let (p0, p1, _p2) = (p[0], p[1], p[2]);
            let u = rng.next_f32();
            let a_idx = if u < p0 {
                0
            } else if u < p0 + p1 {
                1
            } else {
                2
            };

            let old_logp = {
                let _ng = NoGradTrack::new();
                let lp = log_prob_actions(&logits, &[a_idx], 1, action_dim);
                lp.data()[0]
            };

            // Step env
            let (next_state, reward, done) = env.step(a_idx);
            episode_return += reward;

            // Critic value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            batch.push(
                state.data(),
                a_idx,
                old_logp,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                out.push(critic.forward(&s2_t).data()[0]);
            }
            out
        };

        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_vec = batch.actions.clone();
        let old_logp_t = Tensor::from_slice(&batch.old_logps, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Views
                let s_mb = states_t
                    .slice_view(start * state_dim, 1, (end - start) * state_dim)
                    .reshape(vec![(end - start) as i32, state_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let a_slice = &actions_vec[start..end];

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward
                let logits_mb = actor.forward(&s_mb); // [B,A]
                let new_logp_mb = log_prob_actions(&logits_mb, a_slice, end - start, action_dim); // [B,1]
                let ratio = ratio_from_logps(&new_logp_mb, &oldlp_mb);
                let ratio_clipped = clamp_ratio(&ratio, clip_eps);
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy bonus from logits (categorical entropy) ≈ -sum p*logp
                let probs_mb = logits_mb.softmax(1);
                let logp_all = probs_mb.add_scalar(1e-8).log();
                let ent = probs_mb
                    .mul_tensor(&logp_all)
                    .sum_dims(&[1], true)
                    .mul_scalar(-1.0)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&ent);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO discrete training finished ===");
    Ok(())
}

examples/RL_training/ppo_continuous.rs (line 490)

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== PPO Continuous Example (YardEnv) ===");

    let state_dim = 3usize;
    let action_dim = 1usize;

    // Hparams
    let total_steps = std::env::var("PPO_STEPS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(4000usize);
    let horizon = 128usize; // rollout length per update
    let epochs = 4usize; // PPO epochs per update
    let mini_batch_size = 64usize; // minibatch from horizon
    let gamma = 0.99f32;
    let lam = 0.95f32; // GAE lambda
    let clip_eps = 0.2f32;
    let vf_coef = 0.5f32;
    let ent_coef = 0.0f32;
    let max_grad_norm = 1.0f32;

    // Models
    let mut actor = Actor::new(state_dim, action_dim, Some(101));
    let mut critic = Critic::new(state_dim, Some(202));

    // Opts
    let mut actor_opt = Adam::with_learning_rate(3e-4);
    for p in actor.parameters() {
        actor_opt.add_parameter(p);
    }
    let mut critic_opt = Adam::with_learning_rate(3e-4);
    for p in critic.parameters() {
        critic_opt.add_parameter(p);
    }

    // Env and RNG
    let mut env = YardEnv::new(42);
    let mut rng = SmallRng::new(999);
    let mut state = env.reset();

    // Metrics
    let mut episode_return = 0.0f32;
    let mut episode = 0usize;
    let mut ema_return: Option<f32> = None;
    let ema_alpha = 0.05f32;
    let mut best_return = f32::NEG_INFINITY;

    let mut t = 0usize;
    while t < total_steps {
        // Collect a rollout
        let mut batch = RolloutBatch::new(horizon, state_dim);
        for _ in 0..horizon {
            // Policy forward (detached sampling to not blow graph; we use stored log_probs)
            let (mean, log_std_row) = actor.forward(&state);
            let mean_v = mean.data()[0];
            let log_std_v = log_std_row.data()[0];
            let std_v = log_std_v.exp();
            let noise = rng.normal();
            let action_v = (mean_v + std_v * noise).clamp(-1.0, 1.0);

            // Build action tensor [1, A] for log_prob calculation with autograd
            let action_t = Tensor::from_slice(&[action_v], vec![1, action_dim]).unwrap();
            let log_prob_t = gaussian_log_prob(&action_t, &mean, &log_std_row);
            let log_prob_v = log_prob_t.data()[0];

            // Step env
            let (next_state, reward, done) = env.step(action_v);
            episode_return += reward;

            // Value
            let value_t = critic.forward(&state);
            let value_v = value_t.data()[0];

            // Push
            batch.push(
                state.data(),
                action_v,
                log_prob_v,
                reward,
                if done { 1.0 } else { 0.0 },
                value_v,
                next_state.data(),
            );

            // Reset
            state = if done {
                let st = env.reset();
                ema_return = Some(match ema_return {
                    None => episode_return,
                    Some(prev) => prev * (1.0 - ema_alpha) + ema_alpha * episode_return,
                });
                if episode_return > best_return {
                    best_return = episode_return;
                }
                println!(
                    "step {:5} | episode {:4} return={:.3} ema={:.3} best={:.3}",
                    t,
                    episode,
                    episode_return,
                    ema_return.unwrap_or(episode_return),
                    best_return
                );
                episode_return = 0.0;
                episode += 1;
                st
            } else {
                next_state
            };

            t += 1;
            if t >= total_steps {
                break;
            }
        }

        // Bootstrap next values for GAE
        let next_values: Vec<f32> = {
            let mut out = Vec::with_capacity(batch.len());
            for i in 0..batch.len() {
                let s2 = &batch.next_states[i * state_dim..(i + 1) * state_dim];
                let s2_t = Tensor::from_slice(s2, vec![1, state_dim]).unwrap();
                let v2 = critic.forward(&s2_t).data()[0];
                out.push(v2);
            }
            out
        };

        // Compute returns and advantages
        let mut returns = vec![0.0f32; batch.len()];
        let mut adv = vec![0.0f32; batch.len()];
        compute_gae(
            &mut returns,
            &mut adv,
            &batch.rewards,
            &batch.dones,
            &batch.values,
            &next_values,
            gamma,
            lam,
        );
        normalize_in_place(&mut adv, 1e-8);

        // Prepare tensors for training
        let states_t = Tensor::from_slice(&batch.states, vec![batch.len(), state_dim]).unwrap();
        let actions_t = Tensor::from_slice(&batch.actions, vec![batch.len(), action_dim]).unwrap();
        let old_logp_t = Tensor::from_slice(&batch.log_probs, vec![batch.len(), 1]).unwrap();
        let returns_t = Tensor::from_slice(&returns, vec![batch.len(), 1]).unwrap();
        let adv_t = Tensor::from_slice(&adv, vec![batch.len(), 1]).unwrap();

        // PPO epochs over the rollout
        let num_minibatches = batch.len().div_ceil(mini_batch_size);
        for e in 0..epochs {
            for mb in 0..num_minibatches {
                let start = mb * mini_batch_size;
                let end = (start + mini_batch_size).min(batch.len());
                if start >= end {
                    break;
                }

                // Slice views
                let s_mb = states_t.slice_view(start * state_dim, 1, (end - start) * state_dim);
                let s_mb = s_mb.reshape(vec![(end - start) as i32, state_dim as i32]);
                let a_mb = actions_t
                    .slice_view(start * action_dim, 1, (end - start) * action_dim)
                    .reshape(vec![(end - start) as i32, action_dim as i32]);
                let oldlp_mb = old_logp_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let ret_mb = returns_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);
                let adv_mb = adv_t
                    .slice_view(start, 1, end - start)
                    .reshape(vec![(end - start) as i32, 1]);

                // Zero grads
                {
                    let mut ps = actor.parameters();
                    actor_opt.zero_grad(&mut ps);
                }
                {
                    let mut ps = critic.parameters();
                    critic_opt.zero_grad(&mut ps);
                }

                // Forward actor and critic
                let (mean_mb, log_std_row) = actor.forward(&s_mb);
                let logp_mb = gaussian_log_prob(&a_mb, &mean_mb, &log_std_row);
                let ratio = logp_mb.sub_tensor(&oldlp_mb).exp(); // exp(new-old)
                let clip_low =
                    Tensor::from_slice(&vec![1.0 - clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                let clip_high =
                    Tensor::from_slice(&vec![1.0 + clip_eps; end - start], vec![end - start, 1])
                        .unwrap();
                // ratio_clipped = min(max(ratio, low), high) using ReLU identities
                let ratio_ge_low = ratio.sub_tensor(&clip_low).relu().add_tensor(&clip_low);
                let ratio_clipped =
                    clip_high.sub_tensor(&ratio_ge_low.sub_tensor(&clip_high).relu());
                let pg1 = ratio.mul_tensor(&adv_mb);
                let pg2 = ratio_clipped.mul_tensor(&adv_mb);
                // min(pg1, pg2) = pg2 - relu(pg2 - pg1)
                let actor_min = pg2.sub_tensor(&pg2.sub_tensor(&pg1).relu());
                let actor_loss = actor_min.mul_scalar(-1.0).mean();

                let v_pred = critic.forward(&s_mb);
                let v_loss = v_pred
                    .sub_tensor(&ret_mb)
                    .pow_scalar(2.0)
                    .mean()
                    .mul_scalar(vf_coef);

                // Entropy (approx Gaussian entropy per action)
                let entropy = log_std_row
                    .add_scalar(0.5 * (2.0 * std::f32::consts::PI * std::f32::consts::E).ln())
                    .sum_dims(&[1], true)
                    .mean()
                    .mul_scalar(ent_coef);

                let mut loss = actor_loss.add_tensor(&v_loss).sub_tensor(&entropy);
                loss.backward(None);

                // Step actor
                {
                    let params = actor.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        actor_opt.step(&mut with_grads);
                        actor_opt.zero_grad(&mut with_grads);
                    }
                }

                // Step critic
                {
                    let params = critic.parameters();
                    let mut with_grads: Vec<&mut Tensor> = Vec::new();
                    for p in params {
                        if p.grad_owned().is_some() {
                            with_grads.push(p);
                        }
                    }
                    if !with_grads.is_empty() {
                        let _ = grad_global_norm(&mut with_grads);
                        clip_gradients(&mut with_grads, max_grad_norm, 1e-6);
                        critic_opt.step(&mut with_grads);
                        critic_opt.zero_grad(&mut with_grads);
                    }
                }

                // Occasionally log
                if e == 0 && mb == 0 {
                    println!(
                        "update@t={} | actor_loss={:.4} v_loss={:.4}",
                        t,
                        actor_loss.value(),
                        v_loss.value()
                    );
                }

                clear_all_graphs_known();
            }
        }
    }

    println!("=== PPO training finished ===");
    Ok(())
}

impl Tensor

pub fn split(&self, split_size: usize, dim: usize) -> Vec<Tensor>

Split tensor into chunks of equal size along specified dimension

Divides the tensor into multiple smaller tensors along the specified dimension, where each chunk (except possibly the last) has the same size. The last chunk may be smaller if the dimension size is not evenly divisible by the split size.

This operation returns a vector of tensors, where each tensor is a view or copy of a portion of the original tensor. The first chunk is returned as a view when possible (zero-copy), while subsequent chunks may require data copying for non-zero base offsets.

Arguments

• split_size - Size of each chunk along the specified dimension (must be > 0)
• dim - Dimension along which to split the tensor (must be < tensor rank)

Returns

A vector of tensors, each representing a chunk of the original tensor. The number of chunks depends on the dimension size and split size.

Panics

• If tensor rank is 0 (scalar tensors cannot be split)
• If dim is out of bounds for the tensor rank
• If split_size is 0

Examples

use train_station::Tensor;

// Split 2D tensor into equal chunks along dimension 1
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
let parts = tensor.split(1, 1);
assert_eq!(parts.len(), 3);
assert_eq!(parts[0].shape().dims(), vec![2, 1]);
assert_eq!(parts[1].shape().dims(), vec![2, 1]);
assert_eq!(parts[2].shape().dims(), vec![2, 1]);
assert_eq!(parts[0].get(&[0, 0]), 1.0);
assert_eq!(parts[1].get(&[0, 0]), 2.0);
assert_eq!(parts[2].get(&[1, 0]), 6.0);

use train_station::Tensor;

// Split with uneven division (last chunk smaller)
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![1, 5]).unwrap();
let parts = tensor.split(2, 1);
assert_eq!(parts.len(), 3);
assert_eq!(parts[0].shape().dims(), vec![1, 2]);
assert_eq!(parts[1].shape().dims(), vec![1, 2]);
assert_eq!(parts[2].shape().dims(), vec![1, 1]); // Last chunk smaller

use train_station::Tensor;

// Split with gradient tracking
let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
tensor.set_requires_grad(true);

let parts = tensor.split(1, 1);
assert_eq!(parts.len(), 2);
assert!(parts[0].requires_grad());
assert!(parts[1].requires_grad());

use train_station::Tensor;

// Split 1D tensor
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6]).unwrap();
let parts = tensor.split(2, 0);
assert_eq!(parts.len(), 3);
assert_eq!(parts[0].shape().dims(), vec![2]);
assert_eq!(parts[1].shape().dims(), vec![2]);
assert_eq!(parts[2].shape().dims(), vec![2]);

Performance

• First Chunk: O(1) - Returns a view when possible (zero-copy)
• Subsequent Chunks: O(n) - May require data copying for non-zero offsets
• Memory Usage: Minimal allocation for view operations, copying for non-zero offsets
• Gradient Tracking: Each chunk preserves gradient requirements and tracking

Relationship to Other Operations

This operation is related to other tensor transformations:

• split_with_sizes() - More general version with explicit chunk sizes
• cat() - Inverse operation that concatenates tensors back together
• chunk() - Alternative splitting operation with different semantics

Memory Layout

The first chunk maintains the same underlying data as a view when the base offset is zero. Subsequent chunks may require data copying to handle non-zero base offsets, ensuring proper memory layout.

pub fn split_with_sizes(&self, split_sizes: &[usize], dim: usize) -> Vec<Tensor>

Split tensor into chunks with explicit sizes along specified dimension

Divides the tensor into multiple smaller tensors along the specified dimension according to the provided size specifications. Each chunk has the exact size specified in the split_sizes array, and the sum of all sizes must equal the size of the specified dimension.

This operation provides precise control over the size of each resulting chunk, unlike split() which creates equal-sized chunks. The first chunk is returned as a view when possible (zero-copy), while subsequent chunks may require data copying for non-zero base offsets.

Arguments

• split_sizes - Array specifying the size of each chunk along the dimension
• dim - Dimension along which to split the tensor (must be < tensor rank)

Returns

A vector of tensors, each representing a chunk of the original tensor with the specified size. The number of chunks equals the length of split_sizes.

Panics

• If tensor rank is 0 (scalar tensors cannot be split)
• If dim is out of bounds for the tensor rank
• If sum of split_sizes does not equal the size of the specified dimension

Examples

use train_station::Tensor;

// Split with explicit sizes
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![1, 5]).unwrap();
let parts = tensor.split_with_sizes(&[2, 3], 1);
assert_eq!(parts.len(), 2);
assert_eq!(parts[0].shape().dims(), vec![1, 2]);
assert_eq!(parts[1].shape().dims(), vec![1, 3]);
assert_eq!(parts[0].get(&[0, 0]), 1.0);
assert_eq!(parts[0].get(&[0, 1]), 2.0);
assert_eq!(parts[1].get(&[0, 0]), 3.0);

use train_station::Tensor;

// Split 2D tensor with different chunk sizes
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
let parts = tensor.split_with_sizes(&[1, 2], 1);
assert_eq!(parts.len(), 2);
assert_eq!(parts[0].shape().dims(), vec![2, 1]);
assert_eq!(parts[1].shape().dims(), vec![2, 2]);

use train_station::Tensor;

// Split with gradient tracking
let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
tensor.set_requires_grad(true);

let parts = tensor.split_with_sizes(&[1, 1], 1);
assert_eq!(parts.len(), 2);
assert!(parts[0].requires_grad());
assert!(parts[1].requires_grad());

use train_station::Tensor;

// Split 1D tensor with explicit sizes
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6]).unwrap();
let parts = tensor.split_with_sizes(&[2, 2, 2], 0);
assert_eq!(parts.len(), 3);
assert_eq!(parts[0].shape().dims(), vec![2]);
assert_eq!(parts[1].shape().dims(), vec![2]);
assert_eq!(parts[2].shape().dims(), vec![2]);

Performance

• First Chunk: O(1) - Returns a view when possible (zero-copy)
• Subsequent Chunks: O(n) - May require data copying for non-zero offsets
• Memory Usage: Minimal allocation for view operations, copying for non-zero offsets
• Gradient Tracking: Each chunk preserves gradient requirements and tracking

Relationship to Other Operations

This operation is related to other tensor transformations:

• split() - Simplified version with equal-sized chunks
• cat() - Inverse operation that concatenates tensors back together
• chunk() - Alternative splitting operation with different semantics

Memory Layout

The first chunk maintains the same underlying data as a view when the base offset is zero. Subsequent chunks may require data copying to handle non-zero base offsets, ensuring proper memory layout. Zero-sized chunks are handled by creating empty tensors with appropriate shapes.

impl Tensor

pub fn squeeze(&self, dim: Option<usize>) -> Tensor

Remove dimensions of size 1 from the tensor

Removes singleton dimensions (dimensions with size 1) from the tensor, reducing its rank while preserving the total number of elements. This operation is useful for cleaning up tensor shapes and preparing data for operations that expect specific dimensionality.

The squeeze operation can remove either all size-1 dimensions or a specific dimension if it has size 1. When all dimensions are size 1, the result is a scalar tensor with shape [1] rather than an empty tensor to maintain mathematical consistency.

Arguments

• dim - Optional specific dimension to squeeze. If None, all size-1 dimensions are removed. If Some(d), only dimension d is removed if it has size 1.

Returns

A new tensor with size-1 dimensions removed. The total number of elements remains unchanged.

Panics

• If dim is specified but out of bounds for the tensor rank
• If dim is specified but the dimension does not have size 1

Examples

use train_station::Tensor;

// Squeeze all size-1 dimensions
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3, 1]).unwrap();
let squeezed = tensor.squeeze(None);
assert_eq!(squeezed.shape().dims(), vec![3]);
assert_eq!(squeezed.get(&[0]), 1.0);
assert_eq!(squeezed.get(&[1]), 2.0);
assert_eq!(squeezed.get(&[2]), 3.0);

use train_station::Tensor;

// Squeeze specific dimension
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3, 1]).unwrap();
let squeezed = tensor.squeeze(Some(0));
assert_eq!(squeezed.shape().dims(), vec![3, 1]);
assert_eq!(squeezed.get(&[0, 0]), 1.0);
assert_eq!(squeezed.get(&[1, 0]), 2.0);
assert_eq!(squeezed.get(&[2, 0]), 3.0);

use train_station::Tensor;

// Squeeze preserves data integrity
let data = vec![1.0, 2.0, 3.0, 4.0];
let tensor = Tensor::from_slice(&data, vec![1, 2, 1, 2]).unwrap();
let squeezed = tensor.squeeze(None);
assert_eq!(squeezed.shape().dims(), vec![2, 2]);
assert_eq!(squeezed.size(), 4);
assert_eq!(squeezed.get(&[0, 0]), data[0]);
assert_eq!(squeezed.get(&[0, 1]), data[1]);
assert_eq!(squeezed.get(&[1, 0]), data[2]);
assert_eq!(squeezed.get(&[1, 1]), data[3]);

use train_station::Tensor;

// Handle edge case: all dimensions are size 1
let tensor = Tensor::from_slice(&[5.0], vec![1, 1, 1]).unwrap();
let squeezed = tensor.squeeze(None);
assert_eq!(squeezed.shape().dims(), vec![1]); // Not empty!
assert_eq!(squeezed.get(&[0]), 5.0);

use train_station::Tensor;

// Squeeze with gradient tracking
let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![1, 3, 1]).unwrap();
tensor.set_requires_grad(true);

let squeezed = tensor.squeeze(None);
assert!(squeezed.requires_grad());
assert_eq!(squeezed.shape().dims(), vec![3]);

use train_station::Tensor;

// Squeeze and unsqueeze roundtrip
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let unsqueezed = tensor.unsqueeze(0);
assert_eq!(unsqueezed.shape().dims(), vec![1, 3]);

let squeezed = unsqueezed.squeeze(Some(0));
assert_eq!(squeezed.shape().dims(), vec![3]);
assert_eq!(squeezed.get(&[0]), 1.0);
assert_eq!(squeezed.get(&[2]), 3.0);

Performance

• Time Complexity: O(1) - Returns a view through reshape operation
• Memory Usage: No additional memory allocation (view operation)
• Gradient Tracking: Preserves gradient requirements and tracking
• Shape Transformation: Reduces tensor rank by removing singleton dimensions

Relationship to Other Operations

This operation is related to other tensor transformations:

• unsqueeze() - Inverse operation that adds size-1 dimensions
• reshape() - More general shape transformation operation
• flatten() - Reduces tensor to 1D by combining all dimensions

Memory Layout

The squeezed tensor maintains the same underlying data as the original tensor through the reshape operation. This ensures zero-copy behavior when the tensor is contiguous, with only the shape metadata being modified to reflect the reduced dimensionality.

Edge Cases

• All size-1 dimensions: Returns a tensor with shape [1] rather than an empty tensor to maintain mathematical consistency
• No size-1 dimensions: Returns a tensor with the same shape as the input
• Mixed dimensions: Only removes dimensions with size 1, preserving others

impl Tensor

pub fn stack(tensors: &[Tensor], dim: usize) -> Tensor

Stack a list of tensors along a new dimension

Combines multiple tensors by adding a new dimension at the specified position. All input tensors must have identical shapes, and the output tensor will have a new dimension of size equal to the number of input tensors. This operation is similar to PyTorch’s torch.stack function.

The stacking operation creates a new axis in the output tensor, unlike concatenation which operates along existing dimensions. This makes stacking useful for creating batch dimensions, combining feature maps, and implementing operations that require adding new tensor axes.

Arguments

• tensors - Array of tensors to stack. All tensors must have identical shapes.
• dim - Index of the new axis in the output shape (0 <= dim <= rank)

Returns

A new tensor with the stacked data. The output shape is the input shape with a new dimension of size tensors.len() inserted at position dim.

Panics

• If the tensor array is empty
• If any tensor has a different shape than the first tensor
• If dim is out of bounds (dim > rank of input tensors)

Examples

use train_station::Tensor;

// Stack two 1D tensors along dimension 0
let a = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let b = Tensor::from_slice(&[4.0, 5.0, 6.0], vec![3]).unwrap();
let stacked = Tensor::stack(&[a, b], 0);
assert_eq!(stacked.shape().dims(), vec![2, 3]);
assert_eq!(stacked.get(&[0, 0]), 1.0);
assert_eq!(stacked.get(&[1, 2]), 6.0);

use train_station::Tensor;

// Stack multiple 2D tensors along dimension 1
let a = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let b = Tensor::from_slice(&[5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();
let c = Tensor::from_slice(&[9.0, 10.0, 11.0, 12.0], vec![2, 2]).unwrap();
let stacked = Tensor::stack(&[a, b, c], 1);
assert_eq!(stacked.shape().dims(), vec![2, 3, 2]);
assert_eq!(stacked.get(&[0, 0, 0]), 1.0);
assert_eq!(stacked.get(&[1, 2, 1]), 12.0);

use train_station::Tensor;

// Stack with gradient tracking
let mut a = Tensor::from_slice(&[1.0, 2.0], vec![2]).unwrap();
let mut b = Tensor::from_slice(&[3.0, 4.0], vec![2]).unwrap();
a.set_requires_grad(true);
b.set_requires_grad(true);

let stacked = Tensor::stack(&[a, b], 0);
assert!(stacked.requires_grad());
assert_eq!(stacked.shape().dims(), vec![2, 2]);

use train_station::Tensor;

// Stack 3D tensors along the last dimension
let data1: Vec<f32> = (0..8).map(|i| i as f32).collect();
let data2: Vec<f32> = (8..16).map(|i| i as f32).collect();
let a = Tensor::from_slice(&data1, vec![2, 2, 2]).unwrap();
let b = Tensor::from_slice(&data2, vec![2, 2, 2]).unwrap();
let stacked = Tensor::stack(&[a, b], 3);
assert_eq!(stacked.shape().dims(), vec![2, 2, 2, 2]);
assert_eq!(stacked.get(&[0, 0, 0, 0]), 0.0);
assert_eq!(stacked.get(&[1, 1, 1, 1]), 15.0);

Performance

• Time Complexity: O(n) where n is the total number of elements
• Memory Usage: Allocates new contiguous tensor for output
• SIMD Optimization: Uses AVX2 acceleration for large block copies
• Block-wise Copying: Optimized copying strategy for better cache performance
• Gradient Tracking: Preserves gradient requirements and tracking

Relationship to Other Operations

This operation is related to other tensor transformations:

• cat() - Concatenates tensors along existing dimensions
• unsqueeze() - Adds a single dimension of size 1
• reshape() - Changes tensor shape without adding dimensions

Memory Layout

The output tensor is always contiguous, with elements arranged so that the stacked dimension is the fastest-changing index. This ensures optimal performance for subsequent operations and maintains compatibility with SIMD optimizations.

Gradient Computation

During backward passes, gradients are split along the stacked dimension and distributed back to the original input tensors. This is implemented using the same gradient function as concatenation, treating the stack operation as concatenation along a new axis.

impl Tensor

pub fn transpose(&self, dim0: usize, dim1: usize) -> Tensor

Transpose two dimensions of the tensor

Swaps two specified dimensions of the tensor, modifying the shape and memory access pattern. When possible, this operation returns a zero-copy view using stride manipulation. For complex cases or non-contiguous tensors, data is copied to ensure correct transposition.

The transpose operation is its own inverse - applying transpose twice with the same dimensions returns the original tensor.

Arguments

• dim0 - First dimension to swap (must be < tensor rank)
• dim1 - Second dimension to swap (must be < tensor rank)

Returns

A new tensor with the specified dimensions transposed. The total number of elements remains unchanged.

Panics

• If dim0 is out of bounds for the tensor rank
• If dim1 is out of bounds for the tensor rank

Examples

use train_station::Tensor;

// Basic 2D transpose
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
let transposed = tensor.transpose(0, 1);
assert_eq!(transposed.shape().dims(), vec![3, 2]);
assert_eq!(transposed.get(&[0, 0]), 1.0);
assert_eq!(transposed.get(&[0, 1]), 4.0);
assert_eq!(transposed.get(&[1, 0]), 2.0);
assert_eq!(transposed.get(&[1, 1]), 5.0);
assert_eq!(transposed.get(&[2, 0]), 3.0);
assert_eq!(transposed.get(&[2, 1]), 6.0);

use train_station::Tensor;

// 3D tensor transpose
let data: Vec<f32> = (0..24).map(|i| i as f32).collect();
let tensor = Tensor::from_slice(&data, vec![2, 3, 4]).unwrap();
let transposed = tensor.transpose(0, 1);
assert_eq!(transposed.shape().dims(), vec![3, 2, 4]);

use train_station::Tensor;

// Transpose with gradient tracking
let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
tensor.set_requires_grad(true);

let transposed = tensor.transpose(0, 1);
assert!(transposed.requires_grad());
assert_eq!(transposed.shape().dims(), vec![2, 2]);

use train_station::Tensor;

// Transpose same dimension (no change)
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let result = tensor.transpose(1, 1);
assert_eq!(result.shape().dims(), tensor.shape().dims());
assert_eq!(result.get(&[0, 0]), tensor.get(&[0, 0]));

use train_station::Tensor;

// Transpose is its own inverse
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let transposed = tensor.transpose(0, 1);
let double_transposed = transposed.transpose(0, 1);
assert_eq!(double_transposed.shape().dims(), tensor.shape().dims());
assert_eq!(double_transposed.get(&[0, 0]), tensor.get(&[0, 0]));

Performance

• Contiguous tensors: O(1) time complexity, returns a view
• Non-contiguous tensors: O(n) time complexity with data copying
• Memory usage: No additional allocation for view operations
• Gradient tracking: Preserves gradient requirements and tracking

Relationship to Other Operations

This operation is related to other tensor transformations:

• t() - Convenience method for matrix transpose (last two dimensions)
• permute() - More general dimension reordering operation
• reshape() - Changes shape without changing dimension order

Memory Layout

For contiguous tensors, transpose returns a view with modified strides, making the tensor non-contiguous. For non-contiguous tensors or complex cases, data is copied to ensure correct transposition.

Examples found in repository

examples/neural_networks/multi_head_attention.rs (line 91)

    pub fn forward(
        &self,
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        attn_mask: Option<&Tensor>,
    ) -> Tensor {
        let qkv = Self::project_qkv(query, key, value, &self.q_proj, &self.k_proj, &self.v_proj);
        let (q, k, v) = qkv;

        // Split heads: [b, t, e] -> [b, h, t, d]
        let (b, tq, _e) = Self::triple(query);
        let (_b2, tk, _e2) = Self::triple(key);
        let q = Self::split_heads(&q, b, tq, self.num_heads, self.head_dim);
        let k = Self::split_heads(&k, b, tk, self.num_heads, self.head_dim);
        let v = Self::split_heads(&v, b, tk, self.num_heads, self.head_dim);

        // Scaled dot-product attention
        // logits: [b, h, tq, tk]
        let k_t = k.transpose(2, 3);
        let mut logits = q.matmul(&k_t).div_scalar((self.head_dim as f32).sqrt());
        if let Some(mask) = attn_mask {
            let dims = mask.shape().dims().to_vec();
            // If boolean-like mask matching [b,h,tq,tk], apply masked_fill
            if dims.len() == 4 && dims[0] == b && dims[1] == self.num_heads && dims[2] == tq {
                // Interpret mask > 0.5 as keep; we invert to build masked positions
                let cond: Vec<bool> = mask.data().iter().map(|&v| v < 0.5).collect();
                // Apply masked fill on a flattened view, then reshape back
                let flat_logits = logits.view(vec![(b * self.num_heads * tq * tk) as i32]);
                let filled = flat_logits.masked_fill(&cond, f32::NEG_INFINITY);
                logits = filled.view(vec![b as i32, self.num_heads as i32, tq as i32, tk as i32]);
            } else {
                // Fallback: additive mask
                logits = logits.add_tensor(mask);
            }
        }
        let attn = logits.softmax(3);

        // context: [b, h, tq, d]
        let context = attn.matmul(&v);
        let context = context.permute(vec![0, 2, 1, 3]); // [b, tq, h, d]
        let context = context.contiguous().view(vec![
            b as i32,
            tq as i32,
            (self.num_heads * self.head_dim) as i32,
        ]);

        // Output projection (flatten to 2D, project, then restore 3D)
        let flat = context.view(vec![(b * tq) as i32, self.embed_dim as i32]);
        let out2d = self.out_proj.forward(&flat);
        out2d.view(vec![b as i32, tq as i32, self.embed_dim as i32])
    }

pub fn t(&self) -> Tensor

Matrix transpose (transpose last two dimensions)

Convenience method for the common case of matrix transposition. For 2D tensors, this performs a standard matrix transpose. For higher-dimensional tensors, this transposes the last two dimensions, treating the tensor as a batch of matrices.

This method is equivalent to transpose(rank-2, rank-1) where rank is the number of dimensions in the tensor.

Returns

A new tensor with the last two dimensions transposed

Panics

• If the tensor has less than 2 dimensions

Examples

use train_station::Tensor;

// 2D matrix transpose
let matrix = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let transposed = matrix.t();
assert_eq!(transposed.shape().dims(), vec![2, 2]);
assert_eq!(transposed.get(&[0, 0]), 1.0);
assert_eq!(transposed.get(&[0, 1]), 3.0);
assert_eq!(transposed.get(&[1, 0]), 2.0);
assert_eq!(transposed.get(&[1, 1]), 4.0);

use train_station::Tensor;

// 3D tensor (batch of matrices)
let data: Vec<f32> = (0..12).map(|i| i as f32).collect();
let tensor = Tensor::from_slice(&data, vec![2, 2, 3]).unwrap();
let transposed = tensor.t();
assert_eq!(transposed.shape().dims(), vec![2, 3, 2]);

use train_station::Tensor;

// Matrix transpose with gradient tracking
let mut matrix = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
matrix.set_requires_grad(true);

let transposed = matrix.t();
assert!(transposed.requires_grad());
assert_eq!(transposed.shape().dims(), vec![2, 2]);

Performance

• Time Complexity: Same as transpose() - O(1) for views, O(n) for copies
• Memory Usage: Same as transpose() - no allocation for views
• Gradient Tracking: Preserves gradient requirements and tracking

Relationship to Other Operations

This operation is equivalent to:

use train_station::Tensor;

let tensor = Tensor::new(vec![2, 3, 4]);
let rank = tensor.shape().rank();
let transposed1 = tensor.t();
let transposed2 = tensor.transpose(rank - 2, rank - 1);
// transposed1 and transposed2 are identical

impl Tensor

pub fn unsqueeze(&self, dim: usize) -> Tensor

Add a dimension of size 1 at the specified position

Inserts a new dimension of size 1 at the specified position in the tensor’s shape, increasing the rank by 1 while preserving the total number of elements. This operation is useful for preparing tensors for broadcasting, creating batch dimensions, and adapting tensor shapes for specific neural network operations.

The unsqueeze operation is the inverse of squeeze() - unsqueezing a dimension and then squeezing it at the same position returns the original tensor.

Arguments

• dim - Position to insert the new dimension (0 <= dim <= rank)

Returns

A new tensor with an additional dimension of size 1 at the specified position. The total number of elements remains unchanged.

Panics

• If dim is out of bounds (dim > rank of the tensor)

Examples

use train_station::Tensor;

// Add dimension at the beginning
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let unsqueezed = tensor.unsqueeze(0);
assert_eq!(unsqueezed.shape().dims(), vec![1, 3]);
assert_eq!(unsqueezed.get(&[0, 0]), 1.0);
assert_eq!(unsqueezed.get(&[0, 1]), 2.0);
assert_eq!(unsqueezed.get(&[0, 2]), 3.0);

use train_station::Tensor;

// Add dimension at the end
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let unsqueezed = tensor.unsqueeze(1);
assert_eq!(unsqueezed.shape().dims(), vec![3, 1]);
assert_eq!(unsqueezed.get(&[0, 0]), 1.0);
assert_eq!(unsqueezed.get(&[1, 0]), 2.0);
assert_eq!(unsqueezed.get(&[2, 0]), 3.0);

use train_station::Tensor;

// Add dimension in the middle of 2D tensor
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let unsqueezed = tensor.unsqueeze(1);
assert_eq!(unsqueezed.shape().dims(), vec![2, 1, 2]);
assert_eq!(unsqueezed.get(&[0, 0, 0]), 1.0);
assert_eq!(unsqueezed.get(&[0, 0, 1]), 2.0);
assert_eq!(unsqueezed.get(&[1, 0, 0]), 3.0);
assert_eq!(unsqueezed.get(&[1, 0, 1]), 4.0);

use train_station::Tensor;

// Unsqueeze preserves data integrity
let data = vec![1.0, 2.0, 3.0, 4.0];
let tensor = Tensor::from_slice(&data, vec![4]).unwrap();
let unsqueezed = tensor.unsqueeze(0);
assert_eq!(unsqueezed.shape().dims(), vec![1, 4]);
assert_eq!(unsqueezed.size(), 4);
for (i, &d) in data.iter().enumerate() {
    assert_eq!(unsqueezed.get(&[0, i]), d);
}

use train_station::Tensor;

// Unsqueeze with gradient tracking
let mut tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
tensor.set_requires_grad(true);

let unsqueezed = tensor.unsqueeze(0);
assert!(unsqueezed.requires_grad());
assert_eq!(unsqueezed.shape().dims(), vec![1, 3]);

use train_station::Tensor;

// Unsqueeze and squeeze roundtrip
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let unsqueezed = tensor.unsqueeze(0);
assert_eq!(unsqueezed.shape().dims(), vec![1, 3]);

let squeezed = unsqueezed.squeeze(Some(0));
assert_eq!(squeezed.shape().dims(), vec![3]);
assert_eq!(squeezed.get(&[0]), 1.0);
assert_eq!(squeezed.get(&[2]), 3.0);

use train_station::Tensor;

// Multiple unsqueeze operations
let tensor = Tensor::from_slice(&[42.0], vec![1]).unwrap();
let unsqueezed1 = tensor.unsqueeze(0);
assert_eq!(unsqueezed1.shape().dims(), vec![1, 1]);

let unsqueezed2 = unsqueezed1.unsqueeze(0);
assert_eq!(unsqueezed2.shape().dims(), vec![1, 1, 1]);
assert_eq!(unsqueezed2.get(&[0, 0, 0]), 42.0);

Performance

• Time Complexity: O(1) - Returns a view through reshape operation
• Memory Usage: No additional memory allocation (view operation)
• Gradient Tracking: Preserves gradient requirements and tracking
• Shape Transformation: Increases tensor rank by adding singleton dimensions

Relationship to Other Operations

This operation is related to other tensor transformations:

• squeeze() - Inverse operation that removes size-1 dimensions
• reshape() - More general shape transformation operation
• expand() - Broadcasts dimensions to larger sizes

Memory Layout

The unsqueezed tensor maintains the same underlying data as the original tensor through the reshape operation. This ensures zero-copy behavior when the tensor is contiguous, with only the shape metadata being modified to reflect the increased dimensionality.

Broadcasting Applications

Unsqueeze is commonly used for broadcasting operations:

use train_station::Tensor;

// Prepare for broadcasting: [3] -> [1, 3] for row-wise operations
let vector = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let row_vector = vector.unsqueeze(0); // Shape: [1, 3]

// Prepare for broadcasting: [3] -> [3, 1] for column-wise operations
let column_vector = vector.unsqueeze(1); // Shape: [3, 1]

Neural Network Applications

Unsqueeze is essential for neural network operations:

use train_station::Tensor;

// Single sample -> batch dimension for neural network input
let sample = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let batch = sample.unsqueeze(0); // Shape: [1, 3] for batch processing

// Add channel dimension for convolutional operations
let feature_map = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let with_channels = feature_map.unsqueeze(0); // Shape: [1, 2, 2] for conv layers

Trait Implementations

impl Add<&Tensor> for Tensor

fn add(self, other: &Tensor) -> Tensor

Adds a tensor and a tensor reference element-wise

Returns

A new tensor containing the element-wise sum

type Output = Tensor

The resulting type after applying the + operator.

impl Add<&Tensor> for f32

fn add(self, tensor: &Tensor) -> Tensor

Adds a scalar to each element of the tensor (reference version)

Returns

A new tensor with the scalar added to each element

type Output = Tensor

The resulting type after applying the + operator.

impl Add<Tensor> for &Tensor

fn add(self, other: Tensor) -> Tensor

Adds a tensor reference and a tensor element-wise

Returns

A new tensor containing the element-wise sum

type Output = Tensor

The resulting type after applying the + operator.

impl Add<Tensor> for f32

Scalar-tensor addition operator implementations

Provides addition operations between scalars and tensors. All implementations delegate to the underlying add_scalar method.

fn add(self, tensor: Tensor) -> Tensor

Adds a scalar to each element of the tensor

Returns

A new tensor with the scalar added to each element

type Output = Tensor

The resulting type after applying the + operator.

impl Add<f32> for &Tensor

fn add(self, scalar: f32) -> Tensor

Adds a scalar to each element of the tensor (reference version)

Returns

A new tensor with the scalar added to each element

type Output = Tensor

The resulting type after applying the + operator.

impl Add<f32> for Tensor

Tensor-scalar addition operator implementations

Provides addition operations between tensors and scalars. All implementations delegate to the underlying add_scalar method.

fn add(self, scalar: f32) -> Tensor

Adds a scalar to each element of the tensor

Returns

A new tensor with the scalar added to each element

type Output = Tensor

The resulting type after applying the + operator.

impl Add for &Tensor

fn add(self, other: &Tensor) -> Tensor

Adds two tensors element-wise (reference version)

Returns

A new tensor containing the element-wise sum

type Output = Tensor

The resulting type after applying the + operator.

impl Add for Tensor

Tensor addition operator implementations

Provides addition operations between tensors with various reference combinations. All implementations delegate to the underlying add_tensor method for optimal performance.

fn add(self, other: Tensor) -> Tensor

Adds two tensors element-wise

Returns

A new tensor containing the element-wise sum

type Output = Tensor

The resulting type after applying the + operator.

impl AddAssign<&Tensor> for Tensor

fn add_assign(&mut self, other: &Tensor)

Adds another tensor reference to this tensor in-place

impl AddAssign<f32> for Tensor

Tensor-scalar addition assignment operator implementations

Provides in-place addition operations between tensors and scalars.

fn add_assign(&mut self, scalar: f32)

Adds a scalar to each element of this tensor in-place

impl AddAssign for Tensor

Tensor addition assignment operator implementations

Provides in-place addition operations between tensors. All implementations delegate to the underlying add_tensor method.

fn add_assign(&mut self, other: Tensor)

Adds another tensor to this tensor in-place

impl Clone for Tensor

Clone implementation for Tensor

Creates a deep copy of the tensor data but resets gradtrack state (new tensor won’t track gradients unless explicitly set)

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Tensor

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Div<&Tensor> for Tensor

fn div(self, other: &Tensor) -> Tensor

Divides a tensor by a tensor reference element-wise

Returns

A new tensor containing the element-wise quotient

type Output = Tensor

The resulting type after applying the / operator.

impl Div<&Tensor> for f32

fn div(self, tensor: &Tensor) -> Tensor

Divides a scalar by each element of the tensor (reference version)

Returns

A new tensor with the scalar divided by each element

type Output = Tensor

The resulting type after applying the / operator.

impl Div<Tensor> for &Tensor

fn div(self, other: Tensor) -> Tensor

Divides a tensor reference by a tensor element-wise

Returns

A new tensor containing the element-wise quotient

type Output = Tensor

The resulting type after applying the / operator.

impl Div<Tensor> for f32

Scalar-tensor division operator implementations

Provides division operations between scalars and tensors. Computes scalar / tensor by computing the reciprocal of the tensor and multiplying by the scalar.

fn div(self, tensor: Tensor) -> Tensor

Divides a scalar by each element of the tensor

Returns

A new tensor with the scalar divided by each element

type Output = Tensor

The resulting type after applying the / operator.

impl Div<f32> for &Tensor

fn div(self, scalar: f32) -> Tensor

Divides each element of the tensor by a scalar (reference version)

Returns

A new tensor with each element divided by the scalar

type Output = Tensor

The resulting type after applying the / operator.

impl Div<f32> for Tensor

Tensor-scalar division operator implementations

Provides division operations between tensors and scalars. All implementations delegate to the underlying div_scalar method.

fn div(self, scalar: f32) -> Tensor

Divides each element of the tensor by a scalar

Returns

A new tensor with each element divided by the scalar

type Output = Tensor

The resulting type after applying the / operator.

impl Div for &Tensor

fn div(self, other: &Tensor) -> Tensor

Divides two tensors element-wise (reference version)

Returns

A new tensor containing the element-wise quotient

type Output = Tensor

The resulting type after applying the / operator.

impl Div for Tensor

Tensor division operator implementations

Provides element-wise division operations between tensors with various reference combinations. All implementations delegate to the underlying div_tensor method for optimal performance.

fn div(self, other: Tensor) -> Tensor

Divides two tensors element-wise

Returns

A new tensor containing the element-wise quotient

type Output = Tensor

The resulting type after applying the / operator.

impl DivAssign<&Tensor> for Tensor

fn div_assign(&mut self, other: &Tensor)

Divides this tensor by another tensor reference in-place

impl DivAssign<f32> for Tensor

Tensor-scalar division assignment operator implementations

Provides in-place division operations between tensors and scalars.

fn div_assign(&mut self, scalar: f32)

Divides each element of this tensor by a scalar in-place

impl DivAssign for Tensor

Tensor division assignment operator implementations

Provides in-place division operations between tensors. All implementations delegate to the underlying div_tensor method.

fn div_assign(&mut self, other: Tensor)

Divides this tensor by another tensor in-place

impl From<Tensor> for Vec<f32>

fn from(tensor: Tensor) -> Vec<f32>

Convert this tensor into a Vec<f32> in row-major order.

• Contiguous fast path: single optimized copy
• Non-contiguous: materialize a contiguous copy first
• Gradient metadata is ignored; this is a pure data extraction API

impl From<Vec<f32>> for Tensor

fn from(v: Vec<f32>) -> Self

Create a Tensor from a Vec<f32> by copying into an aligned/padded allocation.

Note: We do not adopt the Vec’s allocation to preserve alignment and padding guarantees for SIMD operations. Use Tensor::into_vec() to extract data back.

impl FromFieldValue for Tensor

fn from_field_value( value: FieldValue, field_name: &str, ) -> SerializationResult<Self>

Convert FieldValue to Tensor for use as struct field

Arguments

• value - FieldValue containing tensor data
• field_name - Name of the field for error reporting

Returns

Tensor instance or error if deserialization fails

impl FromIterator<Tensor> for Tensor

fn from_iter<I: IntoIterator<Item = Tensor>>(iter: I) -> Self

Collect element view tensors back into a single tensor

This method reconstructs a tensor from an iterator of element view tensors. It includes optimizations for common patterns and maintains gradient tracking when appropriate.

The collection process automatically detects whether all elements are scalar views (shape [1]) and uses optimized collection strategies accordingly. Gradient tracking is preserved when any input element requires gradients.

Performance

• Optimized Collection: Specialized paths for scalar and mixed views
• Memory Efficient: Direct memory copying without intermediate allocations
• Gradient Preservation: Maintains gradtrack functionality when enabled
• Shape Detection: Automatic detection of element shapes for optimization

Implementation Details

The method performs the following steps:

1. Element Collection: Gathers all element tensors from the iterator
2. Shape Analysis: Determines if all elements are scalar views
3. Optimized Path: Uses specialized collection for scalar views
4. General Path: Handles mixed shapes by flattening into 1D tensor
5. Gradient Setup: Preserves gradient tracking when appropriate

Examples

Basic Collection

use train_station::Tensor;

let original = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let doubled: Tensor = original.iter()
    .map(|elem| elem.mul_scalar(2.0))
    .collect();

assert_eq!(doubled.data(), &[2.0, 4.0, 6.0]);

Collection with Gradient Tracking

use train_station::Tensor;

let original = Tensor::from_slice(&[1.0, 2.0], vec![2])
    .unwrap()
    .with_requires_grad();

let result: Tensor = original.iter()
    .map(|elem| elem.mul_scalar(2.0))
    .collect();

assert!(result.requires_grad());
assert_eq!(result.data(), &[2.0, 4.0]);

Empty Iterator Handling

use train_station::Tensor;

let empty: Tensor = Vec::<Tensor>::new().into_iter().collect();
assert_eq!(empty.size(), 0);
assert_eq!(empty.shape().dims(), vec![0]);

impl FromIterator<f32> for Tensor

fn from_iter<I: IntoIterator<Item = f32>>(iter: I) -> Self

Collect f32 values into a 1D contiguous, SIMD-aligned Tensor

• Streams directly when iterator reports exact size_hint
• Falls back to temporary Vec and optimized_copy otherwise
• No gradient tracking is set on the result

impl<'a> IntoIterator for &'a Tensor

High-performance iterator over tensor elements as view tensors

Each element becomes a proper Tensor view of shape [1] that can use all existing tensor operations and gradient tracking. Implements all standard iterator traits for maximum compatibility with Rust’s ecosystem.

This iterator provides zero-copy access to tensor elements through view tensors, enabling efficient element-wise operations while maintaining full compatibility with Rust’s standard library iterator methods.

Performance

• Zero-Copy Views: Each element is a view tensor sharing memory with source
• O(1) Element Access: Constant-time view creation for each element
• Memory Efficient: ~64 bytes overhead per element view
• SIMD Compatible: All tensor operations use existing optimizations
• Gradient Tracking: Full gradtrack support through element operations

Implementation Details

The iterator creates lightweight view tensors on-demand, sharing the same memory allocation as the source tensor. This ensures zero-copy semantics while maintaining full tensor operation compatibility.

Each element view is created using Tensor::element_view(), which provides a true view of the underlying data without any copying. The view tensors support all standard tensor operations including gradient tracking.

Standard Library Compatibility

This iterator implements all standard iterator traits:

• Iterator: Basic iteration with next() and size_hint()
• ExactSizeIterator: Precise size information with len()
• DoubleEndedIterator: Reverse iteration with next_back()
• FusedIterator: Fused iteration for better performance
• IntoIterator: Automatic conversion for for loops

Examples

Basic Iteration

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();

// Basic iteration
for element in tensor.iter() {
    println!("Element value: {}", element.value());
}

// Standard library methods
let sum: f32 = tensor.iter()
    .map(|elem| elem.value())
    .sum();

assert_eq!(sum, 6.0);

Element Operations

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();

// Tensor operations on elements
let transformed: Tensor = tensor.iter()
    .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0)) // 2x + 1
    .collect();

assert_eq!(transformed.data(), &[3.0, 5.0, 7.0]);

Advanced Iterator Methods

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5]).unwrap();

// Filter and transform
let result: Tensor = tensor.iter()
    .filter(|elem| elem.value() > 2.0)
    .map(|elem| elem.mul_scalar(10.0))
    .collect();

assert_eq!(result.data(), &[30.0, 40.0, 50.0]);

// Reverse iteration
let reversed: Tensor = tensor.iter().rev().collect();
assert_eq!(reversed.data(), &[5.0, 4.0, 3.0, 2.0, 1.0]);

IntoIterator for &Tensor now iterates outermost dimension, yielding sub-tensors (views)

type Item = Tensor

The type of the elements being iterated over.

type IntoIter = TensorDimIterator<'a>

Which kind of iterator are we turning this into?

fn into_iter(self) -> Self::IntoIter

Creates an iterator from a value.

impl IntoIterator for Tensor

IntoIterator for owned Tensor: iterate outermost dimension producing sub-tensors. Enables .into_iter().flatten() patterns on owned tensors.

type Item = Tensor

The type of the elements being iterated over.

type IntoIter = TensorDimOwnedIterator

Which kind of iterator are we turning this into?

fn into_iter(self) -> Self::IntoIter

Creates an iterator from a value.

impl Mul<&Tensor> for Tensor

fn mul(self, other: &Tensor) -> Tensor

Multiplies a tensor and a tensor reference element-wise

Returns

A new tensor containing the element-wise product

type Output = Tensor

The resulting type after applying the * operator.

impl Mul<&Tensor> for f32

fn mul(self, tensor: &Tensor) -> Tensor

Multiplies each element of the tensor by a scalar (reference version)

Returns

A new tensor with each element multiplied by the scalar

type Output = Tensor

The resulting type after applying the * operator.

impl Mul<Tensor> for &Tensor

fn mul(self, other: Tensor) -> Tensor

Multiplies a tensor reference and a tensor element-wise

Returns

A new tensor containing the element-wise product

type Output = Tensor

The resulting type after applying the * operator.

impl Mul<Tensor> for f32

Scalar-tensor multiplication operator implementations

Provides multiplication operations between scalars and tensors. All implementations delegate to the underlying mul_scalar method.

fn mul(self, tensor: Tensor) -> Tensor

Multiplies each element of the tensor by a scalar

Returns

A new tensor with each element multiplied by the scalar

type Output = Tensor

The resulting type after applying the * operator.

impl Mul<f32> for &Tensor

fn mul(self, scalar: f32) -> Tensor

Multiplies each element of the tensor by a scalar (reference version)

Returns

A new tensor with each element multiplied by the scalar

type Output = Tensor

The resulting type after applying the * operator.

impl Mul<f32> for Tensor

Tensor-scalar multiplication operator implementations

Provides multiplication operations between tensors and scalars. All implementations delegate to the underlying mul_scalar method.

fn mul(self, scalar: f32) -> Tensor

Multiplies each element of the tensor by a scalar

Returns

A new tensor with each element multiplied by the scalar

type Output = Tensor

The resulting type after applying the * operator.

impl Mul for &Tensor

fn mul(self, other: &Tensor) -> Tensor

Multiplies two tensors element-wise (reference version)

Returns

A new tensor containing the element-wise product

type Output = Tensor

The resulting type after applying the * operator.

impl Mul for Tensor

Tensor multiplication operator implementations

Provides element-wise multiplication operations between tensors with various reference combinations. All implementations delegate to the underlying mul_tensor method for optimal performance.

fn mul(self, other: Tensor) -> Tensor

Multiplies two tensors element-wise

Returns

A new tensor containing the element-wise product

type Output = Tensor

The resulting type after applying the * operator.

impl MulAssign<&Tensor> for Tensor

fn mul_assign(&mut self, other: &Tensor)

Multiplies this tensor by another tensor reference in-place

impl MulAssign<f32> for Tensor

Tensor-scalar multiplication assignment operator implementations

Provides in-place multiplication operations between tensors and scalars.

fn mul_assign(&mut self, scalar: f32)

Multiplies each element of this tensor by a scalar in-place

impl MulAssign for Tensor

Tensor multiplication assignment operator implementations

Provides in-place multiplication operations between tensors. All implementations delegate to the underlying mul_tensor method.

fn mul_assign(&mut self, other: Tensor)

Multiplies this tensor by another tensor in-place

impl Neg for &Tensor

fn neg(self) -> Tensor

Negates each element of the tensor (reference version)

Returns

A new tensor with each element negated

type Output = Tensor

The resulting type after applying the - operator.

impl Neg for Tensor

Tensor negation operator implementations

Provides unary negation operations for tensors. All implementations delegate to the underlying mul_scalar method with -1.0.

fn neg(self) -> Tensor

Negates each element of the tensor

Returns

A new tensor with each element negated

type Output = Tensor

The resulting type after applying the - operator.

impl Serializable for Tensor

fn to_json(&self) -> SerializationResult<String>

Serialize the tensor to JSON format

This method converts the tensor into a human-readable JSON string representation that includes all tensor data, shape information, device placement, and gradtrack state. The JSON format is suitable for debugging, configuration files, and cross-language interoperability.

Returns

JSON string representation of the tensor on success, or SerializationError on failure

Examples

use train_station::Tensor;
use train_station::serialization::Serializable;

let mut tensor = Tensor::zeros(vec![2, 3]);
tensor.set(&[0, 0], 1.0);
tensor.set(&[1, 2], 5.0);

let json = tensor.to_json().unwrap();
assert!(!json.is_empty());
assert!(json.contains("data"));
assert!(json.contains("shape"));

fn from_json(json: &str) -> SerializationResult<Self>

Deserialize a tensor from JSON format

This method parses a JSON string and reconstructs a tensor with all its data, shape information, device placement, and gradtrack state. The JSON must contain all necessary fields in the expected format.

Arguments

• json - JSON string containing serialized tensor data

Returns

The deserialized tensor on success, or SerializationError on failure

Examples

use train_station::Tensor;
use train_station::serialization::Serializable;

let mut original = Tensor::ones(vec![2, 2]);
original.set(&[0, 1], 3.0);
original.set_requires_grad(true);

let json = original.to_json().unwrap();
let restored = Tensor::from_json(&json).unwrap();

assert_eq!(original.shape().dims(), restored.shape().dims());
assert_eq!(original.get(&[0, 1]), restored.get(&[0, 1]));
assert_eq!(original.requires_grad(), restored.requires_grad());

fn to_binary(&self) -> SerializationResult<Vec<u8>>

Serialize the tensor to binary format

This method converts the tensor into a compact binary representation optimized for storage and transmission. The binary format provides maximum performance and minimal file sizes, making it ideal for large tensors and production use.

Returns

Binary representation of the tensor on success, or SerializationError on failure

Examples

use train_station::Tensor;
use train_station::serialization::Serializable;

let mut tensor = Tensor::zeros(vec![100, 100]);
for i in 0..10 {
    tensor.set(&[i, i], i as f32);
}

let binary = tensor.to_binary().unwrap();
assert!(!binary.is_empty());
// Binary format is more compact than JSON for large tensors

fn from_binary(data: &[u8]) -> SerializationResult<Self>

Deserialize a tensor from binary format

This method parses binary data and reconstructs a tensor with all its data, shape information, device placement, and gradtrack state. The binary data must contain complete serialized information in the expected format.

Arguments

• data - Binary data containing serialized tensor information

Returns

The deserialized tensor on success, or SerializationError on failure

Examples

use train_station::Tensor;
use train_station::serialization::Serializable;

let mut original = Tensor::ones(vec![3, 4]);
original.set(&[2, 3], 7.5);
original.set_requires_grad(true);

let binary = original.to_binary().unwrap();
let restored = Tensor::from_binary(&binary).unwrap();

assert_eq!(original.shape().dims(), restored.shape().dims());
assert_eq!(original.get(&[2, 3]), restored.get(&[2, 3]));
assert_eq!(original.requires_grad(), restored.requires_grad());

fn save<P: AsRef<Path>>( &self, path: P, format: Format, ) -> SerializationResult<()>

Save the object to a file in the specified format

fn save_to_writer<W: Write>( &self, writer: &mut W, format: Format, ) -> SerializationResult<()>

Save the object to a writer in the specified format

fn load<P: AsRef<Path>>(path: P, format: Format) -> SerializationResult<Self>

Load an object from a file in the specified format

fn load_from_reader<R: Read>( reader: &mut R, format: Format, ) -> SerializationResult<Self>

Load an object from a reader in the specified format

impl StructSerializable for Tensor

fn to_serializer(&self) -> StructSerializer

Convert Tensor to StructSerializer for serialization

Serializes tensor data, shape, device, and gradtrack state. Runtime state (id, grad, grad_fn, allocation_owner) is not serialized.

Returns

StructSerializer containing all persistent tensor state

fn from_deserializer( deserializer: &mut StructDeserializer, ) -> SerializationResult<Self>

Create Tensor from StructDeserializer

Reconstructs tensor from serialized data, shape, device, and gradtrack state. Allocates new memory and generates new tensor ID.

Arguments

• deserializer - StructDeserializer containing tensor data

Returns

Reconstructed Tensor instance or error if deserialization fails

fn save_json<P: AsRef<Path>>(&self, path: P) -> SerializationResult<()>

Saves the struct to a JSON file

fn save_binary<P: AsRef<Path>>(&self, path: P) -> SerializationResult<()>

Saves the struct to a binary file

fn load_json<P: AsRef<Path>>(path: P) -> SerializationResult<Self>

Loads the struct from a JSON file

fn load_binary<P: AsRef<Path>>(path: P) -> SerializationResult<Self>

Loads the struct from a binary file

fn to_json(&self) -> SerializationResult<String>

Converts the struct to a JSON string

fn to_binary(&self) -> SerializationResult<Vec<u8>>

Converts the struct to binary data

fn from_json(json: &str) -> SerializationResult<Self>

Creates the struct from a JSON string

fn from_binary(data: &[u8]) -> SerializationResult<Self>

Creates the struct from binary data

impl Sub<&Tensor> for Tensor

fn sub(self, other: &Tensor) -> Tensor

Subtracts a tensor reference from a tensor element-wise

Returns

A new tensor containing the element-wise difference

type Output = Tensor

The resulting type after applying the - operator.

impl Sub<&Tensor> for f32

fn sub(self, tensor: &Tensor) -> Tensor

Subtracts each element of the tensor from the scalar (reference version)

Returns

A new tensor with each element subtracted from the scalar

type Output = Tensor

The resulting type after applying the - operator.

impl Sub<Tensor> for &Tensor

fn sub(self, other: Tensor) -> Tensor

Subtracts a tensor from a tensor reference element-wise

Returns

A new tensor containing the element-wise difference

type Output = Tensor

The resulting type after applying the - operator.

impl Sub<Tensor> for f32

Scalar-tensor subtraction operator implementations

Provides subtraction operations between scalars and tensors. Computes scalar - tensor by negating the tensor and adding the scalar.

fn sub(self, tensor: Tensor) -> Tensor

Subtracts each element of the tensor from the scalar

Returns

A new tensor with each element subtracted from the scalar

type Output = Tensor

The resulting type after applying the - operator.

impl Sub<f32> for &Tensor

fn sub(self, scalar: f32) -> Tensor

Subtracts a scalar from each element of the tensor (reference version)

Returns

A new tensor with the scalar subtracted from each element

type Output = Tensor

The resulting type after applying the - operator.

impl Sub<f32> for Tensor

Tensor-scalar subtraction operator implementations

Provides subtraction operations between tensors and scalars. All implementations delegate to the underlying sub_scalar method.

fn sub(self, scalar: f32) -> Tensor

Subtracts a scalar from each element of the tensor

Returns

A new tensor with the scalar subtracted from each element

type Output = Tensor

The resulting type after applying the - operator.

impl Sub for &Tensor

fn sub(self, other: &Tensor) -> Tensor

Subtracts two tensors element-wise (reference version)

Returns

A new tensor containing the element-wise difference

type Output = Tensor

The resulting type after applying the - operator.

impl Sub for Tensor

Tensor subtraction operator implementations

Provides subtraction operations between tensors with various reference combinations. All implementations delegate to the underlying sub_tensor method for optimal performance.

fn sub(self, other: Tensor) -> Tensor

Subtracts two tensors element-wise

Returns

A new tensor containing the element-wise difference

type Output = Tensor

The resulting type after applying the - operator.

impl SubAssign<&Tensor> for Tensor

fn sub_assign(&mut self, other: &Tensor)

Subtracts another tensor reference from this tensor in-place

impl SubAssign<f32> for Tensor

Tensor-scalar subtraction assignment operator implementations

Provides in-place subtraction operations between tensors and scalars.

fn sub_assign(&mut self, scalar: f32)

Subtracts a scalar from each element of this tensor in-place

impl SubAssign for Tensor

Tensor subtraction assignment operator implementations

Provides in-place subtraction operations between tensors. All implementations delegate to the underlying sub_tensor method.

fn sub_assign(&mut self, other: Tensor)

Subtracts another tensor from this tensor in-place

impl Send for Tensor

impl Sync for Tensor