OptiRS

Advanced ML optimization and hardware acceleration library - Main integration crate for the OptiRS ecosystem.

Overview

OptiRS is a comprehensive Rust library for machine learning optimization that provides state-of-the-art optimization algorithms, hardware acceleration, learned optimizers, neural architecture search, and performance analysis tools. This main crate serves as the unified entry point to the entire OptiRS ecosystem, allowing users to access all functionality through feature gates.

Features

• Core Optimization: Traditional and advanced optimization algorithms (SGD, Adam, AdamW, RMSprop)
• GPU Acceleration: Multi-backend GPU support (CUDA, Metal, OpenCL, WebGPU)
• TPU Coordination: Large-scale distributed optimization on Google Cloud TPUs
• Learned Optimizers: Neural network-based optimization with meta-learning
• Neural Architecture Search: Automated architecture and hyperparameter optimization
• Performance Analysis: Comprehensive benchmarking and profiling tools
• SciRS2 Integration: Built on the SciRS2 scientific computing foundation
• Cross-Platform: Support for Linux, macOS, Windows, and WebAssembly

Quick Start

Add OptiRS to your Cargo.toml:

[dependencies]
optirs = "0.1.0-beta.1"

Basic Example

use optirs::prelude::*;
use optirs::optimizers::Adam;
use ndarray::Array2;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create an Adam optimizer
    let mut optimizer = Adam::new(0.001)
        .beta1(0.9)
        .beta2(0.999)
        .build();

    // Initialize parameters
    let mut params = Array2::<f32>::zeros((100, 50));
    let gradients = Array2::<f32>::ones((100, 50)) * 0.01;

    // Perform optimization step
    optimizer.step(&mut params.view_mut(), &gradients.view())?;

    println!("Optimization step completed!");
    Ok(())
}

Feature Gates

OptiRS uses feature gates to allow selective compilation of functionality:

Core Features

[dependencies]
optirs = { version = "0.1.0-beta.1", features = ["core"] }  # Always included

Hardware Acceleration

[dependencies]
optirs = { version = "0.1.0-beta.1", features = ["gpu"] }  # GPU acceleration

[dependencies]
optirs = { version = "0.1.0-beta.1", features = ["tpu"] }  # TPU coordination

Advanced Optimization

[dependencies]
optirs = { version = "0.1.0-beta.1", features = ["learned"] }  # Learned optimizers

[dependencies]
optirs = { version = "0.1.0-beta.1", features = ["nas"] }  # Neural Architecture Search

Development and Analysis

[dependencies]
optirs = { version = "0.1.0-beta.1", features = ["bench"] }  # Benchmarking tools

Full Feature Set

[dependencies]
optirs = { version = "0.1.0-beta.1", features = ["full"] }  # All features

Architecture Overview

OptiRS Ecosystem
├── optirs-core     │ Core optimization algorithms
├── optirs-gpu      │ GPU acceleration (CUDA, Metal, OpenCL, WebGPU)
├── optirs-tpu      │ TPU coordination and distributed training
├── optirs-learned  │ Learned optimizers and meta-learning
├── optirs-nas      │ Neural Architecture Search
├── optirs-bench    │ Benchmarking and performance analysis
└── optirs          │ Main integration crate (this crate)

Usage Examples

GPU-Accelerated Optimization

use optirs::prelude::*;

#[cfg(feature = "gpu")]
use optirs::gpu::{GpuOptimizer, DeviceManager};

#[tokio::main]
#[cfg(feature = "gpu")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize GPU
    let device_manager = DeviceManager::new().await?;
    let device = device_manager.select_best_device()?;

    // Create GPU optimizer
    let mut gpu_optimizer = GpuOptimizer::new(device)
        .with_optimizer(Adam::new(0.001))
        .build()?;

    // Create GPU tensors
    let mut params = gpu_optimizer.create_tensor(&[1024, 512])?;
    let grads = gpu_optimizer.create_tensor_from_slice(&gradient_data)?;

    // GPU-accelerated optimization step
    gpu_optimizer.step(&mut params, &grads).await?;

    Ok(())
}

Learned Optimizer

use optirs::prelude::*;

#[cfg(feature = "learned")]
use optirs::learned::{TransformerOptimizer, MetaLearner};

#[tokio::main]
#[cfg(feature = "learned")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create learned optimizer
    let mut learned_optimizer = TransformerOptimizer::new()
        .with_hidden_size(256)
        .with_num_layers(4)
        .with_num_heads(8)
        .build()?;

    // Meta-train the optimizer
    let meta_tasks = load_training_tasks()?;
    learned_optimizer.meta_train(&meta_tasks).await?;

    // Use for optimization
    let mut params = create_model_parameters()?;
    let grads = compute_gradients(&params)?;

    learned_optimizer.step(&mut params, &grads).await?;

    Ok(())
}

Neural Architecture Search

use optirs::prelude::*;

#[cfg(feature = "nas")]
use optirs::nas::{BayesianOptimizer, SearchSpace};

#[tokio::main]
#[cfg(feature = "nas")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Define search space
    let search_space = SearchSpace::new()
        .add_continuous("learning_rate", 1e-5, 1e-1)
        .add_discrete("batch_size", &[16, 32, 64, 128])
        .add_categorical("optimizer", &["sgd", "adam", "adamw"])
        .build();

    // Create Bayesian optimizer
    let mut nas = BayesianOptimizer::new()
        .with_search_space(search_space)
        .with_budget(100)
        .build()?;

    // Search for optimal configuration
    let best_config = nas.optimize().await?;

    println!("Best configuration: {:?}", best_config);
    Ok(())
}

Performance Benchmarking

use optirs::prelude::*;

#[cfg(feature = "bench")]
use optirs::bench::{BenchmarkSuite, BenchmarkConfig};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Setup benchmark
    let config = BenchmarkConfig::new()
        .with_iterations(1000)
        .with_warmup_iterations(100)
        .build();

    let mut benchmark = BenchmarkSuite::new()
        .with_config(config)
        .add_optimizer("Adam", Adam::new(0.001))
        .add_optimizer("SGD", SGD::new(0.01))
        .build()?;

    // Run benchmarks
    let results = benchmark.run()?;
    results.print_summary();

    Ok(())
}

Multi-GPU Distributed Training

use optirs::prelude::*;

#[cfg(all(feature = "gpu", feature = "tpu"))]
use optirs::{gpu::MultiGpuOptimizer, tpu::DistributedCoordinator};

#[tokio::main]
#[cfg(all(feature = "gpu", feature = "tpu"))]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Setup distributed training
    let coordinator = DistributedCoordinator::new()
        .with_gpu_nodes(4)
        .with_tpu_pods(2)
        .build().await?;

    let mut distributed_optimizer = coordinator
        .create_distributed_optimizer(Adam::new(0.001))
        .await?;

    // Distributed optimization step
    distributed_optimizer.step_synchronized(&gradients).await?;

    Ok(())
}

Integration with SciRS2

OptiRS is built on the SciRS2 scientific computing ecosystem:

use optirs::prelude::*;
use scirs2_core::Array;
use scirs2_autograd::Variable;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create SciRS2 variables
    let mut params = Variable::new(Array::zeros([100, 50]));

    // OptiRS automatically integrates with SciRS2
    let mut optimizer = Adam::new(0.001)
        .with_scirs2_integration(true)
        .build();

    // Automatic differentiation with SciRS2
    let loss = compute_loss(&params);
    let grads = loss.backward();

    // Optimization with SciRS2 variables
    optimizer.step_scirs2(&mut params)?;

    Ok(())
}

Performance Characteristics

Benchmarks

Scalability

• Single GPU: Up to 10,000 parameters/ms
• Multi-GPU: Linear scaling up to 8 GPUs
• TPU Pods: Scaling to 1000+ cores
• Memory Efficiency: <1MB overhead per optimizer

Platform Support

Documentation

• API Documentation
• User Guide
• Examples
• Benchmarks

Contributing

We welcome contributions! Please see our Contributing Guide for details.

Development Setup

# Clone the repository
git clone https://github.com/cool-japan/optirs.git
cd optirs

# Install dependencies
cargo build

# Run tests
cargo test --all-features

# Run benchmarks
cargo bench

License

This project is dual-licensed under either of:

• Apache License, Version 2.0 (LICENSE-APACHE)
• MIT License (LICENSE-MIT)

at your option.

Acknowledgments

• Built on the SciRS2 scientific computing ecosystem
• Inspired by PyTorch, TensorFlow, and JAX optimization libraries
• Thanks to all contributors and the Rust ML community

OptiRS - Optimizing the future of machine learning in Rust 🚀