axonml 0.6.2

Overview

axonml is the umbrella crate for the AxonML machine-learning framework — a PyTorch-equivalent ML/AI toolkit written in pure Rust. It re-exports every sub-crate under a unified namespace, so you can pull in the whole framework with a single dependency.

The crate is intentionally thin. After the 0.6.1 split, it contains only:

1. Feature-gated re-exports of every sub-crate (axonml-core .. axonml-train)
2. A prelude module with the most-used types
3. The live browser training monitor (TrainingMonitor)
4. version() / features() introspection helpers

Everything else — layers, optimizers, models, data loaders, training infrastructure, HVAC diagnostics, adversarial training — lives in dedicated sibling crates that can also be used standalone.

Last updated: 2026-04-16 — version 0.6.1.

Sub-Crate Architecture

The umbrella crate re-exports the framework sub-crates under short module names:

Core

Neural Networks

Data & I/O

Domain-Specific

Training, Optimization, Deployment

Tooling (workspace crates, not re-exported)

Nine LLM Architectures (feature = "llm")

Plus:

• BPE / WordPiece / Whitespace / Character tokenizers (in axonml-text)
• Generation utilities — TextGenerator, GenerationConfig
• Training scripts for every model in crates/axonml-llm/examples/

The axonml::prelude re-exports only a small subset of LLM types (Bert, BertConfig, BertForMaskedLM, BertForSequenceClassification, GPT2, GPT2Config, GPT2LMHead, GenerationConfig, TextGenerator). For the other architectures pull from axonml::llm::* directly.

Installation

[dependencies]
axonml = "0.6.1"                     # default = full feature set

Select only what you need:

[dependencies]
# Core tensors + autograd
axonml = { version = "0.6.1", default-features = false, features = ["core"] }

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

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

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

# With GPU acceleration
axonml = { version = "0.6.1", features = ["full", "cuda"] }

Usage

Prelude

The prelude module exports the most-used types so you can get started fast:

use axonml::prelude::*;

fn main() -> axonml::core::Result<()> {
    // Tensor + autograd
    let x = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[2, 2])?;
    let var = Variable::new(x, true);

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

    // Optimizer
    let mut optimizer = AdamW::new(model.parameters(), 1e-3);
    optimizer.zero_grad();

    Ok(())
}

Contents of the prelude, feature-gated:

• core: DType, Device, Error, Result, Tensor, Variable, no_grad
• nn: layers (Linear, Conv2d, LSTM, GRU, RNN, MultiHeadAttention, BatchNorm1d/2d, LayerNorm, Dropout, Embedding, MaxPool2d, AvgPool2d, activations, losses, Parameter, Sequential, Module) and optimizers (SGD, Adam, AdamW, RMSprop, Optimizer, LRScheduler, CosineAnnealingLR, ExponentialLR, StepLR)
• data: DataLoader, Dataset, RandomSampler, SequentialSampler, Transform
• vision: LeNet, SimpleCNN, SyntheticMNIST, SyntheticCIFAR, CenterCrop, ImageNormalize, RandomHorizontalFlip, Resize
• text: tokenizers, Vocab, TextDataset, LanguageModelDataset, SyntheticSentimentDataset
• audio: MelSpectrogram, MFCC, Resample, NormalizeAudio, AddNoise, synthetic datasets
• distributed: DDP, DistributedDataParallel, ProcessGroup, World, all_reduce_{mean,sum}, barrier, broadcast
• profile: Profiler, ComputeProfiler, MemoryProfiler, TimelineProfiler, Bottleneck, BottleneckAnalyzer, ProfileGuard, ProfileReport
• llm: GPT2, GPT2Config, GPT2LMHead, Bert, BertConfig, BertForMaskedLM, BertForSequenceClassification, GenerationConfig, TextGenerator
• train: TrainingConfig, TrainingHistory, TrainingMetrics, EarlyStopping, ProgressLogger, Callback, AdversarialTrainer
• jit: CompiledFunction, Graph, JitCompiler, Optimizer as JitOptimizer, TracedValue, trace

Training Loop with Live Monitor

use axonml::prelude::*;
use axonml::TrainingMonitor;

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

    let mut optimizer = AdamW::new(model.parameters(), 1e-3);
    let param_count = model.parameters().iter().map(|p| p.numel()).sum::<usize>();

    // Live browser dashboard — opens Chromium automatically
    let monitor = TrainingMonitor::new("MNIST Classifier", param_count)
        .total_epochs(10)
        .batch_size(32)
        .launch();

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

    for epoch in 0..10 {
        let mut epoch_loss = 0.0f32;
        let mut batches = 0;
        for batch in loader.iter() {
            let output = model.forward(&batch.data);
            let loss = CrossEntropyLoss::new().compute(&output, &batch.targets);
            optimizer.zero_grad();
            loss.backward();
            optimizer.step();
            epoch_loss += loss.data().to_vec()[0];
            batches += 1;
        }
        monitor.log_epoch(epoch + 1, epoch_loss / batches as f32, None, vec![]);
    }

    monitor.set_status("complete");
    Ok(())
}

Training Monitor

axonml::monitor::TrainingMonitor (re-exported as axonml::TrainingMonitor) is a zero-dependency, pure-Rust HTTP server that serves a real-time training dashboard to your browser. The dashboard HTML lives next to the module at crates/axonml/src/monitor_dashboard.html. It is intentionally kept in the umbrella crate so every training script across the workspace can use it with a single import.

use axonml::TrainingMonitor;

let monitor = TrainingMonitor::new("MyModel", param_count)
    .total_epochs(50)
    .batch_size(32)
    .launch();  // opens http://127.0.0.1:<auto-port> in Chromium

// Each epoch
monitor.log_epoch(epoch + 1, train_loss, Some(val_loss), vec![
    ("accuracy", acc),
    ("lr", lr),
]);

// When done
monitor.set_status("complete");

The dashboard shows real-time training loss, optional validation loss, custom metric overlays, current epoch / total epochs, and best loss so far.

Feature Flag Reference

Examples

The crate ships three generic examples; HVAC-specific examples live in axonml-hvac.

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

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

# NLP + audio transform smoke test
cargo run -p axonml --example nlp_audio_test

For model-specific training scripts, see the per-crate examples/ directories:

# LLM training (all 9 architectures)
cargo run -p axonml-llm --example train_gpt2 --release
cargo run -p axonml-llm --example train_trident --release
cargo run -p axonml-llm --example train_hydra --release

# Vision training
cargo run -p axonml-vision --example train_resnet --release
cargo run -p axonml-vision --example train_mnemosyne --release  # Aegis biometric

# HVAC training
cargo run -p axonml-hvac --example train_panoptes --release

Introspection

use axonml::{version, features};

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

Version Information

• Crate version: 0.6.1
• Rust edition: 2024
• MSRV: Rust 1.85+
• 0.6.1 split: hvac (HVAC diagnostic models) and train (trainer / hub / benchmark / adversarial) were extracted from this umbrella into standalone crates to keep the umbrella a thin re-export layer. The live browser TrainingMonitor stayed here.

License

Licensed under either of:

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

at your option.

Part of AxonML — a complete ML/AI framework in pure Rust, by AutomataNexus LLC.