AxonML (named after axons - the nerve fibers that transmit signals between neurons) is an ambitious open-source project to create a complete machine learning framework in Rust. Our goal is to provide the same comprehensive functionality as PyTorch while leveraging Rust’s performance, safety, and concurrency guarantees.

PyTorch Parity: ~92-95%

AxonML provides comprehensive PyTorch-equivalent functionality with 1076+ passing tests.

Key Features

Architecture

Category	Features
Tensor Operations	N-dimensional tensors, broadcasting, views, slicing, matmul, reductions
Automatic Differentiation	Dynamic computational graph, reverse-mode autodiff, AMP, checkpointing
Neural Networks	Linear, Conv1d/2d, BatchNorm, LayerNorm, GroupNorm, Attention, LSTM/GRU
Optimizers	SGD, Adam, AdamW, RMSprop, LAMB with LR schedulers
Distributed Training	DDP, FSDP (ZeRO-2/3), Pipeline Parallelism
Model Formats	ONNX import/export (40+ operators), SafeTensors
GPU Backends	CUDA, Vulkan, Metal, WebGPU

┌─────────────────────────────────────────────────────────────────────┐
│                           AxonML Suite                               │
├─────────────────────────────────────────────────────────────────────┤
│  axonml-cli    │  axonml-server  │  axonml-tui  │  axonml-dashboard │
├────────────────┴─────────────────┴──────────────┴───────────────────┤
│                              axonml                                  │
│                        (Umbrella Crate)                              │
├──────────────┬──────────────┬──────────────┬────────────────────────┤
│ axonml-vision│ axonml-audio │ axonml-text  │ axonml-llm             │
├──────────────┴──────────────┴──────────────┴────────────────────────┤
│ axonml-nn   │ axonml-optim │ axonml-data  │ axonml-distributed     │
├─────────────┴──────────────┴──────────────┴─────────────────────────┤
│           axonml-autograd          │         axonml-serialize       │
├────────────────────────────────────┴────────────────────────────────┤
│                          axonml-tensor                               │
├─────────────────────────────────────────────────────────────────────┤
│                           axonml-core                                │
│              (Device, DType, Storage, Memory)                        │
└─────────────────────────────────────────────────────────────────────┘

Quick Links

Installation

Section	Description
Getting Started	Installation and first model
Tensor Operations	Working with tensors
Neural Networks	Building models
Training	Training loops and optimization
Distributed	Multi-GPU and distributed training
Crate Documentation	All 22 crates

[dependencies]
axonml = "0.2.8"

[dependencies]
axonml = { version = "0.2.8", features = ["cuda", "vision", "llm"] }

Quick Example

use axonml::prelude::*;

fn main() {
    // Create tensors
    let x = Tensor::randn(&[32, 784]);
    let y = Tensor::randn(&[32, 10]);

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

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

    // Training loop
    for epoch in 0..100 {
        let output = model.forward(&Variable::new(x.clone(), false));
        let loss = output.mse_loss(&Variable::new(y.clone(), false));

        optimizer.zero_grad();
        loss.backward();
        optimizer.step();

        println!("Epoch {}: Loss = {:.4}", epoch, loss.data().item());
    }
}

Operation	AxonML	PyTorch	Ratio
MatMul (1024x1024)	2.1ms	1.8ms	1.17x
Conv2d (224x224)	4.3ms	3.9ms	1.10x
LSTM (seq=128)	8.2ms	7.1ms	1.15x
Adam step	0.8ms	0.7ms	1.14x

AxonML Documentation

Overview