Export: training and serving, Julia native
The forward you trained is retraced into the served program, without leaving Julia and Reactant.
using ReactantServerExport # the extension that provides the backendn = Nitro(e; weights = "runs/x/best.jld2", data = (;)) # no training in this processexport_model(n, ReactantServerBundle(); dir = "export_out", name = "mnist_v1")export_model takes a Nitro, which already holds the restored weights, so a model's export code never loads a checkpoint itself. data = (;) skips build_data, so an export process needs no training data.
The wire contract
Export hooks declare what clients send and receive:
export_inputsandexport_outputsdeclare the wire shapes and dtypes.export_preprocessconverts what clients send, usuallyUInt8pixels, into what the model trains on, usually normalizedFloat32. This happens inside the traced graph, so what ships isexport_preprocess ∘ forwardas one program and a client never has to reproduce your normalization.export_postprocessandexport_client_inputs/export_client_outputscover the other side of the seam when the served contract differs from the traced one.
The framework translates axes: an ExportSpec uses 1-based Julia axes, like every other axis in the package, and the backend's IOSpec uses 0-based axes.
What a bundle is
The artifact is a ReactantServer StableHLO bundle:
manifest.yamlweights.safetensors- one
model[.bN].mlirper compiled batch size model.jl, when the model ships a postprocess
Export is a single-device CPU trace, so it needs no GPU, and a handle with a mesh is refused; build the export handle with n_devs = 1. The bundle always records the config, seed, preset, framework version, run directory and checkpoint. Repository state (the commit, the tree hash, and on a dirty tree the working-tree patch) is recorded only when you pass provenance_root, the repository root; without it the export succeeds but cannot say which code produced it.
TensorFlow SavedModel
TFSavedModel writes a SavedModel for TF Serving and the platforms built on it (SageMaker, Vertex AI, KServe). Each traced program is one XlaCallModule op wrapping the StableHLO, the form JAX's native serialization produces, so the program is packaged, not converted.
using PythonCall # the extension; `tensorflow` must be importableexport_model(n, TFSavedModel(); dir = "export_out", name = "mnist_v1", batch_sizes = [1, 8])- One signature per batch size. Every compiled program is static, so
batch_sizes = [1, 8]givesserving_b1andserving_b8, sharing one copy of the weights;serving_defaultis the first. A client picks one withsignature_name. - Batch-first tensors. The Julia shape
(W, H, C, N)is the TF shape[N, C, H, W]. Signature inputs and outputs are named byexport_inputsandexport_outputs. - One platform.
platform = "CUDA"by default; the module runs only there. Use"CPU"for a CPU server. - No postprocess. A SavedModel cannot carry
model.jl. A model declaringexport_client_outputsis refused; any other postprocess is dropped with a warning. - Provenance is written to
assets.extra/provenance.json, with the working-tree patch beside it, which TF Serving ignores.
Verified TensorFlow releases, loading and running every signature on CPU:
| TF | Options |
|---|---|
| 2.17, 2.20 | the defaults, stablehlo_version = v"1.5.0", call_module_version = 9 |
| 2.15 | stablehlo_version = v"0.14.0", call_module_version = 8 |
An older stablehlo_version reaches older releases; a module that uses an op newer than that version then fails at export, not at load. A program containing a custom_call is refused, since XlaCallModule cannot run one portably.
Caveats
A view-shaped program output has to be copyed before it leaves the traced graph, or the bundle carries a wrapper rather than an array. This does not apply to parameters, which the flat layout reconstructs through a copy already.
Each backend is an extension on a weak dependency, ReactantServerExport or PythonCall, so a deployment that never exports does not load either.