Crate axonml 

Axonml - A Complete ML/AI Framework in Pure Rust

Axonml is a comprehensive machine learning framework that provides ~92-95% PyTorch-equivalent functionality in pure Rust. It includes 1076+ passing tests and production-ready features:

Core Features

• Tensors: N-dimensional arrays with broadcasting, views, BLAS operations, sparse tensors
• 70+ Tensor Operations: arithmetic, reduction, sorting (topk, sort, argsort), indexing (gather, scatter, nonzero, unique)
• Autograd: Automatic differentiation with computational graph, AMP (autocast), gradient checkpointing
• Neural Networks: Linear, Conv1d/2d, BatchNorm, LayerNorm, GroupNorm, InstanceNorm, Attention, RNN/LSTM/GRU
• Optimizers: SGD, Adam, AdamW, RMSprop, LAMB with LR schedulers and GradScaler
• Data Loading: Dataset trait, DataLoader, samplers, transforms, parallel loading
• Vision: Image transforms, MNIST/CIFAR datasets, ResNet/VGG/ViT architectures, pretrained hub
• Text: Tokenizers (BPE, WordPiece), vocabularies, text datasets
• Audio: Spectrograms, MFCC, audio transforms, audio datasets
• Distributed: DDP, FSDP (ZeRO-2/3), Pipeline Parallelism, Tensor Parallelism
• LLM: BERT, GPT-2 architectures with pretrained model hub (LLaMA, Mistral, Phi, Qwen)
• GPU Backends: CUDA, Vulkan, Metal, WebGPU with comprehensive test suite

Model Hub & Benchmarking

• Unified model registry across vision and LLM domains
• Model search, filtering, and recommendations
• Throughput testing and memory profiling utilities

Quick Start

use axonml::prelude::*;

// Create a tensor
let x = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]).unwrap();

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

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

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

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

// Training loop
for batch in loader.iter() {
    // Forward pass
    let output = model.forward(&batch.data);

    // Compute loss
    let loss = cross_entropy_loss(&output, &batch.labels);

    // Backward pass
    loss.backward();

    // Update weights
    optimizer.step();
    optimizer.zero_grad();
}

Mixed Precision Training

use axonml_autograd::amp::autocast;
use axonml_optim::GradScaler;

let mut scaler = GradScaler::new();

// Forward with autocast
let loss = autocast(DType::F16, || {
    model.forward(&input)
});

// Scale loss for backward
let scaled_loss = scaler.scale_loss(loss);
scaled_loss.backward();

// Unscale and step
scaler.unscale_grads(&mut gradients);
optimizer.step();
scaler.update();

Feature Flags

• full (default): All features enabled
• core: Core tensor and autograd functionality
• nn: Neural network layers and optimizers
• data: Data loading utilities
• vision: Image processing and vision datasets
• text: Text processing and NLP utilities
• audio: Audio processing utilities
• distributed: Distributed training utilities (DDP, FSDP, Pipeline)
• profile: Performance profiling and bottleneck detection
• llm: LLM architectures (BERT, GPT-2) with pretrained hub
• jit: JIT compilation and tracing
• onnx: ONNX model import and export
• cuda: NVIDIA CUDA GPU backend
• vulkan: Vulkan GPU backend
• metal: Apple Metal GPU backend
• wgpu: WebGPU backend

@version 0.2.6 @author AutomataNexus Development Team

Re-exports

pub use adversarial::fgsm_attack;

pub use adversarial::pgd_attack;

pub use adversarial::adversarial_training_step;

pub use adversarial::AdversarialTrainer;

pub use trainer::Callback;

pub use trainer::EarlyStopping;

pub use trainer::ProgressLogger;

pub use trainer::TrainingConfig;

pub use trainer::TrainingHistory;

pub use trainer::TrainingMetrics;

pub use trainer::TrainingState;

pub use trainer::clip_grad_norm;

pub use trainer::compute_accuracy;

pub use hub::BenchmarkResult;

pub use hub::ModelCategory;

pub use hub::UnifiedModelInfo;

pub use hub::compare_benchmarks;

pub use hub::list_all_models;

pub use hub::models_by_category;

pub use hub::models_by_max_params;

pub use hub::models_by_max_size_mb;

pub use hub::recommended_models;

pub use hub::search_models;

pub use benchmark::print_throughput_results;

pub use benchmark::MemorySnapshot;

pub use benchmark::ThroughputConfig;

pub use benchmark::ThroughputResult;

pub use benchmark::benchmark_model;

pub use benchmark::benchmark_model_named;

pub use benchmark::compare_models;

pub use benchmark::throughput_test;

pub use benchmark::warmup_model;

pub use benchmark::print_memory_profile;

pub use benchmark::profile_model_memory;

pub use axonml_core as core;

pub use axonml_tensor as tensor;

pub use axonml_autograd as autograd;

pub use axonml_nn as nn;

pub use axonml_optim as optim;

pub use axonml_data as data;

pub use axonml_vision as vision;

pub use axonml_text as text;

pub use axonml_audio as audio;

pub use axonml_distributed as distributed;

pub use axonml_profile as profile;

pub use axonml_llm as llm;

pub use axonml_jit as jit;

pub use axonml_onnx as onnx;

Modules

adversarial

Adversarial training utilities (FGSM, PGD, adversarial trainer). Adversarial Training Utilities

benchmark

Model Benchmarking Utilities

hub

Unified Model Hub - Central Registry for All Pretrained Models

hvac

HVAC 8-model diagnostic system (~8.6M total parameters). HVAC 8-Model Diagnostic System

prelude

Common imports for machine learning tasks.

trainer

High-Level Training Utilities

Functions

features

Returns a string describing the enabled features.

version

Returns the version of the Axonml framework.