axonml

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)

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

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:

• Apache License, Version 2.0 (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
• MIT License (LICENSE-MIT or http://opensource.org/licenses/MIT)

at your option.