axonml 0.2.4

A complete ML/AI framework in pure Rust - PyTorch-equivalent functionality
Documentation

Overview

axonml is the umbrella crate for the AxonML machine learning framework - a complete ML/AI toolkit written in pure Rust. It provides PyTorch-equivalent functionality including tensors, automatic differentiation, neural networks, optimizers, data loading, and domain-specific utilities for vision, text, and audio processing.

This crate re-exports all AxonML sub-crates under a unified API, allowing you to build and train deep learning models with a single import.

Features

  • Tensors - N-dimensional arrays with broadcasting, views, slicing, and BLAS-accelerated operations
  • Automatic Differentiation - Computational graph construction with backward pass for gradient computation
  • Neural Networks - Linear, Conv2d, BatchNorm, LayerNorm, Attention, RNN/LSTM/GRU, and more
  • Optimizers - SGD, Adam, AdamW, RMSprop with learning rate schedulers (Step, Exponential, Cosine)
  • Data Loading - Dataset trait, DataLoader with batching, samplers, and transforms
  • Vision - Image transforms, MNIST/CIFAR datasets, CNN architectures (LeNet, SimpleCNN)
  • Text - Tokenizers (BPE, WordPiece, Whitespace), vocabularies, text datasets
  • Audio - Spectrograms, MFCC, audio transforms, synthetic audio datasets
  • Distributed - DDP, all-reduce, broadcast, barrier, process groups
  • Profiling - Memory and compute profilers, timeline analysis, bottleneck detection
  • LLM Architectures - BERT and GPT-2 implementations with generation utilities
  • JIT Compilation - Graph tracing and optimization for compiled inference

Modules (Re-exports)

Module Description Feature Flag
core Error types, Device, DType definitions core
tensor N-dimensional tensor operations core
autograd Automatic differentiation and Variables core
nn Neural network layers and modules nn
optim Optimizers and learning rate schedulers nn
data DataLoader, Dataset trait, samplers, transforms data
vision Image transforms and vision datasets vision
text Tokenizers, vocabularies, text datasets text
audio Audio transforms and audio datasets audio
distributed Distributed training utilities (DDP) distributed
profile Profiling and bottleneck analysis profile
llm BERT and GPT-2 architectures llm
jit JIT compilation and graph optimization jit

Usage

Add axonml to your Cargo.toml:

[dependencies]
axonml = "0.1.0"

Using the Prelude

The prelude module re-exports commonly used types for quick access:

use axonml::prelude::*;

fn main() -> axonml::core::Result<()> {
    // Create a tensor
    let x = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[2, 2])?;

    // Wrap in a Variable for autograd
    let var = Variable::new(x, true);

    // Build a simple model
    let model = Sequential::new()
        .add(Linear::new(784, 128))
        .add(ReLU)
        .add(Linear::new(128, 10));

    // Create an optimizer
    let mut optimizer = Adam::new(model.parameters(), 0.001);

    Ok(())
}

Training Loop Example

use axonml::prelude::*;

fn train() -> axonml::core::Result<()> {
    // Create model
    let model = Sequential::new()
        .add(Linear::new(784, 256))
        .add(ReLU)
        .add(Dropout::new(0.2))
        .add(Linear::new(256, 10));

    // Optimizer
    let mut optimizer = Adam::new(model.parameters(), 0.001);

    // Dataset and loader
    let dataset = SyntheticMNIST::new(1000);
    let loader = DataLoader::new(dataset, 32);

    // Training loop
    for epoch in 0..10 {
        for batch in loader.iter() {
            // Forward pass
            let output = model.forward(&batch.data);
            let loss = CrossEntropyLoss::new().forward(&output, &batch.labels);

            // Backward pass
            loss.backward();

            // Update weights
            optimizer.step();
            optimizer.zero_grad();
        }
        println!("Epoch {} complete", epoch + 1);
    }

    Ok(())
}

Feature Flags

Select only the features you need:

[dependencies]
# Full framework (default)
axonml = "0.1.0"

# Core only (tensors + autograd)
axonml = { version = "0.1.0", default-features = false, features = ["core"] }

# Neural networks without domain-specific modules
axonml = { version = "0.1.0", default-features = false, features = ["nn", "data"] }

# Vision pipeline
axonml = { version = "0.1.0", default-features = false, features = ["vision"] }

# NLP pipeline
axonml = { version = "0.1.0", default-features = false, features = ["text", "llm"] }

Feature Flag Reference

Feature Includes Description
full All features Complete framework (default)
core tensor, autograd Core tensor operations and autodiff
nn core + nn, optim Neural network layers and optimizers
data core + data DataLoader and dataset utilities
vision nn, data + vision Image processing and vision datasets
text nn, data + text Tokenizers and text processing
audio nn, data + audio Audio transforms and datasets
distributed nn + distributed Distributed training (DDP)
profile core + profile Profiling and analysis tools
llm nn + llm BERT and GPT-2 architectures
jit core + jit JIT compilation and tracing

Version Information

use axonml::{version, features};

fn main() {
    println!("AxonML version: {}", version());
    println!("Enabled features: {}", features());
}

Tests

Run the full test suite:

cargo test -p axonml

Run tests for specific features:

cargo test -p axonml --features "core nn data"

Run integration tests:

cargo test -p axonml --test integration_test

Examples

The crate includes example programs:

# Simple training example
cargo run -p axonml --example simple_training

# MNIST digit classification
cargo run -p axonml --example mnist_training

# NLP and audio feature test
cargo run -p axonml --example nlp_audio_test

Sub-Crate Structure

axonml/
  axonml-core       - Error types, Device, DType
  axonml-tensor     - N-dimensional tensor operations
  axonml-autograd   - Automatic differentiation
  axonml-nn         - Neural network layers
  axonml-optim      - Optimizers and schedulers
  axonml-data       - DataLoader and Dataset trait
  axonml-vision     - Image transforms and datasets
  axonml-text       - Tokenizers and text utilities
  axonml-audio      - Audio transforms and datasets
  axonml-distributed- Distributed training (DDP)
  axonml-profile    - Profiling and analysis
  axonml-llm        - BERT and GPT-2 models
  axonml-jit        - JIT compilation
  axonml-serialize  - Model serialization
  axonml-tui        - Terminal user interface
  axonml (this)     - Umbrella crate

License

Licensed under either of:

at your option.