Exporting JAX Models for Production

Training a model in JAX is one thing. Getting it into production is another. JAX provides several paths for model export, with StableHLO being the recommended format.

StableHLO Export (Recommended)

StableHLO is now the only recommended export format from JAX. It produces a portable, versioned representation of your computation that can be consumed by various compilers and runtimes. The older jax2tf with enable_xla=False is deprecated.

import jax
from jax import export
import jax.numpy as jnp
import numpy as np

# 1. Create a JIT-transformed function
@jax.jit
def predict(params, x):
    h = jax.nn.relu(x @ params['w1'] + params['b1'])
    return h @ params['w2'] + params['b2']

# 2. Define input shapes for export
params_shapes = {
    'w1': jax.ShapeDtypeStruct((784, 256), jnp.float32),
    'b1': jax.ShapeDtypeStruct((256,), jnp.float32),
    'w2': jax.ShapeDtypeStruct((256, 10), jnp.float32),
    'b2': jax.ShapeDtypeStruct((10,), jnp.float32),
}
x_shape = jax.ShapeDtypeStruct((1, 784), jnp.float32)

# 3. Export to StableHLO
exported = export.export(predict)(params_shapes, x_shape)

# 4. Get the StableHLO module (MLIR text)
stablehlo_module = exported.mlir_module()

# 5. Serialize for later use
serialized = export.export(predict)(params_shapes, x_shape).serialize()
# Can be saved to disk and loaded in a different process/language

Dynamic Batch Sizes

# Export with dynamic batch dimension
scope = export.SymbolicScope()
dynamic_x = jax.ShapeDtypeStruct(
    export.symbolic_shape("batch, 784", scope=scope),
    jnp.float32,
)

exported_dynamic = export.export(predict)(params_shapes, dynamic_x)
# Now the exported model accepts any batch size

TensorFlow SavedModel (via StableHLO)

# You can pack StableHLO into a TensorFlow SavedModel
# for serving with TensorFlow Serving
from jax.experimental.jax2tf import convert as jax2tf_convert

# Note: the recommended path is now:
# 1. Export to StableHLO via jax.export
# 2. Load StableHLO into TF SavedModel if needed for TF Serving
# 3. Or use StableHLO directly with XLA-compatible runtimes