Skip to main content

Crate scirs2_autograd

Crate scirs2_autograd 

Source
Expand description

Ā§SciRS2 Autograd - Automatic Differentiation for Rust

scirs2-autograd provides PyTorch-style automatic differentiation with lazy tensor evaluation, enabling efficient gradient computation for scientific computing and deep learning.

Ā§šŸŽÆ Key Features

  • Reverse-mode Autodiff: Efficient backpropagation for neural networks
  • Lazy Evaluation: Build computation graphs, evaluate only when needed
  • Higher-order Gradients: Compute derivatives of derivatives
  • Neural Network Ops: Optimized operations for deep learning
  • Optimizers: Adam, SGD, RMSprop with state management
  • Model Persistence: Save and load trained models
  • Variable Management: Namespace-based variable organization

Ā§šŸ“¦ Installation

[dependencies]
scirs2-autograd = { version = "0.1.5", features = ["blas"] }

For fast matrix operations, enable BLAS (uses OxiBLAS - pure Rust):

[dependencies]
scirs2-autograd = { version = "0.1.5", features = ["blas"] }

Ā§šŸš€ Quick Start

Ā§Basic Differentiation

Compute gradients of a simple function:

use scirs2_autograd as ag;
use ag::tensor_ops as T;

ag::run(|ctx: &mut ag::Context<f64>| {
    // Define variables
    let x = ctx.placeholder("x", &[]);
    let y = ctx.placeholder("y", &[]);

    // Build computation graph: z = 2xĀ² + 3y + 1
    let z = 2.0 * x * x + 3.0 * y + 1.0;

    // Compute dz/dy
    let dz_dy = &T::grad(&[z], &[y])[0];
    println!("dz/dy = {:?}", dz_dy.eval(ctx));  // => 3.0

    // Compute dz/dx (feed x=2)
    let dz_dx = &T::grad(&[z], &[x])[0];
    let x_val = scirs2_core::ndarray::arr0(2.0);
    let result = ctx.evaluator()
        .push(dz_dx)
        .feed(x, x_val.view().into_dyn())
        .run()[0].clone();
    println!("dz/dx at x=2: {:?}", result);  // => 8.0

    // Higher-order: dĀ²z/dxĀ²
    let d2z_dx2 = &T::grad(&[dz_dx], &[x])[0];
    println!("dĀ²z/dxĀ² = {:?}", d2z_dx2.eval(ctx));  // => 4.0
});

Ā§Neural Network Training

Train a multi-layer perceptron for MNIST:

use scirs2_autograd as ag;
use ag::optimizers::adam::Adam;
use ag::tensor_ops::*;
use ag::prelude::*;

// Create variable environment
let mut env = ag::VariableEnvironment::new();
let mut rng = ag::ndarray_ext::ArrayRng::<f32>::default();

// Initialize network weights
env.name("w1").set(rng.glorot_uniform(&[784, 256]));
env.name("b1").set(ag::ndarray_ext::zeros(&[1, 256]));
env.name("w2").set(rng.glorot_uniform(&[256, 10]));
env.name("b2").set(ag::ndarray_ext::zeros(&[1, 10]));

// Create Adam optimizer
let var_ids = env.default_namespace().current_var_ids();
let adam = Adam::default("adam", var_ids, &mut env);

// Training loop
for epoch in 0..10 {
    env.run(|ctx| {
        // Define computation graph
        let x = ctx.placeholder("x", &[-1, 784]);
        let y = ctx.placeholder("y", &[-1]);

        let w1 = ctx.variable("w1");
        let b1 = ctx.variable("b1");
        let w2 = ctx.variable("w2");
        let b2 = ctx.variable("b2");

        // Forward pass: x -> hidden -> output
        let hidden = relu(matmul(x, w1) + b1);
        let logits = matmul(hidden, w2) + b2;

        // Loss: cross-entropy
        let loss = reduce_mean(
            sparse_softmax_cross_entropy(logits, &y),
            &[0],
            false
        );

        // Backpropagation
        let params = &[w1, b1, w2, b2];
        let grads = &grad(&[loss], params);

        // Update weights (requires actual data feeding)
        // let mut feeder = ag::Feeder::new();
        // feeder.push(x, x_batch).push(y, y_batch);
        // adam.update(params, grads, ctx, feeder);
    });
}

Ā§Custom Operations

Define custom differentiable operations:

use scirs2_autograd as ag;
use ag::tensor_ops::*;

ag::run::<f64, _, _>(|ctx| {
    let x = ones(&[3, 4], ctx);

    // Apply custom transformations using tensor.map()
    let y = x.map(|arr| arr.mapv(|v: f64| v * 2.0 + 1.0));

    // Hooks for debugging
    let z = x.showshape();  // Print shape
    let w = x.raw_hook(|arr| println!("Tensor value: {}", arr));
});

Ā§šŸ§  Core Concepts

Ā§Tensors

Lazy-evaluated multi-dimensional arrays with automatic gradient tracking:

use scirs2_autograd as ag;
use ag::tensor_ops::*;
use ag::prelude::*;

ag::run::<f64, _, _>(|ctx| {
    // Create tensors
    let a = zeros(&[2, 3], ctx);        // All zeros
    let b = ones(&[2, 3], ctx);         // All ones
    let c = ctx.placeholder("c", &[2, 3]);  // Placeholder (fill later)
    let d = ctx.variable("d");          // Trainable variable
});

Ā§Computation Graphs

Build graphs of operations, evaluate lazily:

use scirs2_autograd as ag;
use ag::tensor_ops as T;

ag::run::<f64, _, _>(|ctx| {
    let x = ctx.placeholder("x", &[2, 2]);
    let y = ctx.placeholder("y", &[2, 2]);

    // Build graph (no computation yet)
    let z = T::matmul(x, y);
    let w = T::sigmoid(z);

    // Evaluate when needed
    // let result = w.eval(ctx);
});

Ā§Gradient Computation

Reverse-mode automatic differentiation:

use scirs2_autograd as ag;
use ag::tensor_ops as T;

ag::run(|ctx| {
    let x = ctx.placeholder("x", &[]);
    let y = x * x * x;  // y = xĀ³

    // Compute dy/dx = 3xĀ²
    let dy_dx = &T::grad(&[y], &[x])[0];

    // Evaluate at x=2: 3(2Ā²) = 12
    let x_val = scirs2_core::ndarray::arr0(2.0);
    let grad_val = ctx.evaluator()
        .push(dy_dx)
        .feed(x, x_val.view().into_dyn())
        .run()[0].clone();
});

Ā§šŸŽØ Available Operations

Ā§Basic Math

  • Arithmetic: +, -, *, /, pow
  • Comparison: equal, not_equal, greater, less
  • Reduction: sum, mean, max, min

Ā§Neural Network Ops

  • Activations: relu, sigmoid, tanh, softmax, gelu
  • Pooling: max_pool2d, avg_pool2d
  • Convolution: conv2d, conv2d_transpose
  • Normalization: batch_norm, layer_norm
  • Dropout: dropout

Ā§Matrix Operations

  • matmul - Matrix multiplication
  • transpose - Matrix transpose
  • reshape - Change tensor shape
  • concat - Concatenate tensors
  • split - Split tensor

Ā§Loss Functions

  • sparse_softmax_cross_entropy - Classification loss
  • sigmoid_cross_entropy - Binary classification
  • softmax_cross_entropy - Multi-class loss

Ā§šŸ”§ Optimizers

Built-in optimization algorithms:

  • Adam: Adaptive moment estimation (recommended)
  • SGD: Stochastic gradient descent with momentum
  • RMSprop: Root mean square propagation
  • Adagrad: Adaptive learning rates

Ā§šŸ’¾ Model Persistence

Save and load trained models:

use scirs2_autograd as ag;

let mut env = ag::VariableEnvironment::<f64>::new();

// After training...
env.save("model.safetensors")?;

// Later, load the model
let env = ag::VariableEnvironment::<f64>::load("model.safetensors")?;

Ā§šŸ“Š Performance

scirs2-autograd is designed for efficiency:

  • Lazy Evaluation: Build graphs without computation overhead
  • Minimal Allocations: Reuse memory where possible
  • BLAS Integration: Fast matrix operations via OxiBLAS (pure Rust)
  • Zero-copy: Efficient data handling with ndarray views

Typical training speed: 0.11 sec/epoch for MNIST MLP (2.7GHz Intel Core i5)

Ā§šŸ”— Integration

  • scirs2-neural: High-level neural network layers
  • scirs2-linalg: Matrix operations
  • scirs2-optimize: Optimization algorithms
  • ndarray: Core array library (re-exported)

Ā§šŸ“š Comparison with PyTorch

FeaturePyTorchscirs2-autograd
Autodiffāœ…āœ…
Dynamic Graphsāœ…āœ…
GPU Supportāœ…āš ļø (limited)
Type SafetyāŒāœ…
Memory Safetyāš ļøāœ…
Pure RustāŒāœ…

Ā§šŸ”’ Version

Current version: 0.1.5 (Released January 15, 2026)

Re-exportsĀ§

pub use crate::ndarray_ext::array_gen;
pub use crate::ndarray_ext::NdArray;
pub use crate::ndarray_ext::NdArrayView;
pub use crate::ndarray_ext::NdArrayViewMut;
pub use crate::evaluation::Evaluator;
pub use crate::evaluation::Feeder;
pub use crate::tensor::Tensor;
pub use crate::error::AutogradError;
pub use crate::error::EvalError;
pub use crate::error::OpError;
pub use crate::error::Result;
pub use crate::graph::run;
pub use crate::graph::Context;
pub use crate::high_performance::memory_efficient_grad_accumulation;
pub use crate::high_performance::parallel_gradient_computation;
pub use crate::high_performance::simd_backward_pass;
pub use crate::high_performance::ultra_backward_pass;
pub use crate::variable::AutogradTensor;
pub use crate::variable::SafeVariable;
pub use crate::variable::SafeVariableEnvironment;
pub use crate::variable::VariableEnvironment;
pub use crate::optimizers::FunctionalAdam;
pub use crate::optimizers::FunctionalOptimizer;
pub use crate::optimizers::FunctionalSGD;

ModulesĀ§

error
Error types for the autograd module
error_helpers
Helper functions to eliminate repetitive expect() patterns throughout the codebase.
evaluation
gradient_clipping
Gradient clipping utilities
graph
high_performance
High-Performance Autograd APIs for ToRSh Integration
hooks
You can register hooks on ag::Tensor objects for debugging.
integration
Integration utilities for working with other SciRS2 modules
ndarray
Complete ndarray re-export for SciRS2 ecosystem
ndarray_ext
A small extension of ndarray
op
Implementing differentiable operations
optimization
Graph optimization and expression simplification for computation graphs
optimizers
A collection of gradient descent optimizers
parallel
Parallel processing and thread pool optimizations
prelude
Exports useful trait implementations
rand
Ultra-advanced random number generation for SCIRS2 ecosystem
schedulers
Learning rate schedulers
tensor
tensor_ops
A collection of functions for manipulating autograd::Tensor objects
test_helper
Provides helper functions for testing.
testing
Numerical stability testing framework for automatic differentiation
tracing
Enhanced tracing and recording capabilities for computation graphs
validation
Domain validation utilities for mathematical operations.
variable
Variable and namespace
visualization
Graph visualization and debugging tools for computation graphs

TraitsĀ§

Float
A primitive type in this crate, which is actually a decorated scirs2_core::numeric::Float.

FunctionsĀ§

rand
Convenience function to generate a random value of the inferred type