Crate axonml 

Axonml - A Complete ML/AI Framework in Pure Rust

Axonml is a comprehensive machine learning framework that provides PyTorch-equivalent functionality in pure Rust. It includes:

• Tensors: N-dimensional arrays with broadcasting, views, and BLAS operations
• Autograd: Automatic differentiation with computational graph
• Neural Networks: Layers, modules, loss functions, and activations
• Optimizers: SGD, Adam, AdamW, RMSprop with learning rate schedulers
• Data Loading: Dataset trait, DataLoader, samplers, and transforms
• Vision: Image transforms, MNIST/CIFAR datasets, CNN architectures
• Text: Tokenizers (BPE, WordPiece), vocabularies, text datasets
• Audio: Spectrograms, MFCC, audio transforms, audio datasets
• Distributed: DDP, all-reduce, broadcast, process groups

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();
}

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
• profile: Performance profiling and bottleneck detection
• llm: LLM architectures (BERT, GPT-2)
• jit: JIT compilation and tracing

@version 0.1.0 @author AutomataNexus Development Team

Re-exports

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;

Modules

prelude

Common imports for machine learning tasks.

Functions

features

Returns a string describing the enabled features.

version

Returns the version of the Axonml framework.