Compiling and Graph Breaks

This guide explains how Transmuter.compile/3 works, what graph breaks are, and how to control them with static inputs.

How compilation works

Transmuter.compile/3 compiles the model graph for the shapes and types of the inputs you give:

run = Transmuter.compile(model, inputs, compiler: EMLX)

It returns a function. The function takes the same inputs and returns the outputs:

outputs = run.(inputs)

Compilation turns the Nx operations in the model into a single program that the backend (EMLX or EXLA) runs. Later calls with the same input shapes run the compiled program directly. Calls with different input shapes compile those shapes on their first call.

The :cache_size option sets how many compiled shapes to keep per graph segment. The default is 8. When you reach the limit, the least recently used shape is removed. Set :cache_size to a larger value when the inputs have many shapes that come back often, such as audio lengths in a speech service:

run = Transmuter.compile(model, inputs, compiler: EMLX, cache_size: 64)

What a graph break is

A compiled graph must know all its shapes before it runs. Some operators use a tensor value to decide a shape or to select a branch. Examples:

• Range uses the upper limit to set the length of its output.
• NonZero returns the positions of all non-zero elements. The count is not known until the values arrive.
• If selects one of two subgraphs based on a condition tensor.
• NonMaxSuppression returns the boxes that pass the threshold. The count depends on the input scores.

When the compiler reaches one of these operators and the value is not yet known (because it depends on the model inputs), it raises Transmuter.UnknownValueError. Transmuter catches this error and creates a graph break at that operator.

Transmuter handles a graph break in three steps:

1. It compiles and runs the operators before the break.
2. It runs the break operator on the real output values of step 1.
3. It compiles the operators after the break for the shapes that the values decide, and continues.

Compiled segments are cached. The key of a segment holds the shapes of its arguments, the values of the static inputs, and the host values that the segment reads (see below). A later call with the same shapes and the same values does not compile again. A new value of a host value that a segment reads compiles that segment again. The break operators still run eagerly on each call. A segment of fewer than 20 nodes after a break is not compiled, and it also runs eagerly on each call.

What goes into the key of a segment

The cache key of a segment has these parts:

• The position of the segment in the graph.
• The shapes and types of its arguments.
• The values of the static inputs and of the host values that it reads.

The small values that a break node reads, such as the positions that a NonZero returns, become host values. Values that the graph computes from them on the host are host values too. A host value goes into the next segment as an argument when no node of that segment reads it on the host. Then its value is not in the key. For example, a Gather that reads the word positions of a text model compiles once for each number of words, and not once for each text.

A host value stays static in a segment when a node of that segment needs it on the host, for example as the shape of a Reshape or the start of a Slice. Then its value is in the key of that segment, and each new value compiles that segment again. The keys of the other segments do not change. Transmuter finds these values when it traces a segment: the node breaks on the argument, and Transmuter traces the segment again with the value static. A small value that an earlier segment computes from host values, such as the first of the positions, is copied to the host when the segment starts. A value that is too large to copy, or that comes from a random operator or from device data in this call, stays an argument, and the node breaks there as a graph break. While Transmuter learns these values, the first calls with new values can compile a segment once more. After that, a new value compiles a segment again only when it is in the key. A segment that runs eagerly does not compile, so a new value in its key costs only a trace. The :host_values of the [:transmuter, :segment, :compile, :stop] event give the names that a trace found (see Transmuter.Telemetry).

Static inputs always stay static, so each new value of a static input is in the key.

Reading the error

When a graph break occurs, Transmuter re-raises the error if it cannot handle it:

Transmuter.UnknownValueError: NonZero NonZero_0: input_ids is only known when the model runs.
Pass the graph inputs it depends on as static_inputs, or call Transmuter.run/4 without compiling.

The error names the operator type, the node name, and the tensor name. To fix it, either pass the tensor as a static input (see below) or use Transmuter.run/4 without compilation.

Static inputs

A static input is a tensor whose value is known when the model compiles. Pass static inputs in the :static_inputs option:

run = Transmuter.compile(
  model,
  inputs,
  static_inputs: %{"use_cache_branch" => Nx.tensor([1], type: :u8, backend: Nx.BinaryBackend)},
  compiler: EMLX
)

Give static inputs as Nx.BinaryBackend tensors. Each new value of a static input compiles the model again.

Common uses for static inputs:

• use_cache_branch in merged decoder exports. It selects an If branch for the prefill or the decode step.
• num_logits_to_keep in some decoder exports. It sets the length of a slice.
• input_ids in Grounding DINO. They feed NonZero, which counts the non-zero elements.

When a static input changes, Transmuter compiles a new version of the model. With :cache_size segments cached, a value that comes back later reuses its compiled version.

See the segments

Transmuter.Graph shows where the graph breaks and why, which nodes fold on the host, and which subgraphs Transmuter.Fusion replaced. With :inputs, it compiles the model and calls it once:

overlay = Transmuter.Graph.overlay(model, inputs: inputs, compile: [compiler: EXLA])
overlay.segments
# => [%{start: 0, stop: 2, compiled: true, reason: "a value that decides a shape"},
#     %{start: 3, stop: :end, compiled: false, reason: nil}]

In Livebook, Transmuter.Graph.as_graph/2 shows the same facts in an interactive view, with one fill color for each segment. Click a group to open it, and click a node to see its details, such as the shapes of its outputs. It needs the optional kino dependency. Outside Livebook, Transmuter.Graph.to_html/3 writes the view to one HTML file:

Transmuter.Graph.as_graph(model, inputs: inputs, compile: [compiler: EXLA])
Transmuter.Graph.to_html(model, "graph.html", inputs: inputs, compile: [compiler: EXLA])

The view works without a network connection. Without :inputs, the view is static only and shows no segments.

The [:transmuter, :segment, :compile, :stop] event also gives the reason of each break (see Transmuter.Telemetry). For the graph as the exporter wrote it, load the model with fuse: false and call Transmuter.Graph.to_html/3 without :inputs, or run mix transmuter.graph REPO --plain.

Limit the number of compiled segments

If the total number of compiled segments grows large (for example, with many unique audio lengths in Kokoro), set a :cache_size that matches the number of expected shapes:

run = Transmuter.compile(model, inputs, compiler: EMLX, cache_size: 64)

When the cache is full, Transmuter removes the least recently used segment. A new audio length triggers one recompilation. Later calls with the same length reuse the compiled segment.

Eager execution

Transmuter.run/4 runs the model without compilation. Call it directly when you do not need compiled speed, when you need to read intermediate values, or when your graph has non-tensor outputs such as strings or sequences:

outputs = Transmuter.run(model, model.params, inputs)

Give Transmuter.run/4 to Nx.Defn.jit/2 to compile it yourself:

jit_run = Nx.Defn.jit(&Transmuter.run(model, &1, &2))
outputs = jit_run.(model.params, inputs)

Strings, sequences, and empty optionals

Nx tensors hold only numbers. Transmuter keeps other ONNX values as Elixir terms:

A compiled function can only return Nx.Tensor values. To read a string or sequence output, use Transmuter.run/4 without compilation.

Give string and sequence inputs as static inputs with :static_inputs.

Random operators

Random operators draw from Nx.Random. Each call of a compiled model draws new values. A node with a seed attribute draws the same values on each call.