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 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_linear_layer.rs (line 156)

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)
    );
}

examples/iterators/element_iteration.rs (line 91)

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(())
}

examples/getting_started/tensor_operators.rs (line 165)

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()
    );
}

examples/getting_started/tensor_basics.rs (line 49)

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()
    );
}

/// Demonstrate basic arithmetic operations
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());
}

/// Demonstrate shape manipulation operations
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()
    );
}

/// Demonstrate data access patterns
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);
    }
}

/// Demonstrate utility functions
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/getting_started/optimizer_basics.rs (line 57)

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 355)

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()
    );
}

/// Demonstrate different configurable architectures
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))
        );
    }
}

/// Demonstrate basic training workflow
fn demonstrate_training_workflow() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Workflow ---");

    // Create a simple classification network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 3],
        output_size: 1,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(46));

    println!("Training network: 2 -> [4, 3] -> 1");

    // Create simple binary classification data: XOR problem
    let x_data = Tensor::from_slice(
        &[
            0.0, 0.0, // -> 0
            0.0, 1.0, // -> 1
            1.0, 0.0, // -> 1
            1.0, 1.0, // -> 0
        ],
        vec![4, 2],
    )
    .unwrap();

    let y_true = Tensor::from_slice(&[0.0, 1.0, 1.0, 0.0], vec![4, 1]).unwrap();

    println!("Training on XOR problem:");
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.1);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

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

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        losses.push(loss.value());

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

    // Test final model
    let final_predictions = network.forward_no_grad(&x_data);
    println!("\nFinal predictions vs targets:");
    for i in 0..4 {
        let pred = final_predictions.data()[i];
        let target = y_true.data()[i];
        let input_x = x_data.data()[i * 2];
        let input_y = x_data.data()[i * 2 + 1];
        println!(
            "  [{:.0}, {:.0}] -> pred: {:.3}, target: {:.0}, error: {:.3}",
            input_x,
            input_y,
            pred,
            target,
            (pred - target).abs()
        );
    }

    Ok(())
}

/// Demonstrate comprehensive training with 100+ steps
fn demonstrate_comprehensive_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Comprehensive Training (100+ Steps) ---");

    // Create a regression network
    let config = FeedForwardConfig {
        input_size: 3,
        hidden_sizes: vec![8, 6, 4],
        output_size: 2,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(47));

    println!("Network architecture: 3 -> [8, 6, 4] -> 2");
    println!("Total parameters: {}", network.parameter_count());

    // Create synthetic regression data
    // Target function: [y1, y2] = [x1 + 2*x2 - x3, x1*x2 + x3]
    let num_samples = 32;
    let mut x_vec = Vec::new();
    let mut y_vec = Vec::new();

    for i in 0..num_samples {
        let x1 = (i as f32 / num_samples as f32) * 2.0 - 1.0; // [-1, 1]
        let x2 = ((i * 2) as f32 / num_samples as f32) * 2.0 - 1.0;
        let x3 = ((i * 3) as f32 / num_samples as f32) * 2.0 - 1.0;

        let y1 = x1 + 2.0 * x2 - x3;
        let y2 = x1 * x2 + x3;

        x_vec.extend_from_slice(&[x1, x2, x3]);
        y_vec.extend_from_slice(&[y1, y2]);
    }

    let x_data = Tensor::from_slice(&x_vec, vec![num_samples, 3]).unwrap();
    let y_true = Tensor::from_slice(&y_vec, vec![num_samples, 2]).unwrap();

    println!("Training data:");
    println!("  {} samples", num_samples);
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer with learning rate scheduling
    let mut optimizer = Adam::with_learning_rate(0.01);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    // Comprehensive training loop (150 epochs)
    let num_epochs = 150;
    let mut losses = Vec::new();
    let mut best_loss = f32::INFINITY;
    let mut patience_counter = 0;
    let patience = 20;

    println!("Starting comprehensive training...");

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        let current_loss = loss.value();
        losses.push(current_loss);

        // Learning rate scheduling
        if epoch > 0 && epoch % 30 == 0 {
            let new_lr = optimizer.learning_rate() * 0.8;
            optimizer.set_learning_rate(new_lr);
            println!("  Reduced learning rate to {:.4}", new_lr);
        }

        // Early stopping logic
        if current_loss < best_loss {
            best_loss = current_loss;
            patience_counter = 0;
        } else {
            patience_counter += 1;
        }

        // Print progress
        if epoch % 25 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:3}: Loss = {:.6}, LR = {:.4}, Best = {:.6}",
                epoch,
                current_loss,
                optimizer.learning_rate(),
                best_loss
            );
        }

        // Early stopping
        if patience_counter >= patience && epoch > 50 {
            println!("Early stopping at epoch {} (patience exceeded)", epoch);
            break;
        }
    }

    // Final evaluation
    let final_predictions = network.forward_no_grad(&x_data);

    // Compute final metrics
    let final_loss = losses[losses.len() - 1];
    let initial_loss = losses[0];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining completed!");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Best loss: {:.6}", best_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final learning rate: {:.4}", optimizer.learning_rate());

    // Sample predictions analysis
    println!("\nSample predictions (first 5):");
    for i in 0..5.min(num_samples) {
        let pred1 = final_predictions.data()[i * 2];
        let pred2 = final_predictions.data()[i * 2 + 1];
        let true1 = y_true.data()[i * 2];
        let true2 = y_true.data()[i * 2 + 1];

        println!(
            "  Sample {}: pred=[{:.3}, {:.3}], true=[{:.3}, {:.3}], error=[{:.3}, {:.3}]",
            i + 1,
            pred1,
            pred2,
            true1,
            true2,
            (pred1 - true1).abs(),
            (pred2 - true2).abs()
        );
    }

    Ok(())
}

Additional examples can be found in:

• examples/getting_started/serialization_basics.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 187)

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 (process in chunks)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();

    let mut streamed_result = Vec::new();
    for chunk_start in (0..size).step_by(chunk_size) {
        let chunk_end = (chunk_start + chunk_size).min(size);
        let chunk: Tensor = tensor
            .iter_range(chunk_start, chunk_end)
            .map(|elem| elem.pow_scalar(2.0).sqrt())
            .collect();
        streamed_result.extend(chunk.data().iter().cloned());
    }
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .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 + batch_size - 1) / 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 115)

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
    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
    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()
    );

    // 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/neural_networks/feedforward_network.rs (line 73)

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

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

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

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

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

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

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

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

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

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

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

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

        Self { layers, config }
    }

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

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

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

        x
    }

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

    /// Get all parameters for optimization
    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
        let mut params = Vec::new();
        for layer in &mut self.layers {
            params.extend(layer.parameters());
        }
        params
    }

    /// Get the number of layers
    pub fn num_layers(&self) -> usize {
        self.layers.len()
    }

    /// Get the total number of parameters
    pub fn parameter_count(&self) -> usize {
        let mut count = 0;
        let mut current_size = self.config.input_size;

        for &hidden_size in &self.config.hidden_sizes {
            count += current_size * hidden_size + hidden_size; // weights + bias
            current_size = hidden_size;
        }

        // Output layer
        count += current_size * self.config.output_size + self.config.output_size;

        count
    }

    /// Save network parameters to JSON
    pub fn save_json(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
        if let Some(parent) = std::path::Path::new(path).parent() {
            fs::create_dir_all(parent)?;
        }

        for (i, layer) in self.layers.iter().enumerate() {
            let layer_path = format!("{}_layer_{}", path, i);
            let weight_path = format!("{}_weight.json", layer_path);
            let bias_path = format!("{}_bias.json", layer_path);

            layer.weight.save_json(&weight_path)?;
            layer.bias.save_json(&bias_path)?;
        }

        println!(
            "Saved feed-forward network to {} ({} layers)",
            path,
            self.layers.len()
        );
        Ok(())
    }

    /// Load network parameters from JSON
    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/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,
        }
    }

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

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 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_by_value().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(())
}

examples/optimizers/adam_configurations.rs (line 92)

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/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());

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/neural_networks/basic_linear_layer.rs (line 158)

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 176)

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 358)

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 394)

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 the accumulated gradients (if any)

Returns a reference to the gradient tensor if gradients have been computed and this tensor has gradient tracking enabled.

Returns

Optional reference to the gradient tensor, or None if no gradients exist

Examples

use train_station::Tensor;

let tensor = Tensor::ones(vec![2, 3]).with_requires_grad();
assert!(tensor.grad().is_none()); // No gradients computed yet

Examples found in repository

examples/iterators/element_iteration.rs (line 183)

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 grad_by_value(&self) -> Option<Tensor>

Get accumulated gradient by value (helper for testing)

Returns the gradient tensor by value, which is useful for testing and when you need to own the gradient data.

Returns

Optional gradient tensor, or None if no gradients exist

Examples

use train_station::Tensor;

let tensor = Tensor::ones(vec![2, 3]).with_requires_grad();
assert!(tensor.grad_by_value().is_none()); // No gradients computed yet

Examples found in repository

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_by_value().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(())
}

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/iterators/element_iteration.rs (line 180)

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 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/optimizers/adam_configurations.rs (line 115)

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 347)

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 138)

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_by_value().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/feedforward_network.rs (line 482)

fn demonstrate_training_workflow() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Workflow ---");

    // Create a simple classification network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 3],
        output_size: 1,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(46));

    println!("Training network: 2 -> [4, 3] -> 1");

    // Create simple binary classification data: XOR problem
    let x_data = Tensor::from_slice(
        &[
            0.0, 0.0, // -> 0
            0.0, 1.0, // -> 1
            1.0, 0.0, // -> 1
            1.0, 1.0, // -> 0
        ],
        vec![4, 2],
    )
    .unwrap();

    let y_true = Tensor::from_slice(&[0.0, 1.0, 1.0, 0.0], vec![4, 1]).unwrap();

    println!("Training on XOR problem:");
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.1);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

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

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        losses.push(loss.value());

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

    // Test final model
    let final_predictions = network.forward_no_grad(&x_data);
    println!("\nFinal predictions vs targets:");
    for i in 0..4 {
        let pred = final_predictions.data()[i];
        let target = y_true.data()[i];
        let input_x = x_data.data()[i * 2];
        let input_y = x_data.data()[i * 2 + 1];
        println!(
            "  [{:.0}, {:.0}] -> pred: {:.3}, target: {:.0}, error: {:.3}",
            input_x,
            input_y,
            pred,
            target,
            (pred - target).abs()
        );
    }

    Ok(())
}

/// Demonstrate comprehensive training with 100+ steps
fn demonstrate_comprehensive_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Comprehensive Training (100+ Steps) ---");

    // Create a regression network
    let config = FeedForwardConfig {
        input_size: 3,
        hidden_sizes: vec![8, 6, 4],
        output_size: 2,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(47));

    println!("Network architecture: 3 -> [8, 6, 4] -> 2");
    println!("Total parameters: {}", network.parameter_count());

    // Create synthetic regression data
    // Target function: [y1, y2] = [x1 + 2*x2 - x3, x1*x2 + x3]
    let num_samples = 32;
    let mut x_vec = Vec::new();
    let mut y_vec = Vec::new();

    for i in 0..num_samples {
        let x1 = (i as f32 / num_samples as f32) * 2.0 - 1.0; // [-1, 1]
        let x2 = ((i * 2) as f32 / num_samples as f32) * 2.0 - 1.0;
        let x3 = ((i * 3) as f32 / num_samples as f32) * 2.0 - 1.0;

        let y1 = x1 + 2.0 * x2 - x3;
        let y2 = x1 * x2 + x3;

        x_vec.extend_from_slice(&[x1, x2, x3]);
        y_vec.extend_from_slice(&[y1, y2]);
    }

    let x_data = Tensor::from_slice(&x_vec, vec![num_samples, 3]).unwrap();
    let y_true = Tensor::from_slice(&y_vec, vec![num_samples, 2]).unwrap();

    println!("Training data:");
    println!("  {} samples", num_samples);
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer with learning rate scheduling
    let mut optimizer = Adam::with_learning_rate(0.01);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    // Comprehensive training loop (150 epochs)
    let num_epochs = 150;
    let mut losses = Vec::new();
    let mut best_loss = f32::INFINITY;
    let mut patience_counter = 0;
    let patience = 20;

    println!("Starting comprehensive training...");

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        let current_loss = loss.value();
        losses.push(current_loss);

        // Learning rate scheduling
        if epoch > 0 && epoch % 30 == 0 {
            let new_lr = optimizer.learning_rate() * 0.8;
            optimizer.set_learning_rate(new_lr);
            println!("  Reduced learning rate to {:.4}", new_lr);
        }

        // Early stopping logic
        if current_loss < best_loss {
            best_loss = current_loss;
            patience_counter = 0;
        } else {
            patience_counter += 1;
        }

        // Print progress
        if epoch % 25 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:3}: Loss = {:.6}, LR = {:.4}, Best = {:.6}",
                epoch,
                current_loss,
                optimizer.learning_rate(),
                best_loss
            );
        }

        // Early stopping
        if patience_counter >= patience && epoch > 50 {
            println!("Early stopping at epoch {} (patience exceeded)", epoch);
            break;
        }
    }

    // Final evaluation
    let final_predictions = network.forward_no_grad(&x_data);

    // Compute final metrics
    let final_loss = losses[losses.len() - 1];
    let initial_loss = losses[0];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining completed!");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Best loss: {:.6}", best_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final learning rate: {:.4}", optimizer.learning_rate());

    // Sample predictions analysis
    println!("\nSample predictions (first 5):");
    for i in 0..5.min(num_samples) {
        let pred1 = final_predictions.data()[i * 2];
        let pred2 = final_predictions.data()[i * 2 + 1];
        let true1 = y_true.data()[i * 2];
        let true2 = y_true.data()[i * 2 + 1];

        println!(
            "  Sample {}: pred=[{:.3}, {:.3}], true=[{:.3}, {:.3}], error=[{:.3}, {:.3}]",
            i + 1,
            pred1,
            pred2,
            true1,
            true2,
            (pred1 - true1).abs(),
            (pred2 - true2).abs()
        );
    }

    Ok(())
}

/// Demonstrate network serialization
fn demonstrate_network_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Network Serialization ---");

    // Create and train a network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 2],
        output_size: 1,
        use_bias: true,
    };
    let mut original_network = FeedForwardNetwork::new(config.clone(), Some(48));

    // Quick training
    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_network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    for _ in 0..20 {
        let y_pred = original_network.forward(&x_data);
        let mut loss = (y_pred.sub_tensor(&y_true)).pow_scalar(2.0).mean();
        loss.backward(None);

        let mut params = original_network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);
    }

    // Test original network
    let test_input = Tensor::from_slice(&[1.0, 1.0], vec![1, 2]).unwrap();
    let original_output = original_network.forward_no_grad(&test_input);

    println!("Original network output: {:?}", original_output.data());

    // Save network
    original_network.save_json("temp_feedforward_network")?;

    // Load network
    let loaded_network = FeedForwardNetwork::load_json("temp_feedforward_network", config)?;
    let loaded_output = loaded_network.forward_no_grad(&test_input);

    println!("Loaded network output: {:?}", loaded_output.data());

    // Verify consistency
    let match_check = original_output
        .data()
        .iter()
        .zip(loaded_output.data().iter())
        .all(|(a, b)| (a - b).abs() < 1e-6);

    println!(
        "Serialization verification: {}",
        if match_check { "PASSED" } else { "FAILED" }
    );

    Ok(())
}

Additional examples can be found in:

• examples/neural_networks/basic_linear_layer.rs
• examples/iterators/performance_optimization.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.

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 is_view(&self) -> bool

Checks if this tensor is a view of another tensor

Returns

true if this tensor is a view (non-contiguous reference)

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 layout(&self) -> &MemoryLayout

Gets the memory layout type for optimization decisions

Returns

Reference to the memory layout information

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

pub fn broadcast_with( &self, other: &Tensor, ) -> Result<(Tensor, Tensor, Shape), BroadcastError>

Broadcast this tensor with another tensor for element-wise operations

Returns a tuple containing:

• Broadcasted view of self
• Broadcasted view of other
• Result shape for the operation

Arguments

• other - The tensor to broadcast with

Returns

A tuple (broadcasted_self, broadcasted_other, result_shape)

Examples

use train_station::Tensor;

let a = Tensor::ones(vec![2, 1, 4]);
let b = Tensor::ones(vec![3, 1]);
let result = a.broadcast_with(&b);
assert!(result.is_ok());

pub fn is_simd_aligned(&self) -> bool

Checks if the tensor data is properly aligned for SIMD operations

Returns

true if the tensor data is aligned to 32-byte boundaries for AVX2

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/getting_started/tensor_operators.rs (line 52)

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/neural_networks/basic_linear_layer.rs (line 173)

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(())
}

examples/iterators/element_iteration.rs (line 82)

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(())
}

examples/getting_started/tensor_basics.rs (line 50)

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()
    );
}

/// Demonstrate basic arithmetic operations
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());
}

/// Demonstrate shape manipulation operations
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()
    );
}

/// Demonstrate data access patterns
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);
    }
}

examples/neural_networks/feedforward_network.rs (line 357)

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()
    );
}

/// Demonstrate different configurable architectures
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))
        );
    }
}

/// Demonstrate basic training workflow
fn demonstrate_training_workflow() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Workflow ---");

    // Create a simple classification network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 3],
        output_size: 1,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(46));

    println!("Training network: 2 -> [4, 3] -> 1");

    // Create simple binary classification data: XOR problem
    let x_data = Tensor::from_slice(
        &[
            0.0, 0.0, // -> 0
            0.0, 1.0, // -> 1
            1.0, 0.0, // -> 1
            1.0, 1.0, // -> 0
        ],
        vec![4, 2],
    )
    .unwrap();

    let y_true = Tensor::from_slice(&[0.0, 1.0, 1.0, 0.0], vec![4, 1]).unwrap();

    println!("Training on XOR problem:");
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.1);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

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

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        losses.push(loss.value());

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

    // Test final model
    let final_predictions = network.forward_no_grad(&x_data);
    println!("\nFinal predictions vs targets:");
    for i in 0..4 {
        let pred = final_predictions.data()[i];
        let target = y_true.data()[i];
        let input_x = x_data.data()[i * 2];
        let input_y = x_data.data()[i * 2 + 1];
        println!(
            "  [{:.0}, {:.0}] -> pred: {:.3}, target: {:.0}, error: {:.3}",
            input_x,
            input_y,
            pred,
            target,
            (pred - target).abs()
        );
    }

    Ok(())
}

/// Demonstrate comprehensive training with 100+ steps
fn demonstrate_comprehensive_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Comprehensive Training (100+ Steps) ---");

    // Create a regression network
    let config = FeedForwardConfig {
        input_size: 3,
        hidden_sizes: vec![8, 6, 4],
        output_size: 2,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(47));

    println!("Network architecture: 3 -> [8, 6, 4] -> 2");
    println!("Total parameters: {}", network.parameter_count());

    // Create synthetic regression data
    // Target function: [y1, y2] = [x1 + 2*x2 - x3, x1*x2 + x3]
    let num_samples = 32;
    let mut x_vec = Vec::new();
    let mut y_vec = Vec::new();

    for i in 0..num_samples {
        let x1 = (i as f32 / num_samples as f32) * 2.0 - 1.0; // [-1, 1]
        let x2 = ((i * 2) as f32 / num_samples as f32) * 2.0 - 1.0;
        let x3 = ((i * 3) as f32 / num_samples as f32) * 2.0 - 1.0;

        let y1 = x1 + 2.0 * x2 - x3;
        let y2 = x1 * x2 + x3;

        x_vec.extend_from_slice(&[x1, x2, x3]);
        y_vec.extend_from_slice(&[y1, y2]);
    }

    let x_data = Tensor::from_slice(&x_vec, vec![num_samples, 3]).unwrap();
    let y_true = Tensor::from_slice(&y_vec, vec![num_samples, 2]).unwrap();

    println!("Training data:");
    println!("  {} samples", num_samples);
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer with learning rate scheduling
    let mut optimizer = Adam::with_learning_rate(0.01);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    // Comprehensive training loop (150 epochs)
    let num_epochs = 150;
    let mut losses = Vec::new();
    let mut best_loss = f32::INFINITY;
    let mut patience_counter = 0;
    let patience = 20;

    println!("Starting comprehensive training...");

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        let current_loss = loss.value();
        losses.push(current_loss);

        // Learning rate scheduling
        if epoch > 0 && epoch % 30 == 0 {
            let new_lr = optimizer.learning_rate() * 0.8;
            optimizer.set_learning_rate(new_lr);
            println!("  Reduced learning rate to {:.4}", new_lr);
        }

        // Early stopping logic
        if current_loss < best_loss {
            best_loss = current_loss;
            patience_counter = 0;
        } else {
            patience_counter += 1;
        }

        // Print progress
        if epoch % 25 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:3}: Loss = {:.6}, LR = {:.4}, Best = {:.6}",
                epoch,
                current_loss,
                optimizer.learning_rate(),
                best_loss
            );
        }

        // Early stopping
        if patience_counter >= patience && epoch > 50 {
            println!("Early stopping at epoch {} (patience exceeded)", epoch);
            break;
        }
    }

    // Final evaluation
    let final_predictions = network.forward_no_grad(&x_data);

    // Compute final metrics
    let final_loss = losses[losses.len() - 1];
    let initial_loss = losses[0];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining completed!");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Best loss: {:.6}", best_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final learning rate: {:.4}", optimizer.learning_rate());

    // Sample predictions analysis
    println!("\nSample predictions (first 5):");
    for i in 0..5.min(num_samples) {
        let pred1 = final_predictions.data()[i * 2];
        let pred2 = final_predictions.data()[i * 2 + 1];
        let true1 = y_true.data()[i * 2];
        let true2 = y_true.data()[i * 2 + 1];

        println!(
            "  Sample {}: pred=[{:.3}, {:.3}], true=[{:.3}, {:.3}], error=[{:.3}, {:.3}]",
            i + 1,
            pred1,
            pred2,
            true1,
            true2,
            (pred1 - true1).abs(),
            (pred2 - true2).abs()
        );
    }

    Ok(())
}

/// Demonstrate network serialization
fn demonstrate_network_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Network Serialization ---");

    // Create and train a network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 2],
        output_size: 1,
        use_bias: true,
    };
    let mut original_network = FeedForwardNetwork::new(config.clone(), Some(48));

    // Quick training
    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_network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    for _ in 0..20 {
        let y_pred = original_network.forward(&x_data);
        let mut loss = (y_pred.sub_tensor(&y_true)).pow_scalar(2.0).mean();
        loss.backward(None);

        let mut params = original_network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);
    }

    // Test original network
    let test_input = Tensor::from_slice(&[1.0, 1.0], vec![1, 2]).unwrap();
    let original_output = original_network.forward_no_grad(&test_input);

    println!("Original network output: {:?}", original_output.data());

    // Save network
    original_network.save_json("temp_feedforward_network")?;

    // Load network
    let loaded_network = FeedForwardNetwork::load_json("temp_feedforward_network", config)?;
    let loaded_output = loaded_network.forward_no_grad(&test_input);

    println!("Loaded network output: {:?}", loaded_output.data());

    // Verify consistency
    let match_check = original_output
        .data()
        .iter()
        .zip(loaded_output.data().iter())
        .all(|(a, b)| (a - b).abs() < 1e-6);

    println!(
        "Serialization verification: {}",
        if match_check { "PASSED" } else { "FAILED" }
    );

    Ok(())
}

examples/getting_started/serialization_basics.rs (line 60)

fn demonstrate_tensor_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("--- Tensor Serialization ---");

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

    // Save tensor in JSON format
    let json_path = "temp_tensor.json";
    original_tensor.save_json(json_path)?;
    println!("Saved tensor to JSON: {}", json_path);

    // Load tensor from JSON
    let loaded_tensor_json = Tensor::load_json(json_path)?;
    println!(
        "Loaded from JSON: shape {:?}, data: {:?}",
        loaded_tensor_json.shape().dims,
        loaded_tensor_json.data()
    );

    // Verify data integrity
    assert_eq!(
        original_tensor.shape().dims,
        loaded_tensor_json.shape().dims
    );
    assert_eq!(original_tensor.data(), loaded_tensor_json.data());
    println!("JSON serialization verification: PASSED");

    // Save tensor in binary format
    let binary_path = "temp_tensor.bin";
    original_tensor.save_binary(binary_path)?;
    println!("Saved tensor to binary: {}", binary_path);

    // Load tensor from binary
    let loaded_tensor_binary = Tensor::load_binary(binary_path)?;
    println!(
        "Loaded from binary: shape {:?}, data: {:?}",
        loaded_tensor_binary.shape().dims,
        loaded_tensor_binary.data()
    );

    // Verify data integrity
    assert_eq!(
        original_tensor.shape().dims,
        loaded_tensor_binary.shape().dims
    );
    assert_eq!(original_tensor.data(), loaded_tensor_binary.data());
    println!("Binary serialization verification: PASSED");

    Ok(())
}

/// Demonstrate optimizer serialization and deserialization
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/iterators/performance_optimization.rs
• examples/getting_started/optimizer_basics.rs
• examples/iterators/advanced_patterns.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/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()
    );
}

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/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/iterators/element_iteration.rs (line 90)

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(())
}

examples/optimizers/adam_configurations.rs (line 102)

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/iterators/performance_optimization.rs (line 106)

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
        let start = Instant::now();
        let iterator_result: Tensor = tensor
            .iter()
            .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0))
            .collect();
        let iterator_time = start.elapsed();

        // Benchmark 3: Chained iterator operations
        let start = Instant::now();
        let _chained_result: Tensor = tensor
            .iter()
            .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();

        // Report results
        println!("  Direct operations: {:?}", direct_time);
        println!("  Iterator operations: {:?}", iterator_time);
        println!("  Chained operations: {:?}", chained_time);

        // Verify correctness
        assert_eq!(direct_result.data(), iterator_result.data());
        println!(
            "  Results match: {}",
            direct_result.data() == iterator_result.data()
        );

        // Performance ratio
        let ratio = iterator_time.as_nanos() as f64 / direct_time.as_nanos() as f64;
        println!("  Iterator/Direct ratio: {:.2}x", ratio);
    }

    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 (process in chunks)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();

    let mut streamed_result = Vec::new();
    for chunk_start in (0..size).step_by(chunk_size) {
        let chunk_end = (chunk_start + chunk_size).min(size);
        let chunk: Tensor = tensor
            .iter_range(chunk_start, chunk_end)
            .map(|elem| elem.pow_scalar(2.0).sqrt())
            .collect();
        streamed_result.extend(chunk.data().iter().cloned());
    }
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .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 + batch_size - 1) / 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/optimizers/learning_rate_scheduling.rs (line 351)

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 144)

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_by_value().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(())
}

Additional examples can be found in:

• examples/neural_networks/feedforward_network.rs
• examples/iterators/advanced_patterns.rs
• examples/neural_networks/basic_linear_layer.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/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/neural_networks/basic_linear_layer.rs (line 353)

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)
    );
}

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 119)

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.data(), &[2.0, 4.0]);

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;
}

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_by_value().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

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_by_value().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_by_value().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

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.

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_by_value().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/neural_networks/feedforward_network.rs (line 74)

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

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

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

examples/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,
        }
    }

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_by_value().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/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(),
    })
}

examples/optimizers/learning_rate_scheduling.rs (line 329)

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,
    })
}

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()
    );
}

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/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/neural_networks/basic_linear_layer.rs (line 169)

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(())
}

examples/iterators/element_iteration.rs (line 81)

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(())
}

examples/getting_started/tensor_basics.rs (line 62)

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()
    );
}

/// Demonstrate basic arithmetic operations
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());
}

/// Demonstrate shape manipulation operations
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()
    );
}

/// Demonstrate data access patterns
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);
    }
}

/// Demonstrate utility functions
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/feedforward_network.rs (line 351)

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()
    );
}

/// Demonstrate different configurable architectures
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))
        );
    }
}

/// Demonstrate basic training workflow
fn demonstrate_training_workflow() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Workflow ---");

    // Create a simple classification network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 3],
        output_size: 1,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(46));

    println!("Training network: 2 -> [4, 3] -> 1");

    // Create simple binary classification data: XOR problem
    let x_data = Tensor::from_slice(
        &[
            0.0, 0.0, // -> 0
            0.0, 1.0, // -> 1
            1.0, 0.0, // -> 1
            1.0, 1.0, // -> 0
        ],
        vec![4, 2],
    )
    .unwrap();

    let y_true = Tensor::from_slice(&[0.0, 1.0, 1.0, 0.0], vec![4, 1]).unwrap();

    println!("Training on XOR problem:");
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.1);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

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

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        losses.push(loss.value());

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

    // Test final model
    let final_predictions = network.forward_no_grad(&x_data);
    println!("\nFinal predictions vs targets:");
    for i in 0..4 {
        let pred = final_predictions.data()[i];
        let target = y_true.data()[i];
        let input_x = x_data.data()[i * 2];
        let input_y = x_data.data()[i * 2 + 1];
        println!(
            "  [{:.0}, {:.0}] -> pred: {:.3}, target: {:.0}, error: {:.3}",
            input_x,
            input_y,
            pred,
            target,
            (pred - target).abs()
        );
    }

    Ok(())
}

/// Demonstrate comprehensive training with 100+ steps
fn demonstrate_comprehensive_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Comprehensive Training (100+ Steps) ---");

    // Create a regression network
    let config = FeedForwardConfig {
        input_size: 3,
        hidden_sizes: vec![8, 6, 4],
        output_size: 2,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(47));

    println!("Network architecture: 3 -> [8, 6, 4] -> 2");
    println!("Total parameters: {}", network.parameter_count());

    // Create synthetic regression data
    // Target function: [y1, y2] = [x1 + 2*x2 - x3, x1*x2 + x3]
    let num_samples = 32;
    let mut x_vec = Vec::new();
    let mut y_vec = Vec::new();

    for i in 0..num_samples {
        let x1 = (i as f32 / num_samples as f32) * 2.0 - 1.0; // [-1, 1]
        let x2 = ((i * 2) as f32 / num_samples as f32) * 2.0 - 1.0;
        let x3 = ((i * 3) as f32 / num_samples as f32) * 2.0 - 1.0;

        let y1 = x1 + 2.0 * x2 - x3;
        let y2 = x1 * x2 + x3;

        x_vec.extend_from_slice(&[x1, x2, x3]);
        y_vec.extend_from_slice(&[y1, y2]);
    }

    let x_data = Tensor::from_slice(&x_vec, vec![num_samples, 3]).unwrap();
    let y_true = Tensor::from_slice(&y_vec, vec![num_samples, 2]).unwrap();

    println!("Training data:");
    println!("  {} samples", num_samples);
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer with learning rate scheduling
    let mut optimizer = Adam::with_learning_rate(0.01);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    // Comprehensive training loop (150 epochs)
    let num_epochs = 150;
    let mut losses = Vec::new();
    let mut best_loss = f32::INFINITY;
    let mut patience_counter = 0;
    let patience = 20;

    println!("Starting comprehensive training...");

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        let current_loss = loss.value();
        losses.push(current_loss);

        // Learning rate scheduling
        if epoch > 0 && epoch % 30 == 0 {
            let new_lr = optimizer.learning_rate() * 0.8;
            optimizer.set_learning_rate(new_lr);
            println!("  Reduced learning rate to {:.4}", new_lr);
        }

        // Early stopping logic
        if current_loss < best_loss {
            best_loss = current_loss;
            patience_counter = 0;
        } else {
            patience_counter += 1;
        }

        // Print progress
        if epoch % 25 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:3}: Loss = {:.6}, LR = {:.4}, Best = {:.6}",
                epoch,
                current_loss,
                optimizer.learning_rate(),
                best_loss
            );
        }

        // Early stopping
        if patience_counter >= patience && epoch > 50 {
            println!("Early stopping at epoch {} (patience exceeded)", epoch);
            break;
        }
    }

    // Final evaluation
    let final_predictions = network.forward_no_grad(&x_data);

    // Compute final metrics
    let final_loss = losses[losses.len() - 1];
    let initial_loss = losses[0];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining completed!");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Best loss: {:.6}", best_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final learning rate: {:.4}", optimizer.learning_rate());

    // Sample predictions analysis
    println!("\nSample predictions (first 5):");
    for i in 0..5.min(num_samples) {
        let pred1 = final_predictions.data()[i * 2];
        let pred2 = final_predictions.data()[i * 2 + 1];
        let true1 = y_true.data()[i * 2];
        let true2 = y_true.data()[i * 2 + 1];

        println!(
            "  Sample {}: pred=[{:.3}, {:.3}], true=[{:.3}, {:.3}], error=[{:.3}, {:.3}]",
            i + 1,
            pred1,
            pred2,
            true1,
            true2,
            (pred1 - true1).abs(),
            (pred2 - true2).abs()
        );
    }

    Ok(())
}

/// Demonstrate network serialization
fn demonstrate_network_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Network Serialization ---");

    // Create and train a network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 2],
        output_size: 1,
        use_bias: true,
    };
    let mut original_network = FeedForwardNetwork::new(config.clone(), Some(48));

    // Quick training
    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_network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    for _ in 0..20 {
        let y_pred = original_network.forward(&x_data);
        let mut loss = (y_pred.sub_tensor(&y_true)).pow_scalar(2.0).mean();
        loss.backward(None);

        let mut params = original_network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);
    }

    // Test original network
    let test_input = Tensor::from_slice(&[1.0, 1.0], vec![1, 2]).unwrap();
    let original_output = original_network.forward_no_grad(&test_input);

    println!("Original network output: {:?}", original_output.data());

    // Save network
    original_network.save_json("temp_feedforward_network")?;

    // Load network
    let loaded_network = FeedForwardNetwork::load_json("temp_feedforward_network", config)?;
    let loaded_output = loaded_network.forward_no_grad(&test_input);

    println!("Loaded network output: {:?}", loaded_output.data());

    // Verify consistency
    let match_check = original_output
        .data()
        .iter()
        .zip(loaded_output.data().iter())
        .all(|(a, b)| (a - b).abs() < 1e-6);

    println!(
        "Serialization verification: {}",
        if match_check { "PASSED" } else { "FAILED" }
    );

    Ok(())
}

examples/getting_started/serialization_basics.rs (line 56)

fn demonstrate_tensor_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("--- Tensor Serialization ---");

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

    // Save tensor in JSON format
    let json_path = "temp_tensor.json";
    original_tensor.save_json(json_path)?;
    println!("Saved tensor to JSON: {}", json_path);

    // Load tensor from JSON
    let loaded_tensor_json = Tensor::load_json(json_path)?;
    println!(
        "Loaded from JSON: shape {:?}, data: {:?}",
        loaded_tensor_json.shape().dims,
        loaded_tensor_json.data()
    );

    // Verify data integrity
    assert_eq!(
        original_tensor.shape().dims,
        loaded_tensor_json.shape().dims
    );
    assert_eq!(original_tensor.data(), loaded_tensor_json.data());
    println!("JSON serialization verification: PASSED");

    // Save tensor in binary format
    let binary_path = "temp_tensor.bin";
    original_tensor.save_binary(binary_path)?;
    println!("Saved tensor to binary: {}", binary_path);

    // Load tensor from binary
    let loaded_tensor_binary = Tensor::load_binary(binary_path)?;
    println!(
        "Loaded from binary: shape {:?}, data: {:?}",
        loaded_tensor_binary.shape().dims,
        loaded_tensor_binary.data()
    );

    // Verify data integrity
    assert_eq!(
        original_tensor.shape().dims,
        loaded_tensor_binary.shape().dims
    );
    assert_eq!(original_tensor.data(), loaded_tensor_binary.data());
    println!("Binary serialization verification: PASSED");

    Ok(())
}

Additional examples can be found in:

• examples/optimizers/adam_configurations.rs
• examples/iterators/performance_optimization.rs
• examples/optimizers/learning_rate_scheduling.rs
• examples/getting_started/optimizer_basics.rs
• examples/iterators/advanced_patterns.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/neural_networks/feedforward_network.rs (line 71)

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

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

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

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

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

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

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

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

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

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

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

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

        Self { layers, config }
    }

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

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

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

        x
    }

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

    /// Get all parameters for optimization
    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
        let mut params = Vec::new();
        for layer in &mut self.layers {
            params.extend(layer.parameters());
        }
        params
    }

    /// Get the number of layers
    pub fn num_layers(&self) -> usize {
        self.layers.len()
    }

    /// Get the total number of parameters
    pub fn parameter_count(&self) -> usize {
        let mut count = 0;
        let mut current_size = self.config.input_size;

        for &hidden_size in &self.config.hidden_sizes {
            count += current_size * hidden_size + hidden_size; // weights + bias
            current_size = hidden_size;
        }

        // Output layer
        count += current_size * self.config.output_size + self.config.output_size;

        count
    }

    /// Save network parameters to JSON
    pub fn save_json(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
        if let Some(parent) = std::path::Path::new(path).parent() {
            fs::create_dir_all(parent)?;
        }

        for (i, layer) in self.layers.iter().enumerate() {
            let layer_path = format!("{}_layer_{}", path, i);
            let weight_path = format!("{}_weight.json", layer_path);
            let bias_path = format!("{}_bias.json", layer_path);

            layer.weight.save_json(&weight_path)?;
            layer.bias.save_json(&bias_path)?;
        }

        println!(
            "Saved feed-forward network to {} ({} layers)",
            path,
            self.layers.len()
        );
        Ok(())
    }

    /// Load network parameters from JSON
    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 })
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Feed-Forward Network Example ===\n");

    demonstrate_network_creation();
    demonstrate_forward_pass();
    demonstrate_configurable_architectures();
    demonstrate_training_workflow()?;
    demonstrate_comprehensive_training()?;
    demonstrate_network_serialization()?;
    cleanup_temp_files()?;

    println!("\n=== Example completed successfully! ===");
    Ok(())
}

/// Demonstrate creating different network configurations
fn demonstrate_network_creation() {
    println!("--- Network Creation ---");

    // Default configuration
    let config = FeedForwardConfig::default();
    let network = FeedForwardNetwork::new(config.clone(), Some(42));

    println!("Default network configuration:");
    println!("  Input size: {}", config.input_size);
    println!("  Hidden sizes: {:?}", config.hidden_sizes);
    println!("  Output size: {}", config.output_size);
    println!("  Number of layers: {}", network.num_layers());
    println!("  Total parameters: {}", network.parameter_count());

    // Custom configurations
    let configs = [
        FeedForwardConfig {
            input_size: 2,
            hidden_sizes: vec![4],
            output_size: 1,
            use_bias: true,
        },
        FeedForwardConfig {
            input_size: 8,
            hidden_sizes: vec![16, 8, 4],
            output_size: 3,
            use_bias: true,
        },
        FeedForwardConfig {
            input_size: 10,
            hidden_sizes: vec![20, 15, 10, 5],
            output_size: 2,
            use_bias: true,
        },
    ];

    for (i, config) in configs.iter().enumerate() {
        let network = FeedForwardNetwork::new(config.clone(), Some(42 + i as u64));
        println!("\nCustom network {}:", i + 1);
        println!(
            "  Architecture: {} -> {:?} -> {}",
            config.input_size, config.hidden_sizes, config.output_size
        );
        println!("  Layers: {}", network.num_layers());
        println!("  Parameters: {}", network.parameter_count());
    }
}

/// Demonstrate forward pass through the network
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()
    );
}

/// Demonstrate different configurable architectures
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_linear_layer.rs (line 56)

    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 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/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_by_value().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(())
}

/// Demonstrate error handling for serialization operations
fn demonstrate_error_handling() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Error Handling ---");

    // Test loading non-existent file
    match Tensor::load_json("nonexistent_file.json") {
        Ok(_) => println!("Unexpected: Successfully loaded non-existent file"),
        Err(e) => println!("Expected error loading non-existent file: {}", e),
    }

    // Test loading with wrong format
    let tensor = Tensor::randn(vec![2, 2], Some(47));
    tensor.save_binary("temp_binary.bin")?;

    match Tensor::load_json("temp_binary.bin") {
        Ok(_) => println!("Unexpected: Successfully loaded binary as JSON"),
        Err(e) => println!("Expected error loading binary as JSON: {}", e),
    }

    // Test loading corrupted file
    fs::write("temp_invalid.json", "invalid json content")?;
    match Tensor::load_json("temp_invalid.json") {
        Ok(_) => println!("Unexpected: Successfully loaded invalid JSON"),
        Err(e) => println!("Expected error loading invalid JSON: {}", e),
    }

    println!("Error handling verification: PASSED");

    Ok(())
}

examples/optimizers/adam_configurations.rs (line 92)

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 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 iter(&self) -> TensorElementIterator<'_>

Create an iterator over tensor elements as view tensors

Each element becomes a Tensor of shape [1] that supports all tensor operations and gradient tracking. This is the main entry point for element-wise iteration with full tensor operation support.

The 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.

Returns

An iterator that yields view tensors for each element

Performance

• Zero-Copy Views: Each element is a view 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

Examples

Basic Element Operations

use train_station::Tensor;

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

// Use any std iterator method
let result: Tensor = tensor.iter()
    .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0)) // 2x + 1
    .filter(|elem| elem.value() > 3.0)                // Keep values > 3
    .collect();

assert_eq!(result.data(), &[5.0, 7.0]);

Advanced Iterator Chains

use train_station::Tensor;

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

// Chain with enumerate, zip, etc.
let indexed: Tensor = tensor.iter()
    .enumerate()
    .map(|(i, elem)| elem.add_scalar(i as f32))
    .collect();

assert_eq!(indexed.data(), &[1.0, 3.0, 5.0, 7.0, 9.0]);

Double-Ended Iteration

use train_station::Tensor;

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

// Use double-ended iterator
let reversed: Tensor = tensor.iter()
    .rev()
    .collect();

assert_eq!(reversed.data(), &[4.0, 3.0, 2.0, 1.0]);

Gradient Tracking

use train_station::Tensor;

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

let result: Tensor = tensor.iter()
    .map(|elem| elem.mul_scalar(2.0))
    .collect();

assert!(result.requires_grad());
assert_eq!(result.data(), &[2.0, 4.0]);

Examples found in repository

examples/iterators/element_iteration.rs (line 86)

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(())
}

examples/iterators/performance_optimization.rs (line 96)

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
        let start = Instant::now();
        let iterator_result: Tensor = tensor
            .iter()
            .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0))
            .collect();
        let iterator_time = start.elapsed();

        // Benchmark 3: Chained iterator operations
        let start = Instant::now();
        let _chained_result: Tensor = tensor
            .iter()
            .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();

        // Report results
        println!("  Direct operations: {:?}", direct_time);
        println!("  Iterator operations: {:?}", iterator_time);
        println!("  Chained operations: {:?}", chained_time);

        // Verify correctness
        assert_eq!(direct_result.data(), iterator_result.data());
        println!(
            "  Results match: {}",
            direct_result.data() == iterator_result.data()
        );

        // Performance ratio
        let ratio = iterator_time.as_nanos() as f64 / direct_time.as_nanos() as f64;
        println!("  Iterator/Direct ratio: {:.2}x", ratio);
    }

    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 (process in chunks)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();

    let mut streamed_result = Vec::new();
    for chunk_start in (0..size).step_by(chunk_size) {
        let chunk_end = (chunk_start + chunk_size).min(size);
        let chunk: Tensor = tensor
            .iter_range(chunk_start, chunk_end)
            .map(|elem| elem.pow_scalar(2.0).sqrt())
            .collect();
        streamed_result.extend(chunk.data().iter().cloned());
    }
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .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 + batch_size - 1) / 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 100)

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
    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
    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()
    );

    // 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 iter_range(&self, start: usize, end: usize) -> TensorElementIterator<'_>

Create an iterator over a range of elements

Creates an iterator that yields view tensors for elements in the specified range. The range is automatically clamped to valid tensor bounds for safety.

Arguments

• start - Starting index (inclusive)
• end - Ending index (exclusive)

Returns

An iterator that yields view tensors for elements in the specified range

Safety

The range is automatically clamped to valid tensor bounds:

• start is clamped to [0, tensor.size()]
• end is clamped to [start, tensor.size()]
• Empty ranges (start >= end) are handled gracefully

Performance

• O(1) Creation: Constant-time iterator initialization
• Bounds Checking: Automatic range validation and clamping
• Zero-Copy Views: Each element is a view sharing memory with source
• Memory Efficient: Minimal overhead for range iteration

Examples

Basic Range Iteration

use train_station::Tensor;

let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], vec![5]).unwrap();
let middle: Tensor = tensor.iter_range(1, 4)
    .map(|elem| elem.mul_scalar(2.0))
    .collect();

assert_eq!(middle.data(), &[4.0, 6.0, 8.0]);

Range with Operations

use train_station::Tensor;

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

// Apply complex operations to range
let result: Tensor = tensor.iter_range(0, 3)
    .enumerate()
    .map(|(i, elem)| elem.add_scalar(i as f32))
    .collect();

assert_eq!(result.data(), &[1.0, 3.0, 5.0]);

Out of Bounds Handling

use train_station::Tensor;

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

// Out of bounds range is clamped
let empty: Tensor = tensor.iter_range(5, 10).collect();
assert_eq!(empty.size(), 0);

// Partial out of bounds
let partial: Tensor = tensor.iter_range(1, 10).collect();
assert_eq!(partial.data(), &[2.0, 3.0]);

Examples found in repository

examples/iterators/performance_optimization.rs (line 154)

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 (process in chunks)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();

    let mut streamed_result = Vec::new();
    for chunk_start in (0..size).step_by(chunk_size) {
        let chunk_end = (chunk_start + chunk_size).min(size);
        let chunk: Tensor = tensor
            .iter_range(chunk_start, chunk_end)
            .map(|elem| elem.pow_scalar(2.0).sqrt())
            .collect();
        streamed_result.extend(chunk.data().iter().cloned());
    }
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .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 + batch_size - 1) / 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 350)

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
    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()
    );

    // 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 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/neural_networks/feedforward_network.rs (line 86)

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

examples/neural_networks/basic_linear_layer.rs (line 74)

    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/getting_started/tensor_basics.rs (line 96)

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 211)

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/advanced_patterns.rs (line 121)

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(())
}

pub fn add_scalar(&self, scalar: f32) -> Tensor

Broadcast addition with a scalar value.

Adds the scalar to every element: output[i] = self[i] + scalar

Arguments

• scalar - Value to add to each element

Returns

A new tensor with the scalar added to each element

Examples

Basic Scalar Addition

use train_station::Tensor;

let a = Tensor::from_slice(&[1.0, 2.0, 3.0], vec![3]).unwrap();
let b = a.add_scalar(10.0);
assert_eq!(b.shape().dims, vec![3]);
assert_eq!(b.get(&[0]), 11.0);
assert_eq!(b.get(&[1]), 12.0);
assert_eq!(b.get(&[2]), 13.0);

Multi-dimensional Scalar Addition

use train_station::Tensor;

let a = Tensor::ones(vec![2, 3]);
let b = a.add_scalar(5.0);
assert_eq!(b.shape().dims, vec![2, 3]);
assert_eq!(b.get(&[0, 0]), 6.0);
assert_eq!(b.get(&[1, 2]), 6.0);

Examples found in repository

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 99)

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(())
}

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()
    );
}

examples/iterators/performance_optimization.rs (line 90)

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
        let start = Instant::now();
        let iterator_result: Tensor = tensor
            .iter()
            .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0))
            .collect();
        let iterator_time = start.elapsed();

        // Benchmark 3: Chained iterator operations
        let start = Instant::now();
        let _chained_result: Tensor = tensor
            .iter()
            .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();

        // Report results
        println!("  Direct operations: {:?}", direct_time);
        println!("  Iterator operations: {:?}", iterator_time);
        println!("  Chained operations: {:?}", chained_time);

        // Verify correctness
        assert_eq!(direct_result.data(), iterator_result.data());
        println!(
            "  Results match: {}",
            direct_result.data() == iterator_result.data()
        );

        // Performance ratio
        let ratio = iterator_time.as_nanos() as f64 / direct_time.as_nanos() as f64;
        println!("  Iterator/Direct ratio: {:.2}x", ratio);
    }

    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 (process in chunks)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();

    let mut streamed_result = Vec::new();
    for chunk_start in (0..size).step_by(chunk_size) {
        let chunk_end = (chunk_start + chunk_size).min(size);
        let chunk: Tensor = tensor
            .iter_range(chunk_start, chunk_end)
            .map(|elem| elem.pow_scalar(2.0).sqrt())
            .collect();
        streamed_result.extend(chunk.data().iter().cloned());
    }
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .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 + batch_size - 1) / 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 223)

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
    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
    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()
    );

    // 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 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/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/iterators/advanced_patterns.rs (line 101)

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
    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
    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()
    );

    // 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/iterators/element_iteration.rs (line 224)

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/iterators/element_iteration.rs (line 226)

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 operation following NumPy semantics

Performs matrix multiplication between this tensor and another tensor with intelligent kernel selection based on matrix dimensions and hardware capabilities. The operation follows broadcasting rules and supports all common matrix multiplication patterns found in machine learning workloads.

Supported Operations

• 1D @ 1D: Dot product returning scalar tensor
• 1D @ 2D: Vector-matrix multiplication (v^T * M) returning 1D tensor
• 2D @ 1D: Matrix-vector multiplication (M * v) returning 1D tensor
• 2D @ 2D: Standard matrix multiplication with cache-optimized blocking
• ND @ ND: Batched matrix multiplication on last two dimensions with broadcasting

Performance Characteristics

The implementation automatically selects optimal kernels based on matrix dimensions:

• Small matrices (<64 elements): Direct computation with minimal overhead
• Medium matrices (64-256 elements): Cache-optimized blocking for L1/L2 cache
• Large matrices (256+ elements): Memory bandwidth optimized with hierarchical blocking
• AVX2 acceleration: 8x SIMD operations for compatible hardware
• Scalar fallbacks: Optimized scalar implementations for non-SIMD platforms

Automatic Differentiation

This operation supports automatic differentiation when either operand requires gradients. Gradient computation follows PyTorch semantics with proper accumulation and chain rule application through the gradtrack engine. Gradients are computed for both operands when requires_grad is set.

Arguments

• other - The tensor to multiply with (must have compatible dimensions)

Returns

A new tensor containing the matrix multiplication result with appropriate shape determined by broadcasting rules and matrix multiplication semantics

Panics

Panics if the inner dimensions don’t match for matrix multiplication:

• For 2D @ 2D: self.shape()[1] != other.shape()[0]
• For 1D @ 2D: self.shape()[0] != other.shape()[0]
• For 2D @ 1D: self.shape()[1] != other.shape()[0]
• For ND @ ND: Last two dimensions must be compatible 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); // [2, 2] @ [2, 2] -> [2, 2]

assert_eq!(result.shape().dims, vec![2, 2]);
assert_eq!(result.data(), &[19.0, 22.0, 43.0, 50.0]);

Vector-Matrix Multiplication

use train_station::Tensor;

let v = Tensor::from_slice(&[1.0, 2.0], vec![2]).unwrap();
let m = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
let result = v.matmul(&m); // [2] @ [2, 2] -> [2]

assert_eq!(result.shape().dims, vec![2]);
assert_eq!(result.data(), &[7.0, 10.0]); // 1*1+2*3, 1*2+2*4

Gradient Tracking

use train_station::Tensor;

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

let result = a.matmul(&b);
assert!(result.requires_grad());
assert_eq!(result.shape().dims, vec![2, 2]);

Thread Safety

This operation is thread-safe and can be used concurrently across multiple threads. The implementation uses immutable tensor references and thread-local gradtrack state.

Memory Safety

The implementation uses Tensor::new_uninitialized for performance-critical allocations and handles memory initialization safely through the kernel system. All unsafe operations are validated through comprehensive FFI testing against LibTorch reference implementation.

Examples found in repository

examples/neural_networks/feedforward_network.rs (line 85)

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

examples/neural_networks/basic_linear_layer.rs (line 72)

    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_by_value().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(())
}

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

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(())
}

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/neural_networks/feedforward_network.rs (line 72)

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

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

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

examples/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/getting_started/tensor_basics.rs (line 115)

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 99)

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(())
}

examples/iterators/performance_optimization.rs (line 90)

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
        let start = Instant::now();
        let iterator_result: Tensor = tensor
            .iter()
            .map(|elem| elem.mul_scalar(2.0).add_scalar(1.0))
            .collect();
        let iterator_time = start.elapsed();

        // Benchmark 3: Chained iterator operations
        let start = Instant::now();
        let _chained_result: Tensor = tensor
            .iter()
            .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();

        // Report results
        println!("  Direct operations: {:?}", direct_time);
        println!("  Iterator operations: {:?}", iterator_time);
        println!("  Chained operations: {:?}", chained_time);

        // Verify correctness
        assert_eq!(direct_result.data(), iterator_result.data());
        println!(
            "  Results match: {}",
            direct_result.data() == iterator_result.data()
        );

        // Performance ratio
        let ratio = iterator_time.as_nanos() as f64 / direct_time.as_nanos() as f64;
        println!("  Iterator/Direct ratio: {:.2}x", ratio);
    }

    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 (process in chunks)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();

    let mut streamed_result = Vec::new();
    for chunk_start in (0..size).step_by(chunk_size) {
        let chunk_end = (chunk_start + chunk_size).min(size);
        let chunk: Tensor = tensor
            .iter_range(chunk_start, chunk_end)
            .map(|elem| elem.pow_scalar(2.0).sqrt())
            .collect();
        streamed_result.extend(chunk.data().iter().cloned());
    }
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .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 + batch_size - 1) / 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 178)

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
    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
    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()
    );

    // 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 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/iterators/element_iteration.rs (line 122)

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(),
    })
}

examples/optimizers/learning_rate_scheduling.rs (line 344)

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 135)

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_by_value().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/feedforward_network.rs (line 479)

fn demonstrate_training_workflow() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Workflow ---");

    // Create a simple classification network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 3],
        output_size: 1,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(46));

    println!("Training network: 2 -> [4, 3] -> 1");

    // Create simple binary classification data: XOR problem
    let x_data = Tensor::from_slice(
        &[
            0.0, 0.0, // -> 0
            0.0, 1.0, // -> 1
            1.0, 0.0, // -> 1
            1.0, 1.0, // -> 0
        ],
        vec![4, 2],
    )
    .unwrap();

    let y_true = Tensor::from_slice(&[0.0, 1.0, 1.0, 0.0], vec![4, 1]).unwrap();

    println!("Training on XOR problem:");
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.1);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

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

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        losses.push(loss.value());

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

    // Test final model
    let final_predictions = network.forward_no_grad(&x_data);
    println!("\nFinal predictions vs targets:");
    for i in 0..4 {
        let pred = final_predictions.data()[i];
        let target = y_true.data()[i];
        let input_x = x_data.data()[i * 2];
        let input_y = x_data.data()[i * 2 + 1];
        println!(
            "  [{:.0}, {:.0}] -> pred: {:.3}, target: {:.0}, error: {:.3}",
            input_x,
            input_y,
            pred,
            target,
            (pred - target).abs()
        );
    }

    Ok(())
}

/// Demonstrate comprehensive training with 100+ steps
fn demonstrate_comprehensive_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Comprehensive Training (100+ Steps) ---");

    // Create a regression network
    let config = FeedForwardConfig {
        input_size: 3,
        hidden_sizes: vec![8, 6, 4],
        output_size: 2,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(47));

    println!("Network architecture: 3 -> [8, 6, 4] -> 2");
    println!("Total parameters: {}", network.parameter_count());

    // Create synthetic regression data
    // Target function: [y1, y2] = [x1 + 2*x2 - x3, x1*x2 + x3]
    let num_samples = 32;
    let mut x_vec = Vec::new();
    let mut y_vec = Vec::new();

    for i in 0..num_samples {
        let x1 = (i as f32 / num_samples as f32) * 2.0 - 1.0; // [-1, 1]
        let x2 = ((i * 2) as f32 / num_samples as f32) * 2.0 - 1.0;
        let x3 = ((i * 3) as f32 / num_samples as f32) * 2.0 - 1.0;

        let y1 = x1 + 2.0 * x2 - x3;
        let y2 = x1 * x2 + x3;

        x_vec.extend_from_slice(&[x1, x2, x3]);
        y_vec.extend_from_slice(&[y1, y2]);
    }

    let x_data = Tensor::from_slice(&x_vec, vec![num_samples, 3]).unwrap();
    let y_true = Tensor::from_slice(&y_vec, vec![num_samples, 2]).unwrap();

    println!("Training data:");
    println!("  {} samples", num_samples);
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer with learning rate scheduling
    let mut optimizer = Adam::with_learning_rate(0.01);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    // Comprehensive training loop (150 epochs)
    let num_epochs = 150;
    let mut losses = Vec::new();
    let mut best_loss = f32::INFINITY;
    let mut patience_counter = 0;
    let patience = 20;

    println!("Starting comprehensive training...");

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        let current_loss = loss.value();
        losses.push(current_loss);

        // Learning rate scheduling
        if epoch > 0 && epoch % 30 == 0 {
            let new_lr = optimizer.learning_rate() * 0.8;
            optimizer.set_learning_rate(new_lr);
            println!("  Reduced learning rate to {:.4}", new_lr);
        }

        // Early stopping logic
        if current_loss < best_loss {
            best_loss = current_loss;
            patience_counter = 0;
        } else {
            patience_counter += 1;
        }

        // Print progress
        if epoch % 25 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:3}: Loss = {:.6}, LR = {:.4}, Best = {:.6}",
                epoch,
                current_loss,
                optimizer.learning_rate(),
                best_loss
            );
        }

        // Early stopping
        if patience_counter >= patience && epoch > 50 {
            println!("Early stopping at epoch {} (patience exceeded)", epoch);
            break;
        }
    }

    // Final evaluation
    let final_predictions = network.forward_no_grad(&x_data);

    // Compute final metrics
    let final_loss = losses[losses.len() - 1];
    let initial_loss = losses[0];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining completed!");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Best loss: {:.6}", best_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final learning rate: {:.4}", optimizer.learning_rate());

    // Sample predictions analysis
    println!("\nSample predictions (first 5):");
    for i in 0..5.min(num_samples) {
        let pred1 = final_predictions.data()[i * 2];
        let pred2 = final_predictions.data()[i * 2 + 1];
        let true1 = y_true.data()[i * 2];
        let true2 = y_true.data()[i * 2 + 1];

        println!(
            "  Sample {}: pred=[{:.3}, {:.3}], true=[{:.3}, {:.3}], error=[{:.3}, {:.3}]",
            i + 1,
            pred1,
            pred2,
            true1,
            true2,
            (pred1 - true1).abs(),
            (pred2 - true2).abs()
        );
    }

    Ok(())
}

/// Demonstrate network serialization
fn demonstrate_network_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Network Serialization ---");

    // Create and train a network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 2],
        output_size: 1,
        use_bias: true,
    };
    let mut original_network = FeedForwardNetwork::new(config.clone(), Some(48));

    // Quick training
    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_network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    for _ in 0..20 {
        let y_pred = original_network.forward(&x_data);
        let mut loss = (y_pred.sub_tensor(&y_true)).pow_scalar(2.0).mean();
        loss.backward(None);

        let mut params = original_network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);
    }

    // Test original network
    let test_input = Tensor::from_slice(&[1.0, 1.0], vec![1, 2]).unwrap();
    let original_output = original_network.forward_no_grad(&test_input);

    println!("Original network output: {:?}", original_output.data());

    // Save network
    original_network.save_json("temp_feedforward_network")?;

    // Load network
    let loaded_network = FeedForwardNetwork::load_json("temp_feedforward_network", config)?;
    let loaded_output = loaded_network.forward_no_grad(&test_input);

    println!("Loaded network output: {:?}", loaded_output.data());

    // Verify consistency
    let match_check = original_output
        .data()
        .iter()
        .zip(loaded_output.data().iter())
        .all(|(a, b)| (a - b).abs() < 1e-6);

    println!(
        "Serialization verification: {}",
        if match_check { "PASSED" } else { "FAILED" }
    );

    Ok(())
}

examples/iterators/performance_optimization.rs (line 155)

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 (process in chunks)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();

    let mut streamed_result = Vec::new();
    for chunk_start in (0..size).step_by(chunk_size) {
        let chunk_end = (chunk_start + chunk_size).min(size);
        let chunk: Tensor = tensor
            .iter_range(chunk_start, chunk_end)
            .map(|elem| elem.pow_scalar(2.0).sqrt())
            .collect();
        streamed_result.extend(chunk.data().iter().cloned());
    }
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .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 + batch_size - 1) / 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(())
}

Additional examples can be found in:

• examples/iterators/advanced_patterns.rs
• examples/neural_networks/basic_linear_layer.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 48)

    pub fn forward(input: &Tensor) -> Tensor {
        input.relu()
    }

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

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

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/iterators/performance_optimization.rs (line 155)

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 (process in chunks)
    println!("\nPattern 1: Streaming Processing");
    let chunk_size = 1000;
    let start = Instant::now();

    let mut streamed_result = Vec::new();
    for chunk_start in (0..size).step_by(chunk_size) {
        let chunk_end = (chunk_start + chunk_size).min(size);
        let chunk: Tensor = tensor
            .iter_range(chunk_start, chunk_end)
            .map(|elem| elem.pow_scalar(2.0).sqrt())
            .collect();
        streamed_result.extend(chunk.data().iter().cloned());
    }
    let streamed_time = start.elapsed();

    // Pattern 2: Full processing
    let start = Instant::now();
    let _full_result: Tensor = tensor
        .iter()
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .take(1000) // Only process first 1000 elements
        .map(|elem| elem.pow_scalar(2.0).sqrt())
        .collect();
    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()
        .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 + batch_size - 1) / 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 178)

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/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());
}

examples/neural_networks/feedforward_network.rs (line 478)

fn demonstrate_training_workflow() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Workflow ---");

    // Create a simple classification network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 3],
        output_size: 1,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(46));

    println!("Training network: 2 -> [4, 3] -> 1");

    // Create simple binary classification data: XOR problem
    let x_data = Tensor::from_slice(
        &[
            0.0, 0.0, // -> 0
            0.0, 1.0, // -> 1
            1.0, 0.0, // -> 1
            1.0, 1.0, // -> 0
        ],
        vec![4, 2],
    )
    .unwrap();

    let y_true = Tensor::from_slice(&[0.0, 1.0, 1.0, 0.0], vec![4, 1]).unwrap();

    println!("Training on XOR problem:");
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.1);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

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

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        losses.push(loss.value());

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

    // Test final model
    let final_predictions = network.forward_no_grad(&x_data);
    println!("\nFinal predictions vs targets:");
    for i in 0..4 {
        let pred = final_predictions.data()[i];
        let target = y_true.data()[i];
        let input_x = x_data.data()[i * 2];
        let input_y = x_data.data()[i * 2 + 1];
        println!(
            "  [{:.0}, {:.0}] -> pred: {:.3}, target: {:.0}, error: {:.3}",
            input_x,
            input_y,
            pred,
            target,
            (pred - target).abs()
        );
    }

    Ok(())
}

/// Demonstrate comprehensive training with 100+ steps
fn demonstrate_comprehensive_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Comprehensive Training (100+ Steps) ---");

    // Create a regression network
    let config = FeedForwardConfig {
        input_size: 3,
        hidden_sizes: vec![8, 6, 4],
        output_size: 2,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(47));

    println!("Network architecture: 3 -> [8, 6, 4] -> 2");
    println!("Total parameters: {}", network.parameter_count());

    // Create synthetic regression data
    // Target function: [y1, y2] = [x1 + 2*x2 - x3, x1*x2 + x3]
    let num_samples = 32;
    let mut x_vec = Vec::new();
    let mut y_vec = Vec::new();

    for i in 0..num_samples {
        let x1 = (i as f32 / num_samples as f32) * 2.0 - 1.0; // [-1, 1]
        let x2 = ((i * 2) as f32 / num_samples as f32) * 2.0 - 1.0;
        let x3 = ((i * 3) as f32 / num_samples as f32) * 2.0 - 1.0;

        let y1 = x1 + 2.0 * x2 - x3;
        let y2 = x1 * x2 + x3;

        x_vec.extend_from_slice(&[x1, x2, x3]);
        y_vec.extend_from_slice(&[y1, y2]);
    }

    let x_data = Tensor::from_slice(&x_vec, vec![num_samples, 3]).unwrap();
    let y_true = Tensor::from_slice(&y_vec, vec![num_samples, 2]).unwrap();

    println!("Training data:");
    println!("  {} samples", num_samples);
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer with learning rate scheduling
    let mut optimizer = Adam::with_learning_rate(0.01);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    // Comprehensive training loop (150 epochs)
    let num_epochs = 150;
    let mut losses = Vec::new();
    let mut best_loss = f32::INFINITY;
    let mut patience_counter = 0;
    let patience = 20;

    println!("Starting comprehensive training...");

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        let current_loss = loss.value();
        losses.push(current_loss);

        // Learning rate scheduling
        if epoch > 0 && epoch % 30 == 0 {
            let new_lr = optimizer.learning_rate() * 0.8;
            optimizer.set_learning_rate(new_lr);
            println!("  Reduced learning rate to {:.4}", new_lr);
        }

        // Early stopping logic
        if current_loss < best_loss {
            best_loss = current_loss;
            patience_counter = 0;
        } else {
            patience_counter += 1;
        }

        // Print progress
        if epoch % 25 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:3}: Loss = {:.6}, LR = {:.4}, Best = {:.6}",
                epoch,
                current_loss,
                optimizer.learning_rate(),
                best_loss
            );
        }

        // Early stopping
        if patience_counter >= patience && epoch > 50 {
            println!("Early stopping at epoch {} (patience exceeded)", epoch);
            break;
        }
    }

    // Final evaluation
    let final_predictions = network.forward_no_grad(&x_data);

    // Compute final metrics
    let final_loss = losses[losses.len() - 1];
    let initial_loss = losses[0];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining completed!");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Best loss: {:.6}", best_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final learning rate: {:.4}", optimizer.learning_rate());

    // Sample predictions analysis
    println!("\nSample predictions (first 5):");
    for i in 0..5.min(num_samples) {
        let pred1 = final_predictions.data()[i * 2];
        let pred2 = final_predictions.data()[i * 2 + 1];
        let true1 = y_true.data()[i * 2];
        let true2 = y_true.data()[i * 2 + 1];

        println!(
            "  Sample {}: pred=[{:.3}, {:.3}], true=[{:.3}, {:.3}], error=[{:.3}, {:.3}]",
            i + 1,
            pred1,
            pred2,
            true1,
            true2,
            (pred1 - true1).abs(),
            (pred2 - true2).abs()
        );
    }

    Ok(())
}

/// Demonstrate network serialization
fn demonstrate_network_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Network Serialization ---");

    // Create and train a network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 2],
        output_size: 1,
        use_bias: true,
    };
    let mut original_network = FeedForwardNetwork::new(config.clone(), Some(48));

    // Quick training
    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_network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    for _ in 0..20 {
        let y_pred = original_network.forward(&x_data);
        let mut loss = (y_pred.sub_tensor(&y_true)).pow_scalar(2.0).mean();
        loss.backward(None);

        let mut params = original_network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);
    }

    // Test original network
    let test_input = Tensor::from_slice(&[1.0, 1.0], vec![1, 2]).unwrap();
    let original_output = original_network.forward_no_grad(&test_input);

    println!("Original network output: {:?}", original_output.data());

    // Save network
    original_network.save_json("temp_feedforward_network")?;

    // Load network
    let loaded_network = FeedForwardNetwork::load_json("temp_feedforward_network", config)?;
    let loaded_output = loaded_network.forward_no_grad(&test_input);

    println!("Loaded network output: {:?}", loaded_output.data());

    // Verify consistency
    let match_check = original_output
        .data()
        .iter()
        .zip(loaded_output.data().iter())
        .all(|(a, b)| (a - b).abs() < 1e-6);

    println!(
        "Serialization verification: {}",
        if match_check { "PASSED" } else { "FAILED" }
    );

    Ok(())
}

examples/neural_networks/basic_linear_layer.rs (line 269)

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(())
}

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/iterators/advanced_patterns.rs (line 101)

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
    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
    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()
    );

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

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.

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/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()
    );
}

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(),
    })
}

examples/optimizers/learning_rate_scheduling.rs (line 344)

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 135)

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_by_value().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/feedforward_network.rs (line 479)

fn demonstrate_training_workflow() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Training Workflow ---");

    // Create a simple classification network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 3],
        output_size: 1,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(46));

    println!("Training network: 2 -> [4, 3] -> 1");

    // Create simple binary classification data: XOR problem
    let x_data = Tensor::from_slice(
        &[
            0.0, 0.0, // -> 0
            0.0, 1.0, // -> 1
            1.0, 0.0, // -> 1
            1.0, 1.0, // -> 0
        ],
        vec![4, 2],
    )
    .unwrap();

    let y_true = Tensor::from_slice(&[0.0, 1.0, 1.0, 0.0], vec![4, 1]).unwrap();

    println!("Training on XOR problem:");
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer
    let mut optimizer = Adam::with_learning_rate(0.1);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

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

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        losses.push(loss.value());

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

    // Test final model
    let final_predictions = network.forward_no_grad(&x_data);
    println!("\nFinal predictions vs targets:");
    for i in 0..4 {
        let pred = final_predictions.data()[i];
        let target = y_true.data()[i];
        let input_x = x_data.data()[i * 2];
        let input_y = x_data.data()[i * 2 + 1];
        println!(
            "  [{:.0}, {:.0}] -> pred: {:.3}, target: {:.0}, error: {:.3}",
            input_x,
            input_y,
            pred,
            target,
            (pred - target).abs()
        );
    }

    Ok(())
}

/// Demonstrate comprehensive training with 100+ steps
fn demonstrate_comprehensive_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Comprehensive Training (100+ Steps) ---");

    // Create a regression network
    let config = FeedForwardConfig {
        input_size: 3,
        hidden_sizes: vec![8, 6, 4],
        output_size: 2,
        use_bias: true,
    };
    let mut network = FeedForwardNetwork::new(config, Some(47));

    println!("Network architecture: 3 -> [8, 6, 4] -> 2");
    println!("Total parameters: {}", network.parameter_count());

    // Create synthetic regression data
    // Target function: [y1, y2] = [x1 + 2*x2 - x3, x1*x2 + x3]
    let num_samples = 32;
    let mut x_vec = Vec::new();
    let mut y_vec = Vec::new();

    for i in 0..num_samples {
        let x1 = (i as f32 / num_samples as f32) * 2.0 - 1.0; // [-1, 1]
        let x2 = ((i * 2) as f32 / num_samples as f32) * 2.0 - 1.0;
        let x3 = ((i * 3) as f32 / num_samples as f32) * 2.0 - 1.0;

        let y1 = x1 + 2.0 * x2 - x3;
        let y2 = x1 * x2 + x3;

        x_vec.extend_from_slice(&[x1, x2, x3]);
        y_vec.extend_from_slice(&[y1, y2]);
    }

    let x_data = Tensor::from_slice(&x_vec, vec![num_samples, 3]).unwrap();
    let y_true = Tensor::from_slice(&y_vec, vec![num_samples, 2]).unwrap();

    println!("Training data:");
    println!("  {} samples", num_samples);
    println!("  Input shape: {:?}", x_data.shape().dims);
    println!("  Target shape: {:?}", y_true.shape().dims);

    // Create optimizer with learning rate scheduling
    let mut optimizer = Adam::with_learning_rate(0.01);
    let params = network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    // Comprehensive training loop (150 epochs)
    let num_epochs = 150;
    let mut losses = Vec::new();
    let mut best_loss = f32::INFINITY;
    let mut patience_counter = 0;
    let patience = 20;

    println!("Starting comprehensive training...");

    for epoch in 0..num_epochs {
        // Forward pass
        let y_pred = network.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 and zero grad
        let mut params = network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);

        let current_loss = loss.value();
        losses.push(current_loss);

        // Learning rate scheduling
        if epoch > 0 && epoch % 30 == 0 {
            let new_lr = optimizer.learning_rate() * 0.8;
            optimizer.set_learning_rate(new_lr);
            println!("  Reduced learning rate to {:.4}", new_lr);
        }

        // Early stopping logic
        if current_loss < best_loss {
            best_loss = current_loss;
            patience_counter = 0;
        } else {
            patience_counter += 1;
        }

        // Print progress
        if epoch % 25 == 0 || epoch == num_epochs - 1 {
            println!(
                "Epoch {:3}: Loss = {:.6}, LR = {:.4}, Best = {:.6}",
                epoch,
                current_loss,
                optimizer.learning_rate(),
                best_loss
            );
        }

        // Early stopping
        if patience_counter >= patience && epoch > 50 {
            println!("Early stopping at epoch {} (patience exceeded)", epoch);
            break;
        }
    }

    // Final evaluation
    let final_predictions = network.forward_no_grad(&x_data);

    // Compute final metrics
    let final_loss = losses[losses.len() - 1];
    let initial_loss = losses[0];
    let loss_reduction = (initial_loss - final_loss) / initial_loss * 100.0;

    println!("\nTraining completed!");
    println!("  Initial loss: {:.6}", initial_loss);
    println!("  Final loss: {:.6}", final_loss);
    println!("  Best loss: {:.6}", best_loss);
    println!("  Loss reduction: {:.1}%", loss_reduction);
    println!("  Final learning rate: {:.4}", optimizer.learning_rate());

    // Sample predictions analysis
    println!("\nSample predictions (first 5):");
    for i in 0..5.min(num_samples) {
        let pred1 = final_predictions.data()[i * 2];
        let pred2 = final_predictions.data()[i * 2 + 1];
        let true1 = y_true.data()[i * 2];
        let true2 = y_true.data()[i * 2 + 1];

        println!(
            "  Sample {}: pred=[{:.3}, {:.3}], true=[{:.3}, {:.3}], error=[{:.3}, {:.3}]",
            i + 1,
            pred1,
            pred2,
            true1,
            true2,
            (pred1 - true1).abs(),
            (pred2 - true2).abs()
        );
    }

    Ok(())
}

/// Demonstrate network serialization
fn demonstrate_network_serialization() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n--- Network Serialization ---");

    // Create and train a network
    let config = FeedForwardConfig {
        input_size: 2,
        hidden_sizes: vec![4, 2],
        output_size: 1,
        use_bias: true,
    };
    let mut original_network = FeedForwardNetwork::new(config.clone(), Some(48));

    // Quick training
    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_network.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    for _ in 0..20 {
        let y_pred = original_network.forward(&x_data);
        let mut loss = (y_pred.sub_tensor(&y_true)).pow_scalar(2.0).mean();
        loss.backward(None);

        let mut params = original_network.parameters();
        optimizer.step(&mut params);
        optimizer.zero_grad(&mut params);
    }

    // Test original network
    let test_input = Tensor::from_slice(&[1.0, 1.0], vec![1, 2]).unwrap();
    let original_output = original_network.forward_no_grad(&test_input);

    println!("Original network output: {:?}", original_output.data());

    // Save network
    original_network.save_json("temp_feedforward_network")?;

    // Load network
    let loaded_network = FeedForwardNetwork::load_json("temp_feedforward_network", config)?;
    let loaded_output = loaded_network.forward_no_grad(&test_input);

    println!("Loaded network output: {:?}", loaded_output.data());

    // Verify consistency
    let match_check = original_output
        .data()
        .iter()
        .zip(loaded_output.data().iter())
        .all(|(a, b)| (a - b).abs() < 1e-6);

    println!(
        "Serialization verification: {}",
        if match_check { "PASSED" } else { "FAILED" }
    );

    Ok(())
}

examples/iterators/advanced_patterns.rs (line 97)

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
    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
    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()
    );

    // 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(())
}

Additional examples can be found in:

• examples/neural_networks/basic_linear_layer.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_by_value().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 98)

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
    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
    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()
    );

    // 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_by_value().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/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 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/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 387)

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_by_value().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.

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());

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

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.

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]

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.

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 Drop for Tensor

fn drop(&mut self)

Frees the tensor’s memory when it goes out of scope

This ensures proper cleanup and prevents memory leaks.

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<'a> IntoIterator for &'a Tensor

type Item = Tensor

The type of the elements being iterated over.

type IntoIter = TensorElementIterator<'a>

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