Transmuter ⚗️

Transmuter runs Hugging Face ONNX models on any Nx backend, with no per-architecture code. Use EMLX on Apple silicon, or EXLA on CUDA or on the CPU.

The goal of Transmuter is to implement the ONNX operators in Nx. Then ONNX models run without model code.

An ONNX file contains a graph of operators. An operator is one small, standard computation, such as MatMul, Conv, Softmax, Add or Reshape. Each operator has inputs, outputs and attributes, for example the kernel size of a Conv. The ONNX standard defines about 200 operators, and a model connects a few hundred to a few thousand of them. When a model is exported, the exporter traces its forward pass and writes these operators to the file.

Transmuter implements each operator one time, as Nx operations. It then runs a model by running its graph, operator by operator. So Transmuter does not port model architectures. A new architecture needs no new Elixir code. It needs only operators that Transmuter already has.

The gap

Elixir has three ways to run Hugging Face models today. Each way has a limit:

• Bumblebee runs models with great accuracy on all Nx backends. But a person must port each architecture to Axon by hand, and a port takes weeks of work and review. Hugging Face gets new architectures faster than people can port them.
• Ortex runs ONNX models with ONNX Runtime. But it runs outside of Nx: you cannot put it in a defn function, it has no Apple silicon GPU path, and it is in maintenance mode.
• axon_onnx changed ONNX graphs into Axon models. It stopped in 2024. It set all shapes when it loaded a model, and many ONNX graphs need shapes that change.

Transmuter fills this gap. It runs the exported graph directly as Nx operations, so:

• Models that Bumblebee does not have run with no new code. Examples are SAM, Depth Anything, Grounding DINO, OWLv2, Florence-2, RF-DETR, Kokoro and Stable Diffusion.
• The result is plain Nx. It runs on EMLX and EXLA, it works with Nx.Serving, and Nx.Defn.grad can differentiate it.
• Shapes are set when the function compiles, not when the model loads. Thus decoders with a KV cache and graphs with If branches work.

When a model has no ONNX file on Hugging Face, export it one time with Optimum. In the 300 most downloaded transformers models, 62% of downloads already have an ONNX file on Hugging Face.

Example

model = Transmuter.load!({:hf, "Xenova/all-MiniLM-L6-v2", "onnx/model.onnx"}, backend: EMLX.Backend)

inputs = %{
  "input_ids" => Nx.tensor([[101, 7592, 2088, 102]]),
  "attention_mask" => Nx.tensor([[1, 1, 1, 1]]),
  "token_type_ids" => Nx.tensor([[0, 0, 0, 0]])
}

run = Transmuter.compile(model, inputs, compiler: EMLX)
%{"last_hidden_state" => hidden} = run.(inputs)

# Sentence embedding: mean of the token vectors, then L2 normalization.
mask = Nx.new_axis(inputs["attention_mask"], -1)
embedding = Nx.sum(Nx.multiply(hidden, mask), axes: [1]) |> Nx.divide(Nx.sum(mask, axes: [1]))
embedding = Nx.divide(embedding, Nx.LinAlg.norm(embedding, axes: [1], keep_axes: true))

The export sets the names of the inputs and outputs. model.inputs and model.outputs show them. For this export, model.outputs is ["last_hidden_state"].

Transmuter runs the graph that is in the file, and nothing more. The graph of all-MiniLM-L6-v2 stops at the token vectors, because sentence-transformers adds the pooling outside of the model. Thus the example above does the pooling. Bumblebee servings, such as Bumblebee.Text.text_embedding/3, also do the tokenization, pooling and batching for you. Transmuter has no servings yet. For tokenization, use tokenizers.

Transmuter.compile/3 compiles the model for the shapes of the inputs that you give. The function that it returns also takes inputs of other shapes: it compiles them on their first call, and keeps the most recently used shapes compiled. The :cache_size option sets how many, and the default is 8.

Transmuter.run/4 is the function that Transmuter.compile/3 compiles. It reads the ONNX graph and calls one Nx function for each operator. You can use it in two ways:

• Call it directly. Each operator runs immediately, one after the other. This is eager execution. It is slower, but the model can read values while it runs.
• Give it to Nx.Defn.jit/2 or Nx.Defn.value_and_grad/2. Nx then records the operators as one graph, and the backend compiles that graph.

Text to speech with Kokoro

bench/kokoro_tts.exs runs Kokoro-82M from its ONNX export and writes a WAV file. All operators of the export run in Transmuter, and Transmuter.compile/3 compiles the model with no static inputs. Inputs up to the limit of the model work: 505 phonemes gave 26.6 seconds of audio.

Kokoro reads phonemes, not text. The Python and JavaScript versions of Kokoro make the phonemes with misaki or espeak-ng, outside of the model. Elixir has no phonemizer yet, so the example contains phonemes that were made in advance.

The example also fixes one node of the export. The export computes the STFT phase with the atan2 of torch.onnx, which gives -π for atan2(0, x) when x is negative. torch.atan2 gives π, and the model was trained with π. Transmuter computes an exact STFT, so the imaginary part is 0 at DC and Nyquist, and the wrong phase in every frame adds a tone at 4.8 kHz. ONNX Runtime hides the error, because its STFT has small rounding errors at these bins. The example changes the Greater test of the atan2 to GreaterOrEqual, as in torch.atan2.

Installation

def deps do
  [
    {:transmuter, "~> 0.1"},
    {:emlx, "~> 0.5"}
  ]
end

Use EXLA instead of EMLX for CUDA.

Values that decide shapes

Some graphs use a value to set a shape or to select a branch. Examples are the length of the audio in Kokoro text-to-speech, which comes from the predicted duration of each phoneme, and the boxes that NonMaxSuppression keeps in Faster R-CNN. A compiled graph must know all its shapes before it runs, so Transmuter.compile/3 breaks the graph at these operators:

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

The compiled parts are cached, so a later call with the same shapes runs only compiled code. The cache keeps the 8 most recently used compiled parts for each break, or the number in the :cache_size option. A new shape after that replaces the part that was used least recently. When the shapes come back often, as with the audio lengths of a speech service, set :cache_size to the number of shapes, so that calls do not compile the same shapes again. The cache belongs to the process that calls Transmuter.compile/3, for example the Nx.Serving process.

These results are from an Apple silicon GPU:

A new audio length in Kokoro compiles the operators after the break again, and that first call takes about as long as a call without compilation. bench/kokoro_tts.exs rounds the total duration up to steps of 200 ms by making the final pause longer, so sentences of a similar length share one compiled decoder.

EMLX has caches of its own, and each new shape adds to them:

• MLX keeps freed buffers for reuse. Without a limit, this cache grew to 11.7 GB on a 16 GB Mac, and later calls for 20.2 seconds of audio took 22.5 s instead of 0.87 s. Set a limit with EMLX.set_cache_limit(1_000_000_000).
• EMLX.CompileCache keeps each compiled function. Set max_items for it in the :emlx config. EMLX also keeps each MLX program, with no limit.

If you know such a value before the model runs, give it as an Nx.BinaryBackend tensor in the :static_inputs option of Transmuter.compile/3. Then the graph has no break there. Each new value of a static input compiles the model again. Examples:

• use_cache_branch in merged decoder exports selects an If branch.
• num_logits_to_keep in some decoder exports sets the length of a slice.
• input_ids in Grounding DINO go into NonZero.

Transmuter.run/4 without compilation reads all values while the model runs.

Strings, sequences and random values

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

• A string tensor is {:strings, shape, values}, with the values as a flat list of binaries.
• A sequence is a list of tensors.
• An empty optional is :none.
• A tensor with a dimension of 0 is {:empty, shape}, because Nx has no zero-size tensors.

Give string and sequence inputs as :static_inputs. A compiled function returns only tensors, so call Transmuter.run/4 without compilation to read string or sequence outputs.

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. The values cannot be the same as in ONNX Runtime, because the random number generators are different.

ImageDecoder needs the optional stb_image dependency.

Precision and quantization

Half-precision exports (model_fp16.onnx) run in f16 or bf16. Sums, normalizations, softmax and attention scores use f32 inside the operation. ONNX Runtime on the CPU calculates half-precision graphs in f32. To get the same result, use Transmuter.load!(source, compute_type: :f32).

Quantized exports run with their weights dequantized in the graph:

• model_q4.onnx and model_q4f16.onnx use 4-bit weights with MatMulNBits and GatherBlockQuantized.
• model_quantized.onnx and model_int8.onnx use dynamic int8 with DynamicQuantizeLinear, MatMulInteger and ConvInteger.

With backend: EMLX.Backend, Transmuter.load!/2 keeps MatMulNBits weights packed as MLX quantized tensors. Then mx::quantized_matmul reads the 4-bit or 8-bit weights directly. On other backends, each forward pass dequantizes the weights to floats. These results are from an Apple silicon GPU with one forward pass of 8 tokens:

Dynamic int8 quantizes each activation again. A change of one f32 rounding step can move a value to the next int8 step, and the next layers make that change larger. For example, when all weights of SmolLM2 int8 change by a factor of 1 + 1e-7, its logits change by 22. Thus, compare an int8 model with its own noise.

Fine-tuning

Transmuter.run/4 is plain Nx. Thus Nx.Defn.value_and_grad/2 calculates gradients through the graph, and an optimizer such as Polaris updates the params. Load the model with trainable: true. Then small weights, such as a bias with 1 element, are params and not constants.

bench/finetune.exs fine-tunes a MiniLM cross-encoder. Its gradients and its loss curve for 30 AdamW steps are the same as in PyTorch.

Fine-tune an fp32 export. Inference exports have no dropout, and quantized weights have no gradients.

Coverage

bench/coverage runs the most downloaded ONNX exports for each Hugging Face task. It compares each output with ONNX Runtime. The reference is ONNX Runtime with graph optimizations off, because its default optimizer gives wrong results for some exports. For MatMulNBits, the reference uses accuracy_level 0. At level 4, ONNX Runtime quantizes the activations to int8, and that result is an approximation.

Transmuter supports all 203 operators of the ONNX standard domain, and the ONNX Runtime contrib operators that Hugging Face exports use.

Two models do not match value by value, because they make small rounding differences large:

• Grounding DINO selects 900 of 13,294 queries with TopK, and some scores differ by only 9.5e-7. All operators before TopK match. On a real image with a text prompt, Transmuter gives the same detections as ONNX Runtime, with scores within 0.0001.
• Kokoro sums a phase over all audio samples and takes its STFT. ONNX Runtime and MLX round that sum differently. Also, the atan2 of the export gives a different phase for an exact STFT than for the STFT of ONNX Runtime (see Text to speech with Kokoro). When the test gives the Sin and STFT outputs of ONNX Runtime to Transmuter, all other operators match. With the atan2 fix of the example, the spectrum of the speech matches PyTorch, with no tone at 4.8 kHz.

These models also ran from input to final result, and the result agrees with ONNX Runtime:

To run the matrix:

1. Run bench/coverage/select_models.py to select the models.
2. Run bench/coverage/record_references.py to download them and record the ONNX Runtime outputs.
3. Run bench/coverage/run_cases.py to run each case through Transmuter.
4. Run bench/coverage/print_matrix.py to show the matrix, and bench/coverage/list_models.py to make the table of tested models.

If a case does not match, run bench/coverage/trace_ort.py and bench/coverage/find_divergence.exs. They show the first node where Transmuter and ONNX Runtime give different values.

Tests

mix test downloads five models and compares them with ONNX Runtime. test/support/model_references.py shows the reference values.

mix test also runs one small graph for each operator variant in test/fixtures/ops, with no download. test/support/make_op_fixtures.py makes these graphs and records their outputs from ONNX Runtime, or from the ONNX reference evaluator when ONNX Runtime does not have the operator. For random operators, the tests check the shape, the type and the distribution of the values.

On a Mac, TRANSMUTER_TEST_DEVICE=cpu mix test uses the EMLX CPU device. But EMLX 0.5 runs compiled functions on the GPU when Metal is available, also with device: :cpu. Thus only the Linux CI tests the CPU fully.

License

Apache-2.0