ExBurn

Status: Early development. Not yet ready for production use.

ExBurn is a middle layer between Nx and Burn that enables GPU-accelerated ML/DL on mobile and desktop devices.

Architecture

Axon model
↓
Nx.Defn graph
↓
ExBurn.Defn.Compiler (Nx.Defn.Compiler behaviour)
↓
ExBurn.Backend (Nx.Backend behaviour)
↓
ExBurn.Nif (Rustler NIF) ←→ ExCubecl (GPU buffers, kernels, pipelines)
↓
Burn Autodiff<CubeCL> (Rust)
↓
CubeCL kernels
↓
Metal (iOS) / Vulkan (Android) / CUDA → GPU

Status

Version 0.6.0 — Early Alpha

⚠️Note: This library is in early development. The API may change between minor versions. Not yet recommended for production use.

Feature	Status
Nx.Backend behaviour (basic ops)	✅ Implemented
Nx.Backend behaviour (shape ops)	✅ Implemented
Nx.Backend behaviour (reductions)	✅ Implemented
Nx.Backend behaviour (linear algebra)	✅ Implemented
Nx.Defn.Compiler	✅ Implemented
Rust NIF bridge (Burn CubeCL)	✅ Implemented
GPU acceleration (Metal/Vulkan)	✅ Via Burn/CubeCL
Axon model compilation	✅ Implemented
Training loop (SGD/Adam/RMSprop)	✅ Implemented
GPU forward pass (defn compiler)	✅ Implemented
Glorot/Xavier initialization	✅ Implemented
Layer freeze/unfreeze	✅ Implemented
Gradient accumulation	✅ Implemented
Nesterov momentum	✅ Implemented
Weight decay (L2)	✅ Implemented
Model summary	✅ Implemented
Device management (CPU↔GPU)	✅ Implemented
Nx.Serving	✅ Implemented
CUDA backend	✅ Implemented
Precompiled NIF binaries	🚧 Planned
Autodiff gradients	🚧 Planned

Known limitations

• Dtype support: The NIF stores tensors as f32 only. Other dtypes are value-converted to f32 at the backend boundary (the direct tensor API raises instead). Full dtype preservation is planned for a future release.
• Gradient computation: The default gradient method is :numerical (finite differences). Autodiff gradients are planned but not yet implemented.
• Global backend mutation: By default, Training.fit/3 no longer changes the global Nx backend. Pass set_default_backend?: true to restore the previous behavior.
• Evaluation loss: Loss is now correctly weighted by sample count when batches have different sizes.

Direct NIF operation constraints

When calling ExBurn.BurnBridge / ExBurn.Nifdirectly (bypassing the Nx backend, which normalizes these cases), the Rust layer currently imposes:

Operation	Constraint
matmul_tensor/2	Both operands rank ≥ 2
pow_tensor/2	Exponent must be an f32 scalar, not a tensor
transpose_tensor/1	2-D tensors only (use the Nx backend for other ranks)
broadcast_tensor/2	Expands existing dimensions only (e.g. [1,2] → [4,2])
iota_tensor/3	Produces a 1-D iota of length shape[axis]
reshape_tensor/2	Element count of the new shape must match

The Nx.Backend layer (ExBurn.Backend) lifts several of these by computing general cases exactly via Nx — prefer going through Nx unless you need the raw single-NIF-call path.

Features

• Nx Backend: Nx.Backend behaviour implementation — partial drop-in replacement for Nx.BinaryBackend with some dtype limitations
• Nx Defn Compiler: Custom Nx.Defn.Compiler that executes defn expressions on the Burn GPU backend
• GPU Acceleration: Burn's CubeCL backend with CUDA (NVIDIA), Metal (Apple), Vulkan (Android)
• ExCubecl Integration: GPU buffer management, kernel execution, async commands, and pipeline orchestration via ExCubecl
• Autodiff: Planned automatic differentiation via Burn's Autodiff backend decorator; currently uses numerical gradients
• Training Loop: Complete training with Adam, SGD, RMSprop optimizers, LR scheduling, gradient clipping, callbacks
• Model Management: Save/load, serialize, quantize (f16), benchmark
• Structured Errors: ExBurn.Error exception type with operation context

Quick Start

1. Install

Add ex_burn to your mix.exs:

defdepsdo[{:ex_burn,"~> 0.5"},{:nx,">= 0.12.0 and < 2.0.0"},{:axon,"~> 0.8"}]end

mix deps.get
mix compile # first build compiles the Rust NIF — expect a few minutes

The NIF builds CPU-only by default. For GPU acceleration see GPU Backends below.

2. Verify the installation

# In iex -S mix:ExBurn.smoke_test()#=> :ok ← Nx → Backend → NIF → Burn pipeline worksExBurn.summary()#=># ExBurn v0.5.0# ──────────────────────────────# Device: Metal (Apple M2 Pro) ← or "NdArray (CPU)" without GPU features# GPU: available# Backends: metal

3. Tensor operations through Burn

# Set ExBurn as the default Nx backendNx.default_backend(ExBurn.Backend)# …or equivalently: ExBurn.configure!()t=Nx.tensor([1.0,2.0,3.0])Nx.add(t,t)|>Nx.to_list()#=> [2.0, 4.0, 6.0]

4. GPU-accelerated functions with defn

# Set ExBurn as both backend and compilerNx.default_backend(ExBurn.Backend)Nx.Defn.global_default_options(compiler: ExBurn.Defn.Compiler)defmoduleMyMathdoimportNx.Defndefnadd_and_scale(x,y,scale)dox|>Nx.add(y)|>Nx.multiply(scale)endendresult=MyMath.add_and_scale(Nx.tensor([1.0,2.0]),Nx.tensor([3.0,4.0]),Nx.tensor(2.0))Nx.to_list(result)#=> [8.0, 12.0]

To use the compiler for a single call only (no global setting):

Nx.Defn.jit_apply(&MyMath.add_and_scale/3,[x,y,scale],compiler: ExBurn.Defn.Compiler)

5. Train a model with Axon

model=Axon.input("input",shape: {nil,784})|>Axon.dense(256,activation: :relu)|>Axon.dropout(rate: 0.2)|>Axon.dense(10)compiled=ExBurn.Model.compile(model,loss: :cross_entropy,optimizer: :adam,learning_rate: 0.001)trained=ExBurn.Training.fit(compiled,{train_x,train_y},epochs: 10,batch_size: 32,validation_data: {val_x,val_y},callbacks: [&ExBurn.Training.LoggingCallback.log/1]){:ok,output}=ExBurn.Model.forward(trained,input_tensor)

6. Batched inference with Nx.Serving

serving=ExBurn.Serving.new(trained,batch_size: 8,batch_timeout: 10)%Nx.Serving{}=result=Nx.Serving.run(serving,Nx.Batch.stack([input1,input2]))output=result.output

Prerequisites

• Elixir ~> 1.18 and OTP 27+
• Rust stable toolchain (required for NIF compilation)

  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

• For iOS development: Xcode + aarch64-apple-ios target

  rustup target add aarch64-apple-ios

• For Android development: Android NDK + aarch64-linux-android target

  rustup target add aarch64-linux-android

Note: Precompiled NIF binaries are planned. Until then, a Rust toolchain is required to build the NIF from source.

Training on Mobile — Caveats

Burn's Autodiff backend is memory-intensive. On iOS/Android with limited RAM, training even small models may cause out-of-memory errors. Realistic expectations:

• Fine-tuning small models (< 10M parameters) is feasible on modern devices
• Full training of large models is not recommended on mobile
• Inference is the primary use case for mobile deployment
• Minimum recommended: 4GB RAM, A12+ chip (iOS) / Snapdragon 700+ (Android)

The training loop in ExBurn currently uses numerical gradients (finite differences). Two methods are available: :numerical (central differences, more accurate) and :numerical_batch (one-sided, ~2x faster). Burn's autodiff integration is planned for v0.3.0 and will replace numerical gradients entirely.

Examples

# Linear regression (simplest possible ML workflow)
mix run examples/linear_regression.exs
# MNIST-like classifier (full deep learning pipeline)
mix run examples/mnist_simple.exs
# XOR classifier (non-linear problem, early stopping, model summary)
mix run examples/xor_classifier.exs
# Dataset utilities (split, normalize, one-hot, data loaders)
mix run examples/dataset_utils.exs
# BurnBridge direct tensor operations (arithmetic, math, linear algebra)
mix run examples/burn_bridge_ops.exs
# Model management (save/load, quantize, freeze, benchmark, export)
mix run examples/model_management.exs
# Training callbacks (logging, early stopping, checkpoint, LR scheduling)
mix run examples/training_callbacks.exs

Benchmarks

Benchmark scripts in bench/ compare ExBurn (Burn GPU backend) against plain Nx (BinaryBackend) across tensor sizes from 10×10 to 2000×2000.

# Tensor creation (zeros, ones, rand)
mix run bench/tensor_creation_bench.exs
# Element-wise arithmetic (add, mul, exp)
mix run bench/arithmetic_bench.exs
# Linear algebra (matmul, transpose)
mix run bench/linear_algebra_bench.exs
# Nx <-> Burn tensor conversion overhead
mix run bench/conversion_bench.exs
# End-to-end training (small/medium MLPs, optimizer comparison)
mix run bench/training_bench.exs
# Inference latency and throughput (single + batched + Nx.Serving)
mix run bench/serving_bench.exs

Each script prints a formatted table with timing results. All benchmarks include warmup runs and report averaged measurements.

Testing

mix test# full suite (~600 tests, <30s)
mix test test/backend_test.exs # single file
mix test --only nif # only NIF-backed tests
mix test --cover # suite + coverage report (threshold enforced)

Test tags (:nif, :cuda, :metal, :vulkan) are excluded automatically when the NIF or matching GPU hardware isn't available — see test/test_helper.exs.

Troubleshooting

Symptom	Likely cause	Fix
:erlang.nif_error(:nif_not_loaded)	Native library wasn't compiled/linked	mix clean && mix compile; check Rust is installed
dtype :f64 is not supported by the NIF	Non-f32 dtype crossed the raw NIF boundary	Convert via Nx.as_type(t, {:f, 32}) or go through Nx with ExBurn.Backend
{:error, "Erlang error: :nif_panicked"}	Rust-side assertion (e.g. rank/dtype mismatch on a direct NIF call)	Re-run with RUST_BACKTRACE=1; check the NIF constraints table
gpu_available: false	NIF built CPU-only, or drivers missing	Rebuild with ./build.sh metal / cuda / vulkan
First compile takes minutes	Debug-mode Rust build of Burn/CubeCL	Expected; release mode is used for production builds

Contributing

See CONTRIBUTING.md for the full development workflow: environment setup, running tests and coverage, linting, and a step-by-step walkthrough for adding new operations.

Guides

• Getting Started — Installation, basic ops, GPU check
• Training Models — Models, training, callbacks, save/load
• Mobile Deployment — iOS/Android compilation, optimization
• Architecture — Deep-dive into the pipeline
• Training Optimization — Best practices for fast, stable training
• Deep Learning Guide — Step-by-step lessons for learning DL with ExBurn
• Benchmarks — Performance benchmarks and how to run them
• Contributing — Development workflow, testing, adding operations

Project Structure

lib/ex_burn/
ex_burn.ex — Main API (version, configure!, default_device)
defn_compiler.ex — Nx.Defn.Compiler for GPU-accelerated defn
backend.ex — Nx.Backend implementation (delegates to Burn via NIF)
nif.ex — Rustler NIF stubs (40+ functions)
tensor.ex — Nx ↔ Burn tensor conversion utilities
error.ex — Structured error type (ExBurn.Error)
burn_bridge.ex — High-level Burn API (direct tensor ops)
cubecl_bridge.ex — GPU compute via ExCubecl (buffers, kernels, pipelines)
model.ex — Model definition, compilation, save/load
training.ex — Training loop (optimizers, LR schedules, callbacks)
native/ex_burn_nif/
src/lib.rs — Rust NIF with real Burn Autodiff<CubeCL> operations
Cargo.toml — Burn 0.21 + CubeCL + Autodiff dependencies
examples/
linear_regression.exs — Simplest ML workflow
mnist_simple.exs — Full deep learning pipeline
xor_classifier.exs — Non-linear classification with early stopping
dataset_utils.exs — Data preprocessing utilities
burn_bridge_ops.exs — Direct Burn tensor operations
model_management.exs — Save/load, quantize, freeze, benchmark
training_callbacks.exs — All callback types + custom callbacks
bench/
tensor_creation_bench.exs — zeros/ones/rand: Nx vs Burn
arithmetic_bench.exs — add/mul/exp: Nx vs Burn
linear_algebra_bench.exs — matmul/transpose: Nx vs Burn
conversion_bench.exs — Nx<->Burn conversion overhead
training_bench.exs — End-to-end training performance
serving_bench.exs — Inference latency & throughput
guides/
01_getting_started.md — Installation, basic ops, GPU check
02_training.md — Models, training, callbacks, save/load
03_mobile_deployment.md — iOS/Android compilation, optimization
04_architecture.md — Deep-dive into the pipeline

GPU Backends

Platform	Backend	Status
NVIDIA	CUDA	✅
iOS	Metal	✅
Android	Vulkan	✅
macOS	Metal	✅
Linux	Vulkan	✅

CUDA Support

The NIF compiles without a GPU backend by default (CPU-only NdArray). To build with GPU acceleration, use the bundled helper script, which auto-detects the best backend for your platform:

./build.sh # auto-detect: cuda / metal / vulkan / cpu
./build.sh metal # force a specific backend

Or set the Cargo feature manually before compiling:

RUSTLER_NIF_CARGO_FEATURES=cuda mix compile

On systems without the requested GPU hardware, the NIF automatically falls back to the NdArray (CPU) backend at runtime.

Check CUDA availability from Elixir:

ExBurn.cuda_available?()# true if NVIDIA GPU detectedExBurn.device_name()# "CUDA (NVIDIA GPU)" or "NdArray (CPU)"ExBurn.device_info()# full device info map

Error Handling

All operations raise ExBurn.Error with structured context:

raiseExBurn.Error,op: :matmul,reason: "shape mismatch",details: %{lhs: [3,4],rhs: [5,6]}

Dependencies

• Burn — Deep learning framework (Rust)
• Nx — Numerical Elixir
• Axon — Neural network library
• CubeCL — GPU compute language
• ExCubecl v0.5+ — GPU compute runtime for Elixir (buffers, kernels, pipelines, media)

---

Topics: elixir · machine-learning · burn · ios · android · nx · rustler · gpu · deep-learning

License

Apache 2.0