ExBurn.Defn.Compiler (ex_burn v0.3.0)

Copy Markdown View Source

Custom Nx.Defn.Compiler that compiles defn expressions to run on the Burn GPU backend via NIF.

This compiler traverses the Nx.Defn.Expr expression tree, converts it into Burn tensor operations, and executes them through the Rust NIF layer for GPU acceleration.

Usage

Nx.Defn.global_default_options(compiler: ExBurn.Defn.Compiler)

Or per-function:

defn my_fun(x, y, opts \\ []) do
  Nx.add(x, y)
end
compiler: ExBurn.Defn.Compiler

How it works

  1. Nx.Defn traces the function body into an expression tree of Nx.Defn.Expr nodes.
  2. The compiler receives the expression tree via __compile__/4 or __jit__/5.
  3. Each expression node is evaluated by converting parameters to Burn tensors, then dispatching the operation to the NIF.
  4. The result is converted back to Nx.Tensor for the caller.

Options

  • :garbage_collect - when true, garbage collects after evaluating each node (default: false)