ML in Rust Basics

10 examples to get you started with machine learning in Rust - 7 basic and 3 intermediate.

Prerequisites

cargo new ml-playground && cd ml-playground
cargo add candle-core candle-nn --features mkl
cargo add tokenizers
cargo add ort --features download-binaries
cargo add axum tokio serde --features derive

Tooling: Examples target Rust 1.97.0 (edition 2024). Pin ML crate versions in production - APIs evolve quickly.

Basic Examples

1. Load an ONNX model with ort

Run exported models without a Python runtime at inference time.

use ort::session::Session;
 
fn main() -> ort::Result<()> {
    let session = Session::builder()?.commit_from_file("model.onnx")?;
    println!("inputs: {:?}", session.inputs);
    Ok(())
}

• ONNX is the common export path from PyTorch, sklearn, and TensorFlow.
• ort bundles ONNX Runtime for CPU/GPU execution.
• Inspect input names and shapes before building tensors.

Related: ONNX Runtime (ort) - full inference loop

2. Tokenize text for an LLM

Convert prompts to token IDs with the Hugging Face tokenizers crate.

use tokenizers::Tokenizer;
 
fn main() -> tokenizers::Result<()> {
    let tok = Tokenizer::from_file("tokenizer.json")?;
    let enc = tok.encode("Hello, Rust ML!", true)?;
    println!("ids: {:?}", enc.get_ids());
    Ok(())
}

• tokenizer.json ships with most HF model repos.
• add_special_tokens controls BOS/EOS insertion for decoder models.
• Tokenization is often 10-30% of LLM latency - keep it in Rust beside inference.

Related: Tokenizers - batching and truncation

3. Simple tensor with candle

Create and multiply tensors on CPU (or GPU with features enabled).

use candle_core::{Device, Tensor};
 
fn main() -> candle_core::Result<()> {
    let device = Device::Cpu;
    let a = Tensor::new(&[[1f32, 2.0], [3.0, 4.0]], &device)?;
    let b = Tensor::new(&[[2f32, 0.0], [1.0, 2.0]], &device)?;
    let c = a.matmul(&b)?;
    println!("{}", c);
    Ok(())
}

• Device::Cpu is the portable default; CUDA/Metal need feature flags and drivers.
• candle mirrors PyTorch-like tensor ops for small models and research.
• Shape errors surface as Result - validate dimensions early.

Related: candle - nn modules and training

4. Run a single ONNX inference step

Feed f32 input and read logits.

use ndarray::Array;
use ort::{inputs, session::Session, value::Tensor};
 
fn main() -> ort::Result<()> {
    let session = Session::builder()?.commit_from_file("mnist.onnx")?;
    let input = Array::from_shape_vec((1, 1, 28, 28), vec![0.0f32; 784])?;
    let outputs = session.run(inputs![Tensor::from_array(input)?])?;
    let logits = outputs[0].try_extract_tensor::<f32>()?;
    println!("logits len: {}", logits.len());
    Ok(())
}

• Shape must match the ONNX graph - use nhwd conventions from the export script.
• Normalize pixels the same way as training or accuracy collapses.
• Batch dimension 1 still matters for graph compatibility.

5. Serve a health check from Axum

Wrap inference behind HTTP with a minimal router.

use axum::{routing::get, Json, Router};
use serde_json::json;
 
#[tokio::main]
async fn main() {
    let app = Router::new().route("/health", get(|| async { Json(json!({"ok": true})) }));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

• Health endpoints let orchestrators probe before sending traffic.
• Keep model weights loaded once at startup, not per request.
• Add /predict with serde-validated JSON bodies next.

Related: Serving Models with Axum - production handlers

6. Serialize inference requests with serde

Typed API contracts prevent shape mistakes at the edge.

use serde::{Deserialize, Serialize};
 
#[derive(Deserialize)]
struct PredictRequest {
    text: String,
}
 
#[derive(Serialize)]
struct PredictResponse {
    label: String,
    score: f32,
}

• Validate max string length before tokenization to bound memory.
• Return scores with labels for observability in clients.
• Mirror training class order in label strings or IDs.

7. Choose burn for backend-agnostic training

Initialize a backend explicitly when experimenting with burn.

use burn::backend::NdArray;
use burn::tensor::Tensor;
 
type B = NdArray<f32>;
 
fn main() {
    let device = <B as burn::backend::Backend>::Device::default();
    let t: Tensor<B, 2> = Tensor::from_data([[1.0, 2.0], [3.0, 4.0]], &device);
    let _ = t.matmul(t.transpose());
}

• burn swaps NdArray, WGPU, and CUDA backends behind one API.
• Training loops live in Rust - no Python notebook required.
• Export to ONNX when you need cross-language deployment.

Related: burn - training loops and checkpoints

Intermediate Examples

8. Batch embeddings for retrieval

Compute vectors for multiple strings in one forward pass.

// Conceptual layout - pair tokenizers + ort or candle model:
// 1. tokenize with padding/truncation
// 2. build input tensors [batch, seq]
// 3. run session, take pooled output [batch, dim]

• Padding aligns sequence lengths within a batch on GPU.
• L2-normalize embeddings before cosine similarity search.
• Store vectors in a Rust-side index (see embeddings page).

Related: Embeddings & Vector Search

9. Local LLM inference pattern

Load weights once, stream tokens with a Rust LLM crate (llama.cpp bindings or mistral.rs).

// Pseudocode structure for local LLM services:
// let model = LlamaModel::load("weights.gguf", &params)?;
// let mut ctx = model.new_context()?;
// for token in ctx.generate(prompt, max_tokens) { print!("{token}") }

• GGUF weights are common for CPU/GPU local inference.
• KV cache size drives memory - cap max_seq_len per deployment.
• Run generation on a dedicated thread pool to avoid blocking Axum workers.

Related: LLM Inference & Serving

10. GPU feature flags

Enable CUDA or Metal only in release builds that target GPU hosts.

[dependencies]
candle-core = { version = "0.8", features = ["cuda"] }

• CI CPU-only agents should not require GPU features.
• Use separate Cargo features gpu vs default CPU builds.
• Fall back to CPU when Device::cuda_if_available() fails.

Related: GPU & Acceleration

Stack versions: This page was written for Rust 1.97.0 (edition 2024), Tokio 1.x, Axum 0.8, serde 1.0, sqlx 0.8, clap 4, and Polars 0.46+.

49454e4f4d5f434e4f5904434556494d43451b181f56181a181c1a1d