Tensor conversion utilities between Nx and Burn formats.
Handles marshaling of tensor data between Elixir's Nx tensor representation and the Rust/Burn tensor references used by the NIF.
This module is the single source of truth for Nx ↔ Burn dtype mapping. Do not duplicate these conversions elsewhere.
Type Mapping
| Nx Type | Burn Type |
|---|---|
{:f, 32} | :f32 |
{:f, 64} | :f64 |
{:f, 16} | :f16 |
{:bf, 16} | :bf16 |
{:s, 32} | :i32 |
{:s, 64} | :i64 |
{:s, 16} | :i16 |
{:s, 8} | :i8 |
{:u, 8} | :u8 |
Note
The mapping above is faithful, but the current NIF layer only
stores :f32 data. Creating a tensor with any other dtype raises
a structured ExBurn.Error at the NIF boundary rather than
silently reinterpreting bytes.
Summary
Functions
Converts a Burn type tag back to an Nx type tuple.
Frees the underlying Rust tensor.
Creates a Burn tensor from raw binary data, shape, and type.
Converts an Nx.Tensor.t() into an ExBurn.Tensor.t().
Batch converts a list of Nx tensors to Burn tensors.
Returns the total number of elements.
Converts an Nx type tuple to a Burn type tag (faithful, no precision loss).
Returns the rank (number of dimensions).
Returns the NIF reference for a tensor.
Returns the shape of a Burn tensor.
Converts an ExBurn.Tensor.t() back into an Nx.Tensor.t().
Batch converts a list of Burn tensors to Nx tensors.
Returns the Burn element type of a tensor.
Types
@type burn_type() :: :f32 | :f64 | :f16 | :bf16 | :i32 | :i64 | :i16 | :i8 | :u8
Burn element type tag
@type t() :: %ExBurn.Tensor{ ref: reference(), shape: [non_neg_integer()], type: burn_type() }
NIF tensor reference
Functions
Converts a Burn type tag back to an Nx type tuple.
@spec free(t()) :: :ok
Frees the underlying Rust tensor.
@spec from_binary(binary(), [non_neg_integer()], burn_type()) :: {:ok, t()} | {:error, String.t()}
Creates a Burn tensor from raw binary data, shape, and type.
@spec from_nx(Nx.Tensor.t()) :: {:ok, t()} | {:error, String.t()}
Converts an Nx.Tensor.t() into an ExBurn.Tensor.t().
The tensor data is sent to the Rust NIF layer as a flat binary.
Returns {:ok, tensor} or {:error, reason}.
@spec from_nx_batch([Nx.Tensor.t()]) :: {:ok, [t()]} | {:error, String.t()}
Batch converts a list of Nx tensors to Burn tensors.
More efficient than calling from_nx/1 individually when you need
to convert many tensors.
@spec numel(t()) :: non_neg_integer()
Returns the total number of elements.
Converts an Nx type tuple to a Burn type tag (faithful, no precision loss).
@spec rank(t()) :: non_neg_integer()
Returns the rank (number of dimensions).
Returns the NIF reference for a tensor.
@spec shape(t()) :: [non_neg_integer()]
Returns the shape of a Burn tensor.
@spec to_nx(t()) :: {:ok, Nx.Tensor.t()} | {:error, String.t()}
Converts an ExBurn.Tensor.t() back into an Nx.Tensor.t().
Reads the raw data from the Rust NIF layer and reshapes it.
Returns {:ok, tensor} or {:error, reason}.
@spec to_nx_batch([t()]) :: {:ok, [Nx.Tensor.t()]} | {:error, String.t()}
Batch converts a list of Burn tensors to Nx tensors.
Returns the Burn element type of a tensor.