ONNX is commonly treated as an inference interchange format: train a model in one framework, export a frozen graph, and run it through a runtime such as ONNX Runtime. That flow is useful, but it leaves a sharp boundary between imported model weights and the training or graph-building system that might need to adapt the model later.
ONNX-GoMLX demonstrates a different bridge. It parses an ONNX model, transfers its variables into a GoMLX context, rebuilds supported ONNX operations as GoMLX graph nodes, executes the resulting graph through the GoMLX backend, and can copy updated variables back into the ONNX model after training. The imported model is not just a black-box session; it becomes part of a Go-native graph that can be composed, fine-tuned, checkpointed, or exported again.
The transferable technique is graph bridging. Instead of wrapping a foreign runtime behind a narrow inference API, a bridge maps model structure, weights, inputs, and outputs into the host framework’s graph IR. The host framework then owns execution policy, automatic differentiation, backend selection, and training loops while preserving ONNX as the interchange container.
An ONNX file contains a serialized model graph, initializers, named inputs, named outputs, tensor metadata, and operator nodes from one or more opsets. A bridge must parse that representation into an internal model object while preserving enough names and shapes to reconnect caller-provided inputs to graph nodes and exported outputs.
GoMLX separates graph construction from variable storage through a context object. ONNX-GoMLX uses this boundary to copy model initializers into a GoMLX context with stable variable names. Once weights live in the framework context, ordinary training and checkpoint mechanisms can update them without treating the imported ONNX file as immutable.
The central operation is not “call ONNX Runtime.” It is “build equivalent host-framework nodes for the requested ONNX outputs.” The package documentation describes `CallGraph` as building the ONNX graph with GoMLX operations. That means backend compilation, differentiation, batching, and execution behavior follow the host framework’s rules.
The bridge is valuable because it turns a frozen interchange artifact into a framework-owned computation. Supported ONNX operators become graph nodes, ONNX initializers become context variables, and updated context variables can be written back to the ONNX model when the adaptation is complete.
The first step is to read the ONNX file into an internal representation. Large constants and model parameters should become variables rather than inline graph literals, because trainable systems need a place to update and checkpoint them. ONNX-GoMLX exposes `ReadFile` for parsing and `VariablesToContext` for populating the GoMLX context.
modelPath := "/models/all-MiniLM-L6-v2/onnx/model.onnx"
model, err := onnx.ReadFile(modelPath)
if err != nil {
return err
}
ctx := context.New()
if err := model.VariablesToContext(ctx); err != nil {
return err
}This phase should fail early if the model file is invalid, if tensor metadata is unsupported, or if variable names cannot be made safe for the host framework. A production bridge should make those errors explicit before any serving path depends on the imported model.
Execution is expressed as a GoMLX graph-building function. The caller supplies concrete tensors to an executor, but the function itself receives symbolic graph nodes. The bridge maps ONNX input names such as `input_ids`, `attention_mask`, and `token_type_ids` to those nodes, then asks the model to build the requested output graph.
outputs := context.MustExecOnceN(
backends.New(),
ctx,
func(ctx *context.Context, inputs []*Node) []*Node {
return model.CallGraph(ctx, inputs[0].Graph(), map[string]*Node{
"input_ids": inputs[0],
"attention_mask": inputs[1],
"token_type_ids": inputs[2],
}, "last_hidden_state")
},
inputIDs,
attentionMask,
tokenTypeIDs,
)The source README demonstrates this shape with a sentence-transformer model and returns an embedding tensor. The same shape applies to any imported model whose ONNX operators are supported by the bridge: caller inputs become named graph nodes, the ONNX graph is rebuilt in the framework, and concrete tensor outputs come back from the backend executor.
Because the ONNX variables live in the GoMLX context, the imported model can participate in a training loop. The model function can call `CallGraph`, compute a task-specific loss, run automatic differentiation, and update context variables. After training, the bridge can copy the updated context values back into the in-memory ONNX model and save it.
// Inside a training pipeline, after optimizer steps update ctx:
if err := model.ContextToONNX(ctx); err != nil {
return err
}
if err := model.SaveToFile("/models/fine-tuned/model.onnx"); err != nil {
return err
}This closes the loop: ONNX imports seed the host framework, the host framework trains or composes the model, and ONNX remains the deployment or interchange artifact for downstream consumers.
A Go service can load an exported embedding or classification model, tokenize inputs in Go, execute the graph through GoMLX/XLA, and keep application logic in one process. This avoids a Python service boundary while still using an accelerator-backed graph runtime instead of hand-written numeric kernels.
Teams often receive pretrained ONNX models that are close to useful but need domain-specific fine-tuning. A graph bridge lets those weights become trainable variables in the host framework, making it possible to adapt the model without returning to the original training framework.
Once the imported ONNX graph is part of a GoMLX computation, it can be combined with additional GoMLX nodes. That can include normalization, pooling, projection heads, ensemble logic, calibration layers, or task-specific output transforms that would be awkward to bolt onto an opaque inference session.
Writing updated variables back to ONNX preserves interoperability. A team can fine-tune in a Go-native graph environment, save the adapted ONNX artifact, and then deploy it in another runtime if that runtime is the operational standard for a different service.
ONNX import is only as complete as the supported operators and shape rules. Keep a conversion test suite for every model family you intend to support. The test should parse the model, load variables, build outputs, run representative inputs, and compare output shape and sample values against a known-good runtime.
func ValidateImportedModel(path string, inputs Fixture) error {
m, err := onnx.ReadFile(path)
if err != nil {
return err
}
ctx := context.New()
if err := m.VariablesToContext(ctx); err != nil {
return err
}
got := runImportedGraph(m, ctx, inputs)
return compareAgainstGolden(got, inputs.Expected)
}Do not rely on positional conventions unless the model format forces you to. Named maps make it clear which caller tensor feeds which ONNX input, and explicit output names let you select intermediate or final values intentionally. This is especially important for transformer exports with multiple masks, token type inputs, logits, pooled outputs, and hidden states.
The bridge should own parsing, variable transfer, name mapping, and graph construction. Task code should own tokenization, batching, loss computation, optimizer choice, and serving behavior. That separation keeps model import mechanics testable and prevents business logic from depending on ONNX protobuf details.
When saving back to ONNX, treat variable update as a controlled export operation. Record which context variables were trainable, which were frozen, and which were added by the host framework. If the adapted model includes new layers that were not present in the original ONNX graph, decide whether to export a combined graph or keep a separate framework checkpoint.