ONNX and MLX / docs / embedding

Export the settling computation and run it without torch: onnxruntime anywhere, MLX on Apple silicon.

Both of these export or reimplement the decision engine, the part that is unique to vega. The perception backbone is an ordinary transformers model with its own mature export paths, so neither module duplicates that work.

ONNX#

bash

pip install "vegaml[onnx]"

python

from vegaml.onnx_export import export_engine, ONNXEngine

export_engine("path/to/checkpoint", "engine.onnx")

eng = ONNXEngine("engine.onnx")              # onnxruntime, CPU provider by default
out = eng(feats, opts, opt_mask, qtype)      # numpy in, numpy out

What exports is the settling computation: pooled perception features in, well energies and the physics readouts out (logits, zT, kinetic, grad_norm, nearest, steps_used). The integrator unrolls into the graph at the checkpoint's own step count; batch size and option count stay dynamic axes, so one file serves any request shape.

A checkpoint with the long-state memory reader exports with chunks and chunk_mask inputs, and an all-false mask reproduces the chunkless path exactly.

Note

Parity is pinned at 1e-4, score rail included.

The test suite holds the export to 1e-4 against the torch engine, and the export's example batch always contains a score question on purpose: the tracer freezes python branches, so a batch without one would silently drop the score rail from the graph.

MLX#

bash

pip install "vegaml[mlx]"        # Apple silicon

python

from vegaml.mlx_engine import MLXEngine

eng = MLXEngine.from_dir("path/to/checkpoint")   # vega_config.json + engine.safetensors
out = eng(feats, opts, opt_mask, qtype)          # numpy in, numpy out

The engine's forward, reproduced in MLX from the same engine.safetensors, read through safetensors' numpy loader: a process that only runs the settling computation needs neither torch nor a CUDA runtime. Scope is the base engine at eval semantics, including the score rail and the long-state memory reader; gated adapters stay with the torch path.

Precision, stated plainly. device="gpu" (the default) runs on Metal, whose matmul accumulates below full float32. Readouts land within about 1e-2 of the torch numbers while decisions and settling steps are identical, and the test suite pins exactly that claim. device="cpu" is strict float32 and matches torch to 1e-4: use it when you are comparing numbers rather than serving them.

When either is worth it#

Only when importing torch is the problem: a sidecar with a size limit, a process that must start in milliseconds, or a host where the right torch wheel is a fight. Otherwise the torch path is the one that is tested end to end and the one every number on this site was measured on.