Most production machine learning stacks keep model definition, training, and experimentation in Python, then wrap the resulting model from another service language for deployment. That split works, but it creates operational seams: model code lives in one runtime, serving code lives in another, and performance depends on how well the two processes exchange tensors, weights, and metadata.
A Go-native ML framework takes a different approach. It exposes model construction, automatic differentiation, training utilities, checkpoint loading, and inference execution directly in Go while delegating low-level acceleration to mature compiler and runtime layers such as OpenXLA, StableHLO, and PJRT. The high-level API is Go code; the heavy numeric kernels still come from the same compiler/runtime ecosystem used by JAX, TensorFlow, and PyTorch/XLA.
GoMLX is a concrete example of this architecture. Its README describes a framework with differentiable operators, layers, training helpers, ONNX conversion paths, a pure Go backend for portability, and an optimized XLA backend for CPUs, GPUs, and TPUs. The transferable pattern is broader than one repository: when a service language owns the ML graph while an accelerator runtime owns compilation and execution, teams can reduce Python sidecars without reimplementing CUDA, TPU, or compiler internals.
The framework frontend is responsible for the user-facing model API: layers, tensor operations, variable scopes, loss functions, optimizers, checkpoint conventions, and training loops. The accelerator backend is responsible for compiling and executing graph operations efficiently on a target device. Keeping those responsibilities separate lets Go code express the model while XLA, PJRT, or a portable Go executor handles numeric execution.
Instead of calling a Python model from a Go service, the service constructs a computation graph in Go. Functions return graph nodes, not immediate numeric results. The graph can then be compiled, differentiated, cached, executed for training, or wrapped as an inference function. This mirrors the static-graph design of many compiler-backed ML systems while preserving Go’s explicit control flow and error handling.
A practical framework needs more than one execution path. An XLA or StableHLO backend provides high-performance execution through PJRT plugins, but it may impose static-shape constraints and native dependency requirements. A pure Go backend is slower, but it is easier to cross-compile, test, run in WebAssembly, or use on platforms without accelerator plugins.
The durable design move is to replace the framework frontend, not the accelerator ecosystem. A Go-native framework should reuse XLA/PJRT or comparable compiler layers for kernel fusion, device buffers, plugin loading, and accelerator dispatch instead of trying to hand-code every low-level primitive.
The model definition is an ordinary Go function that receives graph inputs and returns graph outputs. The function composes layers and primitive operations, threads graph nodes explicitly, and records enough shape information for the backend to compile or execute the computation.
func C10ConvModel(
ctx *mlxcontext.Context,
spec any,
inputs []*graph.Node,
) []*graph.Node {
images := inputs[0]
g := images.Graph()
dtype := images.DType()
batch := images.Shape().Dimensions[0]
logits := images
logits = layers.Convolution(ctx.In("conv1"), logits).
Filters(32).KernelSize(3).PadSame().Done()
logits = activations.Relu(logits)
logits = graph.MaxPool(logits).Window(2).Done()
logits = graph.Reshape(logits, batch, -1)
logits = layers.Dense(ctx.In("head"), logits, true, 10)
_ = g
_ = dtype
return []*graph.Node{logits}
}This style is more verbose than a high-level Python DSL, but the verbosity buys direct integration with Go tooling, static types, ordinary package boundaries, and normal service deployment. The framework can still add higher-level layer builders on top of the graph API.
At execution time, the framework creates a backend, loads variables or checkpoints into a context, traces the model function into a graph, and calls the compiled executable with concrete tensor inputs. The same pattern can be used for batch inference, serving endpoints, notebook experiments, or training loops.
backend := backends.New()
ctx := mlxcontext.New()
if _, err := checkpoints.Load(ctx).Dir(checkpointDir).Done(); err != nil {
return err
}
exec := mlxcontext.NewExec(
backend,
ctx.In("model"),
func(ctx *mlxcontext.Context, image *graph.Node) *graph.Node {
image = graph.ExpandAxes(image, 0)
logits := C10ConvModel(ctx, nil, []*graph.Node{image})[0]
choice := graph.ArgMax(logits, -1, dtypes.Int32)
return graph.Reshape(choice)
},
)
outputs := exec.Call(inputTensor)When the backend is XLA/PJRT, the graph can be lowered into a compiler representation and dispatched to CPU, GPU, or TPU devices. When the backend is pure Go, the same high-level graph can run with fewer deployment dependencies at lower throughput.
GoMLX also points to companion packages for converting ONNX models and downloading model artifacts. This is important because real systems rarely train every model from scratch. A Go-native framework becomes more useful when it can import pretrained weights, convert supported graph operations, and then fine-tune or serve the model in the same runtime.
A service that already runs in Go can keep request parsing, feature extraction, model execution, and response formatting in one process. This avoids a Python sidecar for simple deployments and removes serialization boundaries between business logic and model inference.
Automatic differentiation and optimizer support make the pattern useful beyond inference. Teams can train classifiers, rerankers, calibration layers, embeddings, or domain-specific models in Go while reusing the same packages they use for data loading, validation, metrics, and deployment.
A portable backend enables environments where native accelerator plugins are unavailable. The pure Go path may be slower, but it can run in WebAssembly demos, constrained devices, test binaries, or cross-platform command-line tools. That portability is valuable for education, local experimentation, and fallback execution.
When model construction lives in the same language as the serving system, engineers can version model code, checkpoint schemas, inference wrappers, and operational logic together. This reduces the distance between ML experimentation and production ownership, especially in teams that are already strong in Go.
Backends often register themselves through package imports or build tags. Keep that selection at application startup so model code does not hard-code one device path. This allows local tests to use a portable backend while production binaries include the XLA backend and any required PJRT plugin configuration.
import (
"github.com/gomlx/gomlx/backends"
_ "github.com/gomlx/gomlx/backends/default"
)
func NewModelRuntime() *Runtime {
return &Runtime{
backend: backends.New(),
ctx: mlxcontext.New(),
}
}Keep graph-building functions pure with respect to request handling. A model function should describe tensor transformations. A separate wrapper should load checkpoints, normalize inputs, call the executable, decode outputs, and expose an application-friendly API.
Variable and checkpoint management needs namespacing. A context object can assign stable names to layers, reuse loaded variables, and detect when the model code expects weights that the checkpoint does not contain. This is especially important when evolving a model architecture without silently loading the wrong tensors.
XLA-style compilation works best when shapes are stable. For text, image batches, or variable-length feature sets, group inputs into a small set of shape buckets. That lets compiled executables be reused instead of triggering a new compilation for every unique request shape.