Struct Adam 

pub struct Adam { /* private fields */ }

Adam optimizer for neural network parameter optimization

Implements the Adam optimization algorithm with PyTorch-compatible interface. Provides adaptive learning rates with momentum for efficient training of neural networks. The optimizer maintains per-parameter state for momentum and velocity estimates, enabling adaptive learning rates that improve convergence across diverse architectures.

Usage Pattern

The optimizer uses ID-based parameter linking for maximum flexibility and thread safety:

• Parameters are linked to the optimizer via add_parameter or add_parameters
• The step method takes mutable references to parameters for thread-safe updates
• Parameter states are maintained by tensor ID, allowing for dynamic parameter management
• Supports serialization and deserialization with parameter re-linking

Dynamic Parameter Management

Parameters can be added, removed, or re-linked at runtime:

• add_parameter: Link a single parameter
• add_parameters: Link multiple parameters at once
• unlink_parameter: Remove parameter state by ID
• clear_states: Remove all parameter states
• is_parameter_linked: Check if a parameter is linked

Serialization Support

The optimizer supports full serialization and deserialization with state preservation:

• Parameter states are saved with their shapes and insertion order for validation
• After deserialization, use relink_parameters to restore saved states to new tensors
• Parameters must be re-linked in the same chronological order they were originally added
• Shape validation ensures consistency between saved and current parameters

Features

• ID-Based Parameter Linking: Dynamic parameter management via tensor IDs
• Thread-Safe Step Method: Takes mutable references for safe concurrent access
• Per-Parameter State: Each parameter maintains its own momentum and velocity buffers
• Bias Correction: Automatically corrects initialization bias in moment estimates
• Weight Decay: Optional L2 regularization with efficient implementation
• AMSGrad Support: Optional AMSGrad variant for improved convergence stability
• SIMD Optimization: AVX2-optimized updates for maximum performance
• Full Serialization: Complete state persistence and restoration

Thread Safety

This type is thread-safe and can be shared between threads. The step method takes mutable references to parameters, ensuring exclusive access during updates.

Implementations

impl Adam

pub fn saved_parameter_count(&self) -> usize

Get the number of saved parameter states for checkpoint validation

This method returns the count of parameter states currently stored in the optimizer, which is essential for validating checkpoint integrity and ensuring proper parameter re-linking after deserialization. The count includes all parameters that have been linked to the optimizer and have accumulated optimization state.

Returns

Number of parameter states currently stored in the optimizer

Usage Patterns

Checkpoint Validation

After deserializing an optimizer, this method helps verify that the expected number of parameters were saved and can guide the re-linking process.

Training Resumption

When resuming training, compare this count with the number of parameters in your model to ensure checkpoint compatibility.

State Management

Use this method to monitor optimizer state growth and memory usage during training with dynamic parameter addition.

Examples

use train_station::{Tensor, optimizers::Adam};
use train_station::serialization::Serializable;

let weight = Tensor::ones(vec![10, 5]).with_requires_grad();
let bias = Tensor::zeros(vec![5]).with_requires_grad();
let mut optimizer = Adam::new();
optimizer.add_parameter(&weight);
optimizer.add_parameter(&bias);

// Check parameter count before serialization
assert_eq!(optimizer.saved_parameter_count(), 2);

// Serialize and deserialize
let json = optimizer.to_json().unwrap();
let loaded_optimizer = Adam::from_json(&json).unwrap();

// Verify parameter count is preserved
assert_eq!(loaded_optimizer.saved_parameter_count(), 2);

Performance

• Time Complexity: O(1) - Direct access to internal state count
• Memory Usage: No additional memory allocation
• Thread Safety: Safe to call from multiple threads concurrently

impl Adam

pub fn new() -> Self

Create a new Adam optimizer with default configuration

Initializes an Adam optimizer with PyTorch-compatible default hyperparameters. Parameters must be linked separately using add_parameter or add_parameters.

Returns

A new Adam optimizer instance with default hyperparameters

Examples found in repository

examples/getting_started/optimizer_basics.rs (line 67)

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

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

pub fn with_config(config: AdamConfig) -> Self

Create a new Adam optimizer with custom configuration

Allows full control over all Adam hyperparameters for specialized training scenarios such as fine-tuning, transfer learning, or research applications. Parameters must be linked separately using add_parameter or add_parameters.

Arguments

• config - Adam configuration with custom hyperparameters

Returns

A new Adam optimizer instance with the specified configuration

Examples found in repository

examples/getting_started/optimizer_basics.rs (line 91)

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/getting_started/serialization_basics.rs (line 125)

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

examples/optimizers/adam_configurations.rs (line 336)

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

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

pub fn with_learning_rate(learning_rate: f32) -> Self

Create a new Adam optimizer with custom learning rate

A convenience constructor that allows setting only the learning rate while using default values for all other hyperparameters. Parameters must be linked separately using add_parameter or add_parameters.

Arguments

• learning_rate - Learning rate for optimization

Returns

A new Adam optimizer instance with the specified learning rate and default values for all other hyperparameters

Examples found in repository

examples/neural_networks/basic_encoder.rs (line 89)

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

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

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

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

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

examples/neural_networks/basic_decoder.rs (line 100)

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

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

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

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

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

examples/neural_networks/basic_transformer.rs (line 155)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

examples/neural_networks/multi_head_attention.rs (line 201)

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

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

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

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

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

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

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

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

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

examples/optimizers/learning_rate_scheduling.rs (line 332)

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

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

    // Create and train a simple layer
    let mut original_layer = LinearLayer::new(2, 1, Some(47));

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

    let mut optimizer = Adam::with_learning_rate(0.01);
    let params = original_layer.parameters();
    for param in &params {
        optimizer.add_parameter(param);
    }

    // Train for a few epochs
    for _ in 0..10 {
        let y_pred = original_layer.forward(&x_data);
        let mut loss = (y_pred.sub_tensor(&y_true)).pow_scalar(2.0).mean();
        loss.backward(None);

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

    println!("Original layer trained");
    println!("  Weight: {:?}", original_layer.weight.data());
    println!("  Bias: {:?}", original_layer.bias.data());

    // Save layer
    original_layer.save_json("temp_linear_layer")?;

    // Load layer
    let loaded_layer = LinearLayer::load_json("temp_linear_layer", 2, 1)?;

    println!("Loaded layer");
    println!("  Weight: {:?}", loaded_layer.weight.data());
    println!("  Bias: {:?}", loaded_layer.bias.data());

    // Verify consistency
    let test_input = Tensor::from_slice(&[1.0, 1.0], vec![1, 2]).unwrap();
    let original_output = original_layer.forward_no_grad(&test_input);
    let loaded_output = loaded_layer.forward_no_grad(&test_input);

    println!("Consistency check:");
    println!("  Original output: {:?}", original_output.data());
    println!("  Loaded output: {:?}", loaded_output.data());
    println!(
        "  Match: {}",
        original_output
            .data()
            .iter()
            .zip(loaded_output.data().iter())
            .all(|(a, b)| (a - b).abs() < 1e-6)
    );

    println!("Serialization verification: PASSED");

    Ok(())
}

Additional examples can be found in:

• examples/getting_started/serialization_basics.rs
• examples/getting_started/optimizer_basics.rs
• examples/supervised_training/../neural_networks/feedforward_network.rs
• examples/supervised_training/supervised_bce.rs
• examples/supervised_training/supervised_classification.rs
• examples/supervised_training/supervised_regression.rs
• examples/RL_training/dqn.rs
• examples/RL_training/ppo_discrete.rs
• examples/RL_training/ppo_continuous.rs
• examples/RL_training/td3.rs

pub fn add_parameter(&mut self, parameter: &Tensor)

Add a single parameter to the optimizer

Links a parameter to the optimizer by creating a new parameter state indexed by the tensor’s ID. The parameter must have requires_grad set to true.

Arguments

• parameter - Reference to the tensor to link

Panics

Panics if the parameter does not have requires_grad set to true

Examples found in repository

examples/neural_networks/basic_encoder.rs (line 92)

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

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

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

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

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

examples/neural_networks/basic_decoder.rs (line 103)

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

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

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

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

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

examples/neural_networks/basic_transformer.rs (line 159)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

examples/getting_started/optimizer_basics.rs (line 74)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        // Backward pass
        loss.backward(None);

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

        losses.push(loss.value());

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

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

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

    Ok(())
}

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

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

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

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

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

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

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

        // Backward pass
        loss.backward(None);

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

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

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

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

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

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

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

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

    Ok(())
}

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

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

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

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

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

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

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

        // Backward pass
        loss.backward(None);

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

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

        losses.push(loss.value());

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

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

    Ok(())
}

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

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

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

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

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

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

        // Backward pass
        loss.backward(None);

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

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

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

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

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

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

    Ok(())
}

examples/getting_started/serialization_basics.rs (line 126)

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

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

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

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

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

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

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

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

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

    Ok(())
}

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

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

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

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

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

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

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

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

    Ok(())
}

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

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

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

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

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

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

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

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

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

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

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

    println!("Checkpointing verification: PASSED");

    Ok(())
}

examples/neural_networks/multi_head_attention.rs (line 204)

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

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

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

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

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

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

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

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

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

Additional examples can be found in:

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

pub fn add_parameters(&mut self, parameters: &[&Tensor])

Add multiple parameters to the optimizer

Links multiple parameters to the optimizer by creating parameter states indexed by each tensor’s ID. All parameters must have requires_grad set to true.

Arguments

• parameters - Slice of references to tensors to link

Panics

Panics if any parameter does not have requires_grad set to true

pub fn unlink_parameter(&mut self, parameter: &Tensor) -> bool

Remove a parameter from the optimizer

Unlinks a parameter by removing its state from the optimizer. The parameter ID is used for identification.

Arguments

• parameter - Reference to the tensor to unlink

Returns

True if the parameter was linked and removed, false if it was not linked

pub fn clear_states(&mut self)

Remove all parameter states from the optimizer

Clears all parameter states, effectively unlinking all parameters. This is useful for resetting the optimizer or preparing for parameter re-linking.

pub fn is_parameter_linked(&self, parameter: &Tensor) -> bool

Check if a parameter is linked to the optimizer

Returns true if the parameter has an associated state in the optimizer.

Arguments

• parameter - Reference to the tensor to check

Returns

True if the parameter is linked, false otherwise

pub fn parameter_count(&self) -> usize

Get the number of linked parameters

Returns the count of parameters currently linked to the optimizer.

Returns

Number of linked parameters

Examples found in repository

examples/getting_started/optimizer_basics.rs (line 78)

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/getting_started/serialization_basics.rs (line 131)

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

pub fn relink_parameters( &mut self, parameters: &[&Tensor], ) -> Result<(), String>

Re-link parameters to saved optimizer states in chronological order

After deserializing an optimizer, use this method to restore saved parameter states to new tensors. Parameters must be provided in the same chronological order they were originally added to the optimizer. Shape validation ensures parameter compatibility.

Arguments

• parameters - Slice of parameter references in chronological order

Returns

Result indicating success or failure with detailed error message

Panics

Panics if any parameter does not have requires_grad set to true

pub fn config(&self) -> &AdamConfig

Get the current optimizer configuration

Returns a reference to the current configuration, allowing inspection of all hyperparameters without modification.

Returns

Reference to the current Adam configuration

Trait Implementations

impl Default for Adam

fn default() -> Self

Returns the “default value” for a type.

impl FromFieldValue for Adam

fn from_field_value( value: FieldValue, field_name: &str, ) -> SerializationResult<Self>

Create Adam from FieldValue

Arguments

• value - FieldValue containing optimizer data
• field_name - Name of the field being deserialized (for error messages)

Returns

Reconstructed Adam instance or error if deserialization fails

impl Optimizer for Adam

fn step(&mut self, parameters: &mut [&mut Tensor])

Perform a single optimization step

Updates all provided parameters based on their accumulated gradients using the Adam algorithm. Each parameter is updated according to the Adam update rule with bias correction and optional AMSGrad variant if enabled. All parameters must be linked to the optimizer before calling this method.

Arguments

• parameters - Mutable slice of parameter references to update

Thread Safety

This method is thread-safe as it takes mutable references to parameters, ensuring exclusive access during updates.

Performance

• Uses SIMD optimization (AVX2) when available for 8x vectorization
• Processes parameters in sequence for optimal cache usage
• Maintains per-parameter state for momentum and velocity estimates

Panics

Panics if any parameter is not linked to the optimizer

fn zero_grad(&mut self, parameters: &mut [&mut Tensor])

Zero out all parameter gradients

Clears accumulated gradients for all provided parameters. This should be called before each backward pass to prevent gradient accumulation across multiple forward/backward passes. Also clears the global autograd gradient map.

Arguments

• parameters - Mutable slice of parameter references to clear gradients for

Performance

• Efficiently clears gradients using optimized tensor operations
• Clears both per-tensor gradients and global autograd state
• Thread-safe as it takes mutable references to parameters

fn learning_rate(&self) -> f32

Get the current learning rate

Returns the current learning rate used for parameter updates.

Returns

Current learning rate as f32

fn set_learning_rate(&mut self, lr: f32)

Set the learning rate for all parameters

Updates the learning rate for all parameters in the optimizer. This allows dynamic learning rate scheduling during training.

Arguments

• lr - New learning rate value

impl Serializable for Adam

fn to_json(&self) -> SerializationResult<String>

Serialize the Adam optimizer to JSON format

This method converts the Adam optimizer into a human-readable JSON string representation that includes all optimizer state, configuration, parameter states, and step counts. The JSON format is suitable for debugging, configuration files, and cross-language interoperability.

Returns

JSON string representation of the optimizer on success, or SerializationError on failure

Examples

use train_station::{Tensor, optimizers::Adam};
use train_station::serialization::Serializable;

let weight = Tensor::ones(vec![2, 3]).with_requires_grad();
let mut optimizer = Adam::new();
optimizer.add_parameter(&weight);

let json = optimizer.to_json().unwrap();
assert!(!json.is_empty());

fn from_json(json: &str) -> SerializationResult<Self>

Deserialize an Adam optimizer from JSON format

This method parses a JSON string and reconstructs an Adam optimizer with all saved state. Parameters must be re-linked after deserialization using add_parameter or relink_parameters.

Arguments

• json - JSON string containing serialized optimizer

Returns

The deserialized optimizer on success, or SerializationError on failure

Examples

use train_station::{Tensor, optimizers::Adam};
use train_station::serialization::Serializable;

let weight = Tensor::ones(vec![2, 3]).with_requires_grad();
let mut optimizer = Adam::new();
optimizer.add_parameter(&weight);

let json = optimizer.to_json().unwrap();
let loaded_optimizer = Adam::from_json(&json).unwrap();
assert_eq!(loaded_optimizer.saved_parameter_count(), 1);

fn to_binary(&self) -> SerializationResult<Vec<u8>>

Serialize the Adam optimizer to binary format

This method converts the optimizer into a compact binary representation optimized for storage and transmission. The binary format provides maximum performance and minimal file sizes compared to JSON.

Returns

Binary representation of the optimizer on success, or SerializationError on failure

Examples

use train_station::{Tensor, optimizers::Adam};
use train_station::serialization::Serializable;

let weight = Tensor::ones(vec![2, 3]).with_requires_grad();
let mut optimizer = Adam::new();
optimizer.add_parameter(&weight);

let binary = optimizer.to_binary().unwrap();
assert!(!binary.is_empty());

fn from_binary(data: &[u8]) -> SerializationResult<Self>

Deserialize an Adam optimizer from binary format

This method parses binary data and reconstructs an Adam optimizer with all saved state. Parameters must be re-linked after deserialization using add_parameter or relink_parameters.

Arguments

• data - Binary data containing serialized optimizer

Returns

The deserialized optimizer on success, or SerializationError on failure

Examples

use train_station::{Tensor, optimizers::Adam};
use train_station::serialization::Serializable;

let weight = Tensor::ones(vec![2, 3]).with_requires_grad();
let mut optimizer = Adam::new();
optimizer.add_parameter(&weight);

let binary = optimizer.to_binary().unwrap();
let loaded_optimizer = Adam::from_binary(&binary).unwrap();
assert_eq!(loaded_optimizer.saved_parameter_count(), 1);

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 Adam

fn to_serializer(&self) -> StructSerializer

Convert Adam to StructSerializer for serialization

Serializes all optimizer state including configuration, parameter states, and global step count. Parameter linking is not serialized and must be done after deserialization.

Returns

StructSerializer containing all serializable optimizer state

fn from_deserializer( deserializer: &mut StructDeserializer, ) -> SerializationResult<Self>

Create Adam from StructDeserializer

Reconstructs Adam optimizer from serialized state. Parameters must be linked separately using add_parameter or add_parameters.

Arguments

• deserializer - StructDeserializer containing optimizer data

Returns

Reconstructed Adam instance without parameter links, 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