Dala integration for the Burn deep learning framework.
ExBurn provides a Nx.Backend implementation that delegates tensor operations
to Burn via Rust NIFs, enabling GPU-accelerated ML/DL on mobile and desktop.
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 → GPUQuick Start
# Set ExBurn as the default Nx backend
Dala.ML.Burn.configure!()
# Create and manipulate tensors
t = Nx.tensor([1.0, 2.0, 3.0])
Nx.add(t, t) |> Nx.to_list()
# Define a model with Axon
model =
Axon.input("input", shape: {nil, 784})
|> Axon.dense(256, activation: :relu)
|> Axon.dropout(rate: 0.2)
|> Axon.dense(10)
# Compile for training
compiled = Dala.ML.Burn.compile(model,
loss: :cross_entropy,
optimizer: :adam,
learning_rate: 0.001
)
# Train
Dala.ML.Burn.fit(compiled, {train_x, train_y},
epochs: 10,
batch_size: 32
)Platform GPU Backends
| Platform | Backend | Status |
|---|---|---|
| iOS | Metal | ✅ |
| Android | Vulkan | ✅ |
| macOS | Metal | ✅ |
| Linux | Vulkan | ✅ |
| NVIDIA | CUDA | ✅ |
Integration with Dala.ML
This module complements the existing Dala ML backends:
Dala.ML.EMLX— MLX backend for Apple Silicon (iOS recommended)Dala.ML.CoreML— iOS-native CoreML (Neural Engine)Dala.ML.ONNX— Cross-platform ONNX RuntimeDala.ML.Burn— Burn framework via ExBurn (this module)
Use Dala.ML.available_backends/0 to see all available backends,
and Dala.ML.Burn.available?/0 specifically for Burn support.
Training on Mobile — Caveats
Burn's Autodiff backend is memory-intensive. On iOS/Android with limited RAM:
- 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)
Summary
Functions
Absolute value of a Burn tensor.
Adds two Burn tensors.
Checks whether ExBurn is available (loaded and functional).
Returns a list of available GPU backends on this system.
Benchmarks the model's forward pass on the given input.
Creates a GPU buffer via ExCubecl from a list of values.
Creates a GPU buffer via ExCubecl, raising on error.
Returns the shape of an ExCubecl buffer.
Returns the byte size of an ExCubecl buffer.
Creates a deep copy of the model with identical parameters and configuration.
Compiles an Axon model for training with the ExBurn backend.
Computes gradients for a given mini-batch.
Computes the loss between predictions and targets.
Configures ExBurn for the current platform with Dala-specific defaults.
Cross-entropy loss between predictions and targets.
Checks whether an NVIDIA CUDA GPU is available.
Creates a data loader that yields mini-batches from a dataset.
Returns the default device for tensor operations.
Deserializes model parameters from a binary.
Returns a map with device information including GPU availability, backend name, and available backends.
Returns the name of the active compute device (e.g., "CUDA (NVIDIA GPU)").
Returns a human-readable summary of the GPU device.
Divides two Burn tensors element-wise.
Dropout (identity during inference).
Enables the ExBurn defn compiler for GPU-accelerated Nx.Defn expressions.
Creates an ExBurn.Error struct (non-raising).
Wraps an error tuple in an ExBurn.Error.
Converts an ExBurn.Error to a log string.
Evaluates a model on a dataset.
Exponential of a Burn tensor.
Exports model parameters to a file.
Trains a model on the given dataset.
Formats an ExBurn.Error for logging or display.
Runs a forward pass through the model using the ExBurn GPU defn compiler.
Returns the output shape and type information from the Axon model.
Frees a Burn tensor's underlying GPU/CPU memory.
Freezes the specified layers so their parameters are not updated during training.
Creates a Burn tensor from raw binary data.
Converts an Nx tensor to a Burn tensor.
Batch converts a list of Nx tensors to Burn tensors.
Checks whether a layer is frozen.
Returns the set of frozen layer names.
Checks whether a GPU device is available for Burn operations.
Checks whether a GPU is available via Burn bridge.
Returns GPU memory info as %{total: bytes, used: bytes, free: bytes}.
Imports model parameters from a file saved with export/3.
Returns a map with model information (param count, layer count, device, memory estimate).
Layer normalization.
Loads model parameters from a file.
Natural logarithm of a Burn tensor.
Matrix multiplication of two Burn tensors.
Mean of all elements in a Burn tensor.
Mean squared error between predictions and targets.
Multiplies two Burn tensors element-wise.
Negates a Burn tensor.
Creates a new empty model struct. Useful for incremental model building.
Returns the number of NIF functions registered by the Rust library. Useful for debugging NIF loading issues.
Checks whether the NIF library is loaded and responds to calls.
Creates a tensor filled with ones via Burn.
Returns the current model parameters.
Runs a forward pass through the model using Axon's default backend.
Profiles a single training step, returning detailed timing for each phase.
Quantizes model parameters to a lower precision type (:f16 or :bf16).
Creates a random tensor with uniform distribution via Burn.
Reads data from an ExCubecl buffer.
ReLU activation of a Burn tensor.
Reshapes a Burn tensor.
Saves the model parameters to a file using compressed Erlang term format.
Serializes model parameters to a binary for network transfer or storage.
Sigmoid activation of a Burn tensor.
Performs a quick smoke test of the ExBurn pipeline.
Softmax along a dimension.
Square root of a Burn tensor.
Subtracts two Burn tensors.
Sum of all elements in a Burn tensor.
Returns a summary of the ExBurn environment.
Returns a summary of the model architecture including parameter count.
Returns the total number of elements in a Burn tensor.
Returns the rank (number of dimensions) of a Burn tensor.
Returns the shape of a Burn tensor.
Returns the element type of a Burn tensor.
Moves a Burn tensor to the CPU.
Moves all model parameters to the specified device (:gpu or :cpu).
Moves a Burn tensor to the GPU.
Converts a Burn tensor to an Nx tensor.
Batch converts a list of Burn tensors to Nx tensors.
Performs a single training step: forward + backward + optimizer update.
Transposes a Burn tensor.
Unfreezes the specified layers so their parameters are updated during training.
Returns a new model with updated parameters.
Returns the current ExBurn version.
Creates a tensor filled with zeros via Burn.
Functions
@spec abs(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Absolute value of a Burn tensor.
@spec add(ExBurn.Tensor.t(), ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Adds two Burn tensors.
@spec available?() :: boolean()
Checks whether ExBurn is available (loaded and functional).
@spec available_backends() :: [atom()]
Returns a list of available GPU backends on this system.
@spec benchmark(ExBurn.Model.t(), Nx.Tensor.t(), keyword()) :: map()
Benchmarks the model's forward pass on the given input.
Options
:warmup— Number of warmup runs (default: 3):runs— Number of benchmarked runs (default: 10)
Returns a map with :avg_ms, :min_ms, :max_ms, :median_ms, :std_ms.
@spec buffer(list(), [non_neg_integer()], atom()) :: ExCubecl.buffer_ref()
Creates a GPU buffer via ExCubecl from a list of values.
@spec buffer!(list(), [non_neg_integer()], atom()) :: ExCubecl.buffer_ref()
Creates a GPU buffer via ExCubecl, raising on error.
@spec buffer_shape(ExCubecl.buffer_ref()) :: [non_neg_integer()]
Returns the shape of an ExCubecl buffer.
@spec buffer_size(ExCubecl.buffer_ref()) :: non_neg_integer()
Returns the byte size of an ExCubecl buffer.
@spec clone(ExBurn.Model.t()) :: ExBurn.Model.t()
Creates a deep copy of the model with identical parameters and configuration.
@spec compile( Axon.ModelState.t(), keyword() ) :: ExBurn.Model.t()
Compiles an Axon model for training with the ExBurn backend.
Options
:loss— Loss function::cross_entropy,:mse,:binary_cross_entropy(default::cross_entropy):optimizer— Optimizer::adam,:sgd,:rmsprop(default::adam):learning_rate— Learning rate (default: 0.001):device— Device::cpuor:gpu(default: auto-detected):weight_decay— L2 regularization coefficient (default: 0.0)
@spec compute_gradients(ExBurn.Model.t(), {Nx.Tensor.t(), Nx.Tensor.t()}, keyword()) :: map()
Computes gradients for a given mini-batch.
Options
:grad_method—:numerical(central differences) or:numerical_batch(one-sided, faster):epsilon— Finite difference step size (default: 1.0e-5)
@spec compute_loss(ExBurn.Model.t(), Nx.Tensor.t(), Nx.Tensor.t()) :: {:ok, Nx.Tensor.t()} | {:error, term()}
Computes the loss between predictions and targets.
@spec configure!(keyword()) :: :ok
Configures ExBurn for the current platform with Dala-specific defaults.
Options
:device— Override device (:cpuor:gpu). Auto-detected by default.:backend— Override GPU backend (:metal,:vulkan,:cuda). Auto-detected.
@spec cross_entropy(ExBurn.Tensor.t(), ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Cross-entropy loss between predictions and targets.
@spec cuda_available?() :: boolean()
Checks whether an NVIDIA CUDA GPU is available.
@spec data_loader({Nx.Tensor.t(), Nx.Tensor.t()}, keyword()) :: Enumerable.t()
Creates a data loader that yields mini-batches from a dataset.
Options
:batch_size— Mini-batch size (default: 32):shuffle— Shuffle data each iteration (default: true)
@spec default_device() :: :cpu | :gpu
Returns the default device for tensor operations.
Deserializes model parameters from a binary.
@spec device_info() :: map()
Returns a map with device information including GPU availability, backend name, and available backends.
@spec device_name() :: String.t()
Returns the name of the active compute device (e.g., "CUDA (NVIDIA GPU)").
@spec device_summary() :: String.t()
Returns a human-readable summary of the GPU device.
@spec div(ExBurn.Tensor.t(), ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Divides two Burn tensors element-wise.
@spec dropout(ExBurn.Tensor.t(), float(), boolean()) :: ExBurn.Tensor.t()
Dropout (identity during inference).
@spec enable_defn_compiler!() :: :ok
Enables the ExBurn defn compiler for GPU-accelerated Nx.Defn expressions.
After calling this, all defn functions will be compiled through
ExBurn.Defn.Compiler and executed on the GPU via Burn.
Example
Dala.ML.Burn.enable_defn_compiler!
defmodule MyMath do
import Nx.Defn
defn add_and_scale(x, y, scale) do
x |> Nx.add(y) |> Nx.multiply(scale)
end
end
# Runs on GPU via Burn
MyMath.add_and_scale(Nx.tensor([1.0]), Nx.tensor([2.0]), Nx.tensor(3.0))@spec error(keyword()) :: ExBurn.Error.t()
Creates an ExBurn.Error struct (non-raising).
Example
Dala.ML.Burn.error(op: :forward, reason: "shape mismatch")@spec error_from_tuple({:error, String.t()}, keyword()) :: ExBurn.Error.t()
Wraps an error tuple in an ExBurn.Error.
Example
Dala.ML.Burn.error_from_tuple({:error, "failed"}, op: :predict)@spec error_to_log_string(ExBurn.Error.t()) :: String.t()
Converts an ExBurn.Error to a log string.
@spec evaluate(ExBurn.Model.t(), {Nx.Tensor.t(), Nx.Tensor.t()}, keyword()) :: float() | {float(), float() | nil}
Evaluates a model on a dataset.
Returns the average loss, or {loss, accuracy} when track_accuracy: true.
@spec exp(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Exponential of a Burn tensor.
@spec export(ExBurn.Model.t(), Path.t(), keyword()) :: :ok | {:error, String.t()}
Exports model parameters to a file.
Options
:format—:elixir_terms(default, compressed) or:json(human-readable)
@spec fit(ExBurn.Model.t(), {Nx.Tensor.t(), Nx.Tensor.t()}, keyword()) :: ExBurn.Model.t()
Trains a model on the given dataset.
Options
:epochs— Number of training epochs (default: 10):batch_size— Mini-batch size (default: 32):shuffle— Shuffle training data each epoch (default: true):validation_data— Validation dataset as{inputs, targets}tuple:callbacks— List of callback functions called after each epoch:verbose— Print training progress (default: true):lr_schedule— Learning rate schedule (default: nil):clip_norm— Max gradient norm for clipping (default: nil):clip_value— Max absolute gradient value for clipping (default: nil):weight_decay— L2 regularization coefficient (default: nil):accumulate_gradients— Gradient accumulation steps (default: 1):accuracy— Track and report classification accuracy (default: false):nesterov— Use Nesterov momentum for SGD (default: false)
@spec format_error(ExBurn.Error.t()) :: String.t()
Formats an ExBurn.Error for logging or display.
@spec forward(ExBurn.Model.t(), Nx.Tensor.t()) :: {:ok, Nx.Tensor.t()} | {:error, term()}
Runs a forward pass through the model using the ExBurn GPU defn compiler.
This is the GPU-accelerated path. Requires the model to be compiled.
Returns {:ok, output_tensor} or {:error, reason}.
@spec forward_pattern(ExBurn.Model.t()) :: %{ output_shape: tuple() | nil, output_type: atom() }
Returns the output shape and type information from the Axon model.
@spec free(ExBurn.Tensor.t()) :: :ok
Frees a Burn tensor's underlying GPU/CPU memory.
@spec freeze(ExBurn.Model.t(), [atom() | String.t()]) :: ExBurn.Model.t()
Freezes the specified layers so their parameters are not updated during training.
@spec from_binary(binary(), [non_neg_integer()], atom()) :: {:ok, ExBurn.Tensor.t()} | {:error, term()}
Creates a Burn tensor from raw binary data.
@spec from_nx(Nx.Tensor.t()) :: {:ok, ExBurn.Tensor.t()} | {:error, term()}
Converts an Nx tensor to a Burn tensor.
@spec from_nx_batch([Nx.Tensor.t()]) :: {:ok, [ExBurn.Tensor.t()]} | {:error, term()}
Batch converts a list of Nx tensors to Burn tensors.
@spec frozen?(ExBurn.Model.t(), atom() | String.t()) :: boolean()
Checks whether a layer is frozen.
@spec frozen_layers(ExBurn.Model.t()) :: MapSet.t()
Returns the set of frozen layer names.
@spec gpu?() :: boolean()
Checks whether a GPU device is available for Burn operations.
@spec gpu_available?() :: boolean()
Checks whether a GPU is available via Burn bridge.
Returns GPU memory info as %{total: bytes, used: bytes, free: bytes}.
@spec import_params(ExBurn.Model.t(), Path.t(), keyword()) :: {:ok, ExBurn.Model.t()} | {:error, String.t()}
Imports model parameters from a file saved with export/3.
Options
:format—:elixir_terms(default) or:json
@spec info(ExBurn.Model.t()) :: map()
Returns a map with model information (param count, layer count, device, memory estimate).
@spec layer_norm(ExBurn.Tensor.t(), non_neg_integer(), float()) :: ExBurn.Tensor.t()
Layer normalization.
@spec load(ExBurn.Model.t(), Path.t()) :: {:ok, ExBurn.Model.t()} | {:error, term()}
Loads model parameters from a file.
@spec log(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Natural logarithm of a Burn tensor.
@spec matmul(ExBurn.Tensor.t(), ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Matrix multiplication of two Burn tensors.
@spec mean(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Mean of all elements in a Burn tensor.
@spec mse(ExBurn.Tensor.t(), ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Mean squared error between predictions and targets.
@spec mul(ExBurn.Tensor.t(), ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Multiplies two Burn tensors element-wise.
@spec neg(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Negates a Burn tensor.
@spec new_model(keyword()) :: ExBurn.Model.t()
Creates a new empty model struct. Useful for incremental model building.
@spec nif_function_count() :: non_neg_integer()
Returns the number of NIF functions registered by the Rust library. Useful for debugging NIF loading issues.
@spec nif_loaded?() :: boolean()
Checks whether the NIF library is loaded and responds to calls.
@spec ones([non_neg_integer()], atom()) :: ExBurn.Tensor.t()
Creates a tensor filled with ones via Burn.
@spec parameters(ExBurn.Model.t()) :: map()
Returns the current model parameters.
@spec predict(ExBurn.Model.t(), Nx.Tensor.t()) :: {:ok, Nx.Tensor.t()} | {:error, term()}
Runs a forward pass through the model using Axon's default backend.
For GPU execution, use forward/2 which uses the ExBurn defn compiler.
Returns {:ok, output_tensor} or {:error, reason}.
@spec profile_step(ExBurn.Model.t(), {Nx.Tensor.t(), Nx.Tensor.t()}, keyword()) :: map()
Profiles a single training step, returning detailed timing for each phase.
Returns a map with :forward_ms, :backward_ms, :optimizer_ms, :total_ms.
@spec quantize(ExBurn.Model.t(), :f16 | :bf16) :: ExBurn.Model.t()
Quantizes model parameters to a lower precision type (:f16 or :bf16).
Useful for reducing model size and speeding up inference on devices with limited compute.
@spec rand([non_neg_integer()], atom(), float(), float()) :: ExBurn.Tensor.t()
Creates a random tensor with uniform distribution via Burn.
@spec read_buffer(ExCubecl.buffer_ref()) :: binary()
Reads data from an ExCubecl buffer.
@spec relu(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
ReLU activation of a Burn tensor.
@spec reshape(ExBurn.Tensor.t(), [non_neg_integer()]) :: ExBurn.Tensor.t()
Reshapes a Burn tensor.
@spec save(ExBurn.Model.t(), Path.t()) :: :ok | {:error, term()}
Saves the model parameters to a file using compressed Erlang term format.
@spec serialize_params(ExBurn.Model.t()) :: binary()
Serializes model parameters to a binary for network transfer or storage.
@spec sigmoid(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Sigmoid activation of a Burn tensor.
@spec smoke_test() :: :ok | {:error, String.t()}
Performs a quick smoke test of the ExBurn pipeline.
Returns :ok on success or {:error, reason} on failure.
@spec softmax(ExBurn.Tensor.t(), non_neg_integer()) :: ExBurn.Tensor.t()
Softmax along a dimension.
@spec sqrt(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Square root of a Burn tensor.
@spec sub(ExBurn.Tensor.t(), ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Subtracts two Burn tensors.
@spec sum(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Sum of all elements in a Burn tensor.
@spec summary() :: String.t()
Returns a summary of the ExBurn environment.
@spec summary(ExBurn.Model.t()) :: String.t()
Returns a summary of the model architecture including parameter count.
@spec tensor_numel(ExBurn.Tensor.t()) :: non_neg_integer()
Returns the total number of elements in a Burn tensor.
@spec tensor_rank(ExBurn.Tensor.t()) :: non_neg_integer()
Returns the rank (number of dimensions) of a Burn tensor.
@spec tensor_shape(ExBurn.Tensor.t()) :: [non_neg_integer()]
Returns the shape of a Burn tensor.
@spec tensor_type(ExBurn.Tensor.t()) :: atom()
Returns the element type of a Burn tensor.
@spec to_cpu(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Moves a Burn tensor to the CPU.
@spec to_device(ExBurn.Model.t(), :cpu | :gpu) :: ExBurn.Model.t()
Moves all model parameters to the specified device (:gpu or :cpu).
@spec to_gpu(ExBurn.Tensor.t()) :: ExBurn.Tensor.t()
Moves a Burn tensor to the GPU.
@spec to_nx(ExBurn.Tensor.t()) :: {:ok, Nx.Tensor.t()} | {:error, term()}
Converts a Burn tensor to an Nx tensor.
@spec to_nx_batch([ExBurn.Tensor.t()]) :: {:ok, [Nx.Tensor.t()]} | {:error, term()}
Batch converts a list of Burn tensors to Nx tensors.
@spec train_step(ExBurn.Model.t(), {Nx.Tensor.t(), Nx.Tensor.t()}, keyword()) :: {float(), ExBurn.Model.t()}
Performs a single training step: forward + backward + optimizer update.
Useful for custom training loops. Returns {loss, updated_model}.
@spec transpose(ExBurn.Tensor.t(), non_neg_integer(), non_neg_integer()) :: ExBurn.Tensor.t()
Transposes a Burn tensor.
@spec unfreeze(ExBurn.Model.t(), [atom() | String.t()]) :: ExBurn.Model.t()
Unfreezes the specified layers so their parameters are updated during training.
@spec update_params(ExBurn.Model.t(), map()) :: ExBurn.Model.t()
Returns a new model with updated parameters.
@spec version() :: String.t()
Returns the current ExBurn version.
@spec zeros([non_neg_integer()], atom()) :: ExBurn.Tensor.t()
Creates a tensor filled with zeros via Burn.