Chapter 9: Export and Deployment

import os, re, sys, json, time, shutil, textwrap, subprocess, tempfile, importlib.util
from importlib.metadata import requires, version
from pathlib import Path

import numpy as np
import jax
import jax.numpy as jnp
from jax import export
import jaxlib, flax, optax
from flax import nnx
import orbax.checkpoint as ocp
import absl.logging
absl.logging.set_verbosity(absl.logging.ERROR)        # keep only Orbax error logs so INFO/WARNING don't flood the output

print("Python", sys.version.split()[0], "| JAX", jax.__version__, "| jaxlib", jaxlib.__version__,
      "| Flax", flax.__version__, "| Optax", optax.__version__, "| Orbax", version("orbax-checkpoint"))
print("Devices:", jax.devices())

# Export artifact directory: next to this notebook (export_artifacts/ in the notebook's folder), cleared on every run
ART = Path("export_artifacts").resolve()
if ART.exists():
    shutil.rmtree(ART)
ART.mkdir(parents=True)
SCRATCH = Path(tempfile.mkdtemp(prefix="ext03_"))     # intermediate files go in a temp dir, not the repo
print("Artifact directory:", ART.name + "/")

def run_py(src, *args, env_extra=None, cwd=None):
    """Run script src in a brand-new subprocess (python -I isolated mode) and return stdout; on failure raise an exception that includes stderr."""
    env = {**os.environ, **(env_extra or {})}
    r = subprocess.run([sys.executable, "-I", "-c", textwrap.dedent(src), *map(str, args)],
                       capture_output=True, text=True, env=env, cwd=str(cwd or SCRATCH))
    if r.returncode != 0:
        raise RuntimeError(f"Subprocess failed (return code {r.returncode}):\n{r.stderr[-3000:]}")
    return r.stdout

def sds(tree):
    """PyTree of arrays → ShapeDtypeStruct PyTree with the same structure (shape and dtype only)."""
    return jax.tree.map(lambda a: jax.ShapeDtypeStruct(np.shape(a), a.dtype), tree)

def max_diff(a, b):
    """Largest elementwise absolute difference between two PyTrees of the same structure (compared in float64)."""
    la, lb = jax.tree.leaves(a), jax.tree.leaves(b)
    assert len(la) == len(lb)
    return max(float(np.max(np.abs(np.asarray(x, np.float64) - np.asarray(y, np.float64)))) for x, y in zip(la, lb))

def tree_bytes(tree):
    return sum(np.asarray(x).nbytes for x in jax.tree.leaves(tree))

def dir_bytes(p):
    p = Path(p)
    return p.stat().st_size if p.is_file() else sum(f.stat().st_size for f in p.rglob("*") if f.is_file())

def main_signature(exp):
    """Pull out the signature line of the exported module's @main (to see argument and result types)."""
    return next(l.strip() for l in exp.mlir_module().splitlines() if "func.func public @main" in l)

# ❌ Wrong: calling serialize() without flatbuffers installed; jax won't install it for you
# We don't actually uninstall the package: in the subprocess we set sys.modules["flatbuffers"] = None, so any later import flatbuffers raises ImportError, which simulates "not installed"
try:
    out = run_py("""
        import sys
        sys.modules["flatbuffers"] = None
        import jax, jax.numpy as jnp
        from jax import export
        exp = export.export(jax.jit(jnp.sin))(jnp.ones(3))
        print("Export itself succeeded, platforms =", exp.platforms)
        try:
            exp.serialize()
            print("RESULT serialization succeeded")
        except ImportError as e:
            print("RESULT ImportError:", e)
    """)
    print("Measured:\n" + out)
except Exception as e:
    print("Measured error:", type(e).__name__, e)
assert "Export itself succeeded" in out
assert "RESULT ImportError: Please install 'flatbuffers' in order to use Exported serialization" in out

# ✅ Correct: install flatbuffers separately (pip install flatbuffers) and add it to requirements.txt
import flatbuffers
deps = (requires("jax") or []) + (requires("jaxlib") or [])
print("flatbuffers version:", flatbuffers.__version__)
print("Do jax/jaxlib declare flatbuffers as a dependency?", any(d.lower().startswith("flatbuffers") for d in deps))
assert not any(d.lower().startswith("flatbuffers") for d in deps)      # it really isn't a jax dependency

blob = export.export(jax.jit(jnp.sin))(jnp.ones(3)).serialize()
print("serialize() works:", type(blob).__name__, len(blob), "bytes")
assert isinstance(blob, bytearray) and len(blob) > 0

def linear_apply(params, x):
    """The smallest "model": the program only describes the computation; weights come in as arguments."""
    return x @ params["w"] + params["b"]

k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
params_a = {"w": jax.random.normal(k1, (4, 2)), "b": jnp.zeros(2)}          # first set of weights
params_b = {"w": jax.random.normal(k2, (4, 2)), "b": jnp.ones(2)}           # second set of weights (say, after retraining)
x_demo = jax.random.normal(k3, (3, 4))

# Program: export once (uses only shapes and dtypes)
prog = export.export(jax.jit(linear_apply))(sds(params_a), sds(x_demo))
# Data: save the two sets of weights with Orbax
ckptr = ocp.StandardCheckpointer()
for name, p in [("a", params_a), ("b", params_b)]:
    ckptr.save(SCRATCH / f"linear_{name}", p)
ckptr.wait_until_finished()

for name, p in [("a", params_a), ("b", params_b)]:
    restored = ckptr.restore(SCRATCH / f"linear_{name}", sds(p))           # restore from an abstract template
    y = prog.call(restored, x_demo)
    print(f"Weights {name}: output of the exported program = {np.round(np.asarray(y[0]), 4)}")
    assert max_diff(y, linear_apply(p, x_demo)) < 1e-6
assert max_diff(prog.call(params_a, x_demo), prog.call(params_b, x_demo)) > 0.1
print("One exported program + two sets of Orbax weights → two different models, both matching the original function ✓")

def mlp_apply(params, x):
    """Pure-function version of a two-layer MLP. The weights params are an argument, not closed-over constants."""
    h = jnp.tanh(x @ params["w1"] + params["b1"])
    return h @ params["w2"] + params["b2"]

def init_mlp(key, din=4, dh=8, dout=3):
    k1, k2 = jax.random.split(key)
    return {"w1": jax.random.normal(k1, (din, dh)) / np.sqrt(din), "b1": jnp.zeros(dh),
            "w2": jax.random.normal(k2, (dh, dout)) / np.sqrt(dh), "b2": jnp.zeros(dout)}

params = init_mlp(jax.random.key(0))
x2 = jax.random.normal(jax.random.key(1), (2, 4))

exp = export.export(jax.jit(mlp_apply))(params, x2)
print("fun_name     :", exp.fun_name)
print("in_tree      :", exp.in_tree)
print("in_avals     :", exp.in_avals)
print("out_avals    :", exp.out_avals)
print("platforms    :", exp.platforms, "| nr_devices:", exp.nr_devices, "| calling_convention_version:", exp.calling_convention_version)
print("@main signature:", main_signature(exp)[:200], "...")

assert exp.fun_name == "mlp_apply" and exp.platforms == ("cpu",) and exp.nr_devices == 1
assert [a.shape for a in exp.in_avals] == [(8,), (3,), (4, 8), (8, 3), (2, 4)]       # b1, b2, w1, w2, x
assert exp.out_avals[0].shape == (2, 3) and exp.in_tree.num_leaves == 5
assert "stablehlo.dot_general" in exp.mlir_module() and "stablehlo.tanh" in exp.mlir_module()
assert max_diff(exp.call(params, x2), mlp_apply(params, x2)) < 1e-6

# Exporting with only ShapeDtypeStructs works too and produces exactly the same program
exp_abs = export.export(jax.jit(mlp_apply))(sds(params), jax.ShapeDtypeStruct((2, 4), jnp.float32))
assert exp_abs.mlir_module_serialized == exp.mlir_module_serialized
print("Exporting with only ShapeDtypeStructs: the bytecode is identical to exporting with real arrays ✓")

# ❌ Wrong: exporting a plain Python function directly (not wrapped in jax.jit)
try:
    export.export(mlp_apply)(params, x2)
    print("Measured: it runs")
except ValueError as e:
    print("Measured error: ValueError:", e)
    err_nojit = str(e)
assert "must be the result of `jit`" in err_nojit

# ✅ Correct: export.export only accepts the result of jax.jit; jit options such as static_argnums and donate take effect too
exp_ok = export.export(jax.jit(mlp_apply))(params, x2)
assert max_diff(exp_ok.call(params, x2), mlp_apply(params, x2)) < 1e-6
print("export.export(jax.jit(f)) ✓")

blob = exp.serialize()
exp_path = SCRATCH / "mlp_static.jaxexport"
exp_path.write_bytes(blob)
io_path = SCRATCH / "mlp_static_io.npz"
np.savez(io_path, x=np.asarray(x2), **{k: np.asarray(v) for k, v in params.items()})
print(f"Export file {exp_path.name}: {len(blob)} bytes; header {bytes(blob[:8])!r}")

sub_src = """
    import sys, numpy as np, jax
    from jax import export
    exp = export.deserialize(open(sys.argv[1], "rb").read())     # only reads bytes; doesn't import any model code
    d = np.load(sys.argv[2])
    params = {k: d[k] for k in ("w1", "b1", "w2", "b2")}
    np.save(sys.argv[3], np.asarray(exp.call(params, d["x"])))
    print("subprocess fun_name =", exp.fun_name, "| in_avals =", exp.in_avals)
    print("subprocess sys.path[0] =", sys.path[0])
"""
out_path = SCRATCH / "mlp_static_out.npy"
print(run_py(sub_src, exp_path, io_path, out_path))

y_sub = np.load(out_path)
print("Max difference between subprocess output and the original function:", max_diff(y_sub, mlp_apply(params, x2)))
assert max_diff(y_sub, mlp_apply(params, x2)) < 1e-6
for token in ("mlp_apply", "tanh", "@", "def "):                 # the subprocess script contains no computation logic
    assert token not in sub_src, token

# In the same process, the deserialized object has the same signature as the original
exp_back = export.deserialize(blob)
assert exp_back.in_avals == exp.in_avals and exp_back.out_avals == exp.out_avals and exp_back.in_tree == exp.in_tree
print("serialize → file → deserialize in a new subprocess → call: the result matches the original function ✓")

# ❌ Wrong: baking the weights into the exported program as closed-over constants (program and data welded together, a bloated file, and it triggers the CPU backend bug)
rng = np.random.default_rng(0)
n = 128
W_np = rng.standard_normal((n, n)).astype(np.float32)
X_np = rng.standard_normal((n, n)).astype(np.float32)
ref = float(np.sum(W_np.astype(np.float64) @ X_np.astype(np.float64)))      # float64 reference value
W_const = jnp.asarray(W_np)
try:
    exp_closure = export.export(jax.jit(lambda x: jnp.sum(W_const @ x)))(jax.ShapeDtypeStruct((n, n), jnp.float32))
    r_closure = float(exp_closure.call(jnp.asarray(X_np)))
    size_closure = len(exp_closure.serialize())
    print(f"Closure version: export file {size_closure:,} bytes; result {r_closure:.4f}, reference {ref:.4f}, error {abs(r_closure - ref):.2e}")
    print("Did the closure version trigger the CPU bug:", abs(r_closure - ref) > 1.0)
except Exception as e:
    print("Measured error:", type(e).__name__, e)
# The weights really are embedded in the StableHLO as a 128×128 constant
assert re.search(r"stablehlo\.constant dense<.*> : tensor<128x128xf32>", exp_closure.mlir_module())

# ✅ Correct: pass the weights in as arguments; the exported program contains no weights at all
exp_param = export.export(jax.jit(lambda w, x: jnp.sum(w @ x)))(
    jax.ShapeDtypeStruct((n, n), jnp.float32), jax.ShapeDtypeStruct((n, n), jnp.float32))
r_param = float(exp_param.call(jnp.asarray(W_np), jnp.asarray(X_np)))
size_param = len(exp_param.serialize())
print(f"Argument version: export file {size_param:,} bytes; result {r_param:.4f}, error {abs(r_param - ref):.2e}")
print(f"The closure version's file is {size_closure / size_param:.0f}× the argument version's (the {W_np.nbytes:,} bytes of weights were baked in)")

assert abs(r_param - ref) < 1e-2                                  # the argument version is correct
assert size_closure > W_np.nbytes > 10 * size_param               # the closure version's size is mostly that copy of the weights
assert not re.search(r"tensor<128x128xf32> = stablehlo\.constant|constant dense<[^>]{100,}", exp_param.mlir_module())

x_poly = jax.ShapeDtypeStruct(export.symbolic_shape("b, 4"), jnp.float32)
print("Symbolic shape:", x_poly.shape, "| type:", type(x_poly.shape[0]).__name__, "| is_symbolic_dim:", export.is_symbolic_dim(x_poly.shape[0]))

exp_poly = export.export(jax.jit(mlp_apply))(sds(params), x_poly)
print("in_avals :", exp_poly.in_avals)
print("out_avals:", exp_poly.out_avals)
print("@main signature:", main_signature(exp_poly)[-120:])
assert str(exp_poly.out_avals[0].shape) == "(b, 3)"
assert "tensor<?x4xf32>" in main_signature(exp_poly) and "shape_assertion" in exp_poly.mlir_module()

exp_poly_back = export.deserialize(exp_poly.serialize())           # symbolic shapes can be serialized too
for batch in [1, 5, 17, 256]:
    xb = jax.random.normal(jax.random.key(batch), (batch, 4))
    yb = exp_poly_back.call(params, xb)
    assert yb.shape == (batch, 3) and max_diff(yb, mlp_apply(params, xb)) < 1e-5
print("One exported artifact; batch = 1 / 5 / 17 / 256 all match the original function ✓")

# ❌ Wrong: calling with a shape / dtype that doesn't match the export signature
errs_call = {}
for name, bad_x in [("feature dim 5 ≠ 4", jnp.ones((3, 5))), ("dtype int32", jnp.ones((3, 4), jnp.int32))]:
    try:
        exp_poly.call(params, bad_x)
        errs_call[name] = "it runs"
    except ValueError as e:
        errs_call[name] = f"ValueError: {e}"
    print(f"{name}: {errs_call[name]}")
assert "Shape mismatch for args[1].shape[1]" in errs_call["feature dim 5 ≠ 4"]
assert "Dtype mismatch for args[1]: expected float32 and called with int32" in errs_call["dtype int32"]

# ✅ Correct: outside the export boundary, convert the input to the shape and dtype the export signature requires
raw = np.arange(12, dtype=np.int64).reshape(3, 4)            # data from outside may well be int64
y_ok = exp_poly.call(params, jnp.asarray(raw, jnp.float32))
assert y_ok.shape == (3, 3)
print("Convert the input to float32, (b, 4) first, then call ✓")

# ❌ Wrong: reshape to (-1, 3); 4b isn't necessarily divisible by 3
try:
    export.export(jax.jit(lambda x: x.reshape(-1, 3)))(x_poly)
    print("Measured: it exports")
except Exception as e:
    err_div = f"{type(e).__name__}: {e}"
    print("Measured error:", err_div[:220])
assert err_div.startswith("InconclusiveDimensionOperation") and "Cannot divide evenly" in err_div

# ✅ Correct: only do reshapes that work for every b; the output shape contains the symbolic expression 2*b
exp_r2 = export.export(jax.jit(lambda x: x.reshape(-1, 2)))(x_poly)
print("Output shape of reshape(-1, 2):", exp_r2.out_avals[0].shape)
assert str(exp_r2.out_avals[0].shape[0]) == "2*b"
assert exp_r2.call(jnp.ones((5, 4))).shape == (10, 2)

# ❌ Wrong: a Python conditional on a symbolic dimension (b > 2 holds for some b and not for others)
def scale_if_large(x):
    if x.shape[0] > 2:            # tracing needs a definite Python bool here
        return x * 2.0
    return x

try:
    export.export(jax.jit(scale_if_large))(x_poly)
    print("Measured: it exports")
except Exception as e:
    err_cmp = f"{type(e).__name__}: {e}"
    print("Measured error:", err_cmp.split("\n")[0])
assert err_cmp.startswith("InconclusiveDimensionOperation") and "'b' > '2' is inconclusive" in err_cmp

# ✅ Correct: declare b >= 3 with constraints, and the comparison has a definite answer; the constraint is checked at call time
x_poly_ge3 = jax.ShapeDtypeStruct(export.symbolic_shape("b, 4", constraints=("b >= 3",)), jnp.float32)
exp_ge3 = export.export(jax.jit(scale_if_large))(x_poly_ge3)
assert max_diff(exp_ge3.call(jnp.ones((5, 4))), 2.0 * jnp.ones((5, 4))) == 0.0
try:
    exp_ge3.call(jnp.ones((2, 4)))                       # b = 2 violates the constraint
except ValueError as e:
    err_ge3 = str(e)
    print("Error when calling with b = 2:", err_ge3[:110], "...")
assert "do not match the symbolic shape constraint b >= 3" in err_ge3
print("With the constraint declared it exports; calls that violate it are stopped by shape_assertion ✓")

# ❌ Wrong: floordiv(b, 2) silently assumes "b is even"; it exports, but calling with an odd b fails during shape refinement
half_reshape = lambda x: x.reshape(x.shape[0] // 2, 2, 4)
exp_floor = export.export(jax.jit(half_reshape))(x_poly)
print("Export succeeded, output shape:", exp_floor.out_avals[0].shape)
print("Call with b = 4:", exp_floor.call(jnp.ones((4, 4))).shape)
try:
    exp_floor.call(jnp.ones((3, 4)))
    print("b = 3: it runs")
except ValueError as e:
    err_refine = str(e)
    print("Error when calling with b = 3: ValueError:", err_refine)
assert str(exp_floor.out_avals[0].shape[0]) == "floordiv(b, 2)"
assert "Error refining shapes" in err_refine

# ✅ Correct: put "even" into the shape spec "2*k, 4"; b // 2 simplifies to k, and odd inputs are clearly rejected at call time
x_even = jax.ShapeDtypeStruct(export.symbolic_shape("2*k, 4"), jnp.float32)
exp_even = export.export(jax.jit(half_reshape))(x_even)
print("Input spec:", exp_even.in_avals[0].shape, "→ output shape:", exp_even.out_avals[0].shape)
assert exp_even.call(jnp.ones((6, 4))).shape == (3, 2, 4)
try:
    exp_even.call(jnp.ones((5, 4)))
except ValueError as e:
    err_even = str(e)
    print("Error when calling with b = 5:", err_even[:150], "...")
assert str(exp_even.out_avals[0].shape[0]) == "k"
assert "Division had remainder 1 when computing the value of 'k'" in err_even

mask_fn = lambda x, m: x * m[:, None]

# ❌ Wrong: two independent symbolic_shape calls, so the two "b"s belong to different scopes
try:
    spec_x = jax.ShapeDtypeStruct(export.symbolic_shape("b, 4"), jnp.float32)
    spec_m = jax.ShapeDtypeStruct(export.symbolic_shape("b"), jnp.float32)
    export.export(jax.jit(mask_fn))(spec_x, spec_m)
    print("Measured: it exports")
except ValueError as e:
    err_scope = str(e)
    print("Measured error: ValueError:", err_scope.split("\n")[0])
assert "Invalid mixing of symbolic scopes" in err_scope

# ✅ Correct: symbolic_args_specs generates the specs for all arguments in one scope ("_" keeps the example argument's dimension)
specs = export.symbolic_args_specs((jnp.ones((2, 4)), jnp.ones((2,))), ("b, _", "b"))
print("Generated specs:", specs)
exp_mask = export.export(jax.jit(mask_fn))(*specs)
xm, mm = jnp.arange(12.).reshape(3, 4), jnp.array([1., 0., 2.])
assert max_diff(exp_mask.call(xm, mm), mask_fn(xm, mm)) == 0.0
try:
    exp_mask.call(jnp.ones((3, 4)), jnp.ones(5))           # the two arguments disagree on b
except ValueError as e:
    print("Error when b is inconsistent:", str(e)[:120], "...")
    assert "Found inconsistency between dimension size args[1].shape[0] (= 5)" in str(e)

# Check the first two rules above
try:
    a_, b_ = export.symbolic_shape("a, b")
    export.export(jax.jit(lambda x: x))(jax.ShapeDtypeStruct((a_ + b_, 4), jnp.float32))
except ValueError as e:
    print("Spec a+b:", str(e)[:105], "...")
    assert "Cannot solve for values of dimension variables" in str(e)

b_sym, = export.symbolic_shape("b")
print("b == 1 →", b_sym == 1, "| b >= 1 →", b_sym >= 1)
assert (b_sym == 1) is False and (b_sym >= 1) is True

# Lowering only; not run on GPU/TPU
exp_cpu = export.export(jax.jit(mlp_apply))(sds(params), x_poly)
exp_mp = export.export(jax.jit(mlp_apply), platforms=("cpu", "cuda", "tpu"))(sds(params), x_poly)
print("platforms:", exp_mp.platforms)
print("Start of the multi-platform @main signature:", main_signature(exp_mp)[:110], "...")
size_cpu, size_mp = len(exp_cpu.serialize()), len(exp_mp.serialize())
print(f"Serialized size: cpu only {size_cpu:,} bytes; cpu+cuda+tpu {size_mp:,} bytes ({size_mp / size_cpu - 1:.0%} more)")

assert exp_mp.platforms == ("cpu", "cuda", "tpu")
assert "platform_index" in main_signature(exp_mp) and "platform_index" not in main_signature(exp_cpu)
assert size_mp < 2 * size_cpu                        # shared parts are stored only once
# Calling on this machine (CPU): the cpu branch is selected automatically and the result matches the original function
xb = jax.random.normal(jax.random.key(7), (6, 4))
assert max_diff(export.deserialize(exp_mp.serialize()).call(params, xb), mlp_apply(params, xb)) < 1e-5
print("The multi-platform artifact called on CPU (cpu branch selected automatically) gives the right result ✓")

# ❌ Wrong: exporting only for cuda, then calling it on CPU
exp_cuda = export.export(jax.jit(mlp_apply), platforms=("cuda",))(sds(params), x_poly)   # lowering only
try:
    exp_cuda.call(params, xb)
    print("Measured: it runs")
except ValueError as e:
    err_plat = str(e)
    print("Measured error: ValueError:", err_plat)
assert "was exported for platforms '('cuda',)' but it is used on '('cpu',)'" in err_plat

# ✅ Correct: include every platform it will run on in platforms (including cpu for local development / testing)
exp_cuda_cpu = export.export(jax.jit(mlp_apply), platforms=("cuda", "cpu"))(sds(params), x_poly)
assert max_diff(exp_cuda_cpu.call(params, xb), mlp_apply(params, xb)) < 1e-5
print("platforms=('cuda', 'cpu'): callable directly on CPU; the cuda branch was only lowered ✓")

import inspect
print("Exported.serialize signature:", inspect.signature(export.Exported.serialize))
assert "vjp_order" in str(inspect.signature(export.Exported.serialize))

def mse_loss(params, x, y):
    return jnp.mean((mlp_apply(params, x) - y) ** 2)

xg = jax.random.normal(jax.random.key(10), (16, 4))
yg = jax.random.normal(jax.random.key(11), (16, 3))
exp_loss = export.export(jax.jit(mse_loss))(sds(params), sds(xg), sds(yg))
print("has_vjp() in the same process:", exp_loss.has_vjp())

g_ref = jax.grad(mse_loss)(params, xg, yg)
g_exp = jax.grad(exp_loss.call)(params, xg, yg)            # differentiate the exported function; the gradient is still a dict
print("Gradient structure:", jax.tree.map(jnp.shape, g_exp))
assert exp_loss.has_vjp() and max_diff(g_exp, g_ref) < 1e-6
vjp_exp = exp_loss.vjp()                                    # the backward function is itself an Exported
print("vjp() takes", len(vjp_exp.in_avals), "inputs (6 leaves of the original inputs + 1 output cotangent) and returns", len(vjp_exp.out_avals), "outputs")
assert len(vjp_exp.in_avals) == 7 and len(vjp_exp.out_avals) == 6

# ❌ Wrong: save with the default serialize() (vjp_order=0), then differentiate after reading it back
loaded0 = export.deserialize(exp_loss.serialize())
try:
    jax.grad(loaded0.call)(params, xg, yg)
    print("Measured: it can be differentiated")
except ValueError as e:
    err_vjp0 = str(e)
    print("Measured error: ValueError:", err_vjp0, "| has_vjp() =", loaded0.has_vjp())
assert err_vjp0 == "No VJP is available" and not loaded0.has_vjp()

# ✅ Correct: declare the order you need, vjp_order=k, on the exporting side; reverse-mode derivatives up to order k are available
sizes = {k: len(exp_loss.serialize(vjp_order=k)) for k in (0, 1, 2)}
print("Serialized size (bytes):", sizes)
assert sizes[0] < sizes[1] < sizes[2]

loaded1 = export.deserialize(exp_loss.serialize(vjp_order=1))
assert max_diff(jax.grad(loaded1.call)(params, xg, yg), g_ref) < 1e-6           # first order OK

def grad_norm2(load):                                       # second order: differentiate the "squared gradient norm" again
    return lambda p: sum(jnp.sum(g ** 2) for g in jax.tree.leaves(jax.grad(load)(p, xg, yg)))
try:
    jax.grad(grad_norm2(loaded1.call))(params)
except ValueError as e:
    print("Second derivative with vjp_order=1: ValueError:", e)
    assert str(e) == "No VJP is available"

loaded2 = export.deserialize(exp_loss.serialize(vjp_order=2))
gg_ref = jax.grad(grad_norm2(mse_loss))(params)
gg_exp = jax.grad(grad_norm2(loaded2.call))(params)
assert max_diff(gg_exp, gg_ref) < 1e-5
print("vjp_order=1: first derivative correct, second derivative fails; vjp_order=2: second derivative matches the original function ✓")

# ❌ Wrong: vmap / jvp over Exported.call
x_single = jax.random.normal(jax.random.key(12), (2, 4))
x_stack = jax.random.normal(jax.random.key(13), (5, 2, 4))      # 5 groups, each (2, 4)
errs_tf = {}
for name, fn in [("vmap", lambda: jax.vmap(exp.call, in_axes=(None, 0))(params, x_stack)),
                 ("jvp", lambda: jax.jvp(lambda x: exp.call(params, x), (x_single,), (x_single,)))]:
    try:
        fn()
        errs_tf[name] = "it runs"
    except Exception as e:
        errs_tf[name] = f"{type(e).__name__}: {e}"
    print(f"{name}: {errs_tf[name]}")
assert errs_tf["vmap"].startswith("NotImplementedError: Batching rule for 'call_exported' not implemented")
assert errs_tf["jvp"].startswith("TypeError") and "custom_vjp" in errs_tf["jvp"]

# ✅ Correct: build the batch dimension into the program before exporting, by exporting the vmapped function (plus a symbolic dimension n here)
batched = jax.vmap(mlp_apply, in_axes=(None, 0))
x_stack_spec = jax.ShapeDtypeStruct(export.symbolic_shape("n, 2, 4"), jnp.float32)
exp_batched = export.export(jax.jit(batched))(sds(params), x_stack_spec)
assert max_diff(exp_batched.call(params, x_stack), batched(params, x_stack)) < 1e-5

# jit can nest: the exported function can be part of a larger jit function (pre/post-processing written in jnp)
@jax.jit
def predict_labels(params, x):
    return jnp.argmax(exp_poly.call(params, x), axis=-1)
assert bool(jnp.all(predict_labels(params, xb) == jnp.argmax(mlp_apply(params, xb), -1)))
print("Export vmap(f) instead of vmap(exp.call) ✓; Exported.call can be nested inside an outer jit ✓")
jaxpr_txt = str(jax.make_jaxpr(predict_labels)(params, xb))
print("The outer jaxpr contains the call_exported primitive:", "call_exported" in jaxpr_txt)
assert "call_exported" in jaxpr_txt

import jax._src.export.serialization as _ser            # private module; we only read one constant here to look at it, don't depend on it in your code
from jax.extend.mlir.dialects import stablehlo as shlo

print("Serialization format version _SERIALIZATION_VERSION =", _ser._SERIALIZATION_VERSION)
print("Calling convention: default", jax.config.jax_export_calling_convention_version,
      "| supported range", export.minimum_supported_calling_convention_version, "–", export.maximum_supported_calling_convention_version)
target = shlo.get_version_from_compatibility_requirement(shlo.StablehloCompatibilityRequirement.WEEK_4)
print("StableHLO: current version", shlo.get_current_version(), "| minimum readable version", shlo.get_minimum_version(), "| serialization target version (≥4 weeks old)", target)
bc = exp.mlir_module_serialized
print("Portable bytecode header:", bc[:4], "(MLIR bytecode magic number 'ML\\xefR')")

assert isinstance(_ser._SERIALIZATION_VERSION, int) and _ser._SERIALIZATION_VERSION >= 12
assert export.minimum_supported_calling_convention_version <= exp.calling_convention_version <= export.maximum_supported_calling_convention_version
assert exp.calling_convention_version == jax.config.jax_export_calling_convention_version == 10
ver = lambda s: tuple(int(t) for t in s.split("."))
assert ver(shlo.get_minimum_version()) <= ver(target) <= ver(shlo.get_current_version())
assert bc[:4] == b"ML\xefR"

cc9_path = SCRATCH / "mlp_cc9.jaxexport"
out = run_py("""
    import sys, jax, jax.numpy as jnp
    from jax import export
    w = jax.ShapeDtypeStruct((4, 3), jnp.float32)
    x = jax.ShapeDtypeStruct(export.symbolic_shape("b, 4"), jnp.float32)
    e = export.export(jax.jit(lambda w, x: jnp.tanh(x @ w)))(w, x)
    open(sys.argv[1], "wb").write(e.serialize())
    print("CC", e.calling_convention_version)
""", cc9_path, env_extra={"JAX_EXPORT_CALLING_CONVENTION_VERSION": "9"})
print("Subprocess output:", out.strip())
e9 = export.deserialize(cc9_path.read_bytes())
w9 = jax.random.normal(jax.random.key(3), (4, 3))
x9 = jax.random.normal(jax.random.key(4), (5, 4))
print("Read back in the current process: calling_convention_version =", e9.calling_convention_version)
assert out.strip() == "CC 9" and e9.calling_convention_version == 9
assert max_diff(e9.call(w9, x9), jnp.tanh(x9 @ w9)) < 1e-6
print("A version-9 artifact runs fine in an environment that defaults to version 10 ✓")

# Measured: a Pallas TPU kernel needs no DisabledSafetyCheck, a custom FFI target does (lowering only, not run on TPU)
from jax.experimental import pallas as pl
from jax._src.export._export import _CUSTOM_CALL_TARGETS_GUARANTEED_STABLE as _STABLE   # private constant; we only read it to look at it

def _add_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]
s8 = jax.ShapeDtypeStruct((8, 128), jnp.float32)
exp_pl = export.export(jax.jit(pl.pallas_call(_add_kernel, out_shape=s8)), platforms=("tpu",))(s8, s8)
blob_pl = exp_pl.serialize()                                     # no disabled_checks passed
assert "tpu_custom_call" in exp_pl.mlir_module() and len(blob_pl) > 0
print("Pallas TPU kernel: export + serialization succeeded (lowering only)")

ffi_fn = lambda x: jax.ffi.ffi_call("my_custom_target", jax.ShapeDtypeStruct(x.shape, x.dtype))(x)
try:
    export.export(jax.jit(ffi_fn))(jnp.ones(3))
    err_ffi = None
except ValueError as e:
    err_ffi = str(e)
    print("Custom FFI target: ValueError:", err_ffi.splitlines()[0][:95], "...")
assert err_ffi is not None and "no compatibility guarantees" in err_ffi
exp_ffi = export.export(jax.jit(ffi_fn),
                        disabled_checks=[export.DisabledSafetyCheck.custom_call("my_custom_target")])(jnp.ones(3))
assert len(exp_ffi.serialize()) > 0
print("With an explicit DisabledSafetyCheck.custom_call it exports (the consumer must register a target with the same name to run it)")
print("Stable list: tpu_custom_call", "tpu_custom_call" in _STABLE, "| mosaic_gpu_v2", "mosaic_gpu_v2" in _STABLE,
      "| __gpu$xla.gpu.triton", "__gpu$xla.gpu.triton" in _STABLE)
assert {"tpu_custom_call", "mosaic_gpu_v2"} <= _STABLE and "__gpu$xla.gpu.triton" not in _STABLE

shard_src = """
    import os, sys
    import numpy as np, jax, jax.numpy as jnp
    from jax import export, lax
    from jax.sharding import PartitionSpec as P
    assert jax.device_count() == 4
    def check(name, ok, info=""):
        print(("PASS " if ok else "FAIL ") + name, info)

    f = lambda w, x: jnp.tanh(x @ w)
    w = jnp.full((4, 3), 0.1); x = jnp.arange(32.).reshape(8, 4)
    mesh = jax.make_mesh((4,), ("data",))
    with jax.set_mesh(mesh):
        xs, ws = jax.device_put(x, P("data")), jax.device_put(w, P())
        e = export.export(jax.jit(f))(ws, xs)                    # input split along the batch dim across 4 devices
        check("1 nr_devices==4", e.nr_devices == 4, e.nr_devices)
        check("1 input sharding recorded", str(e.in_shardings_jax(mesh)[1].spec) == str(P("data")), e.in_shardings_jax(mesh)[1].spec)
        y = e.call(ws, xs)
        check("1 call in the same mesh is correct", np.allclose(y, f(w, x), atol=1e-6), y.sharding.spec)
    blob = e.serialize()
    e2 = export.deserialize(blob)
    try:                                                         # call outside any mesh
        e2.call(w, x); check("1 call without a mesh should fail", False)
    except ValueError as err:
        check("1 call without a mesh → ValueError", "is not equal to the size of the mesh 4" in str(err), str(err)[:90])
    mesh2 = jax.make_mesh((2,), ("data",), devices=jax.devices()[:2])
    try:                                                         # a mesh with only 2 devices
        with jax.set_mesh(mesh2):
            e2.call(jax.device_put(w, P()), jax.device_put(x, P("data")))
        check("1 call on 2 devices should fail", False)
    except ValueError as err:
        check("1 call on a 2-device mesh → ValueError", "Length of device assignment 2" in str(err), str(err)[:90])
    mesh_rev = jax.make_mesh((4,), ("data",), devices=jax.devices()[::-1])
    with jax.set_mesh(mesh_rev):                                 # the same 4 devices in a different order
        y = e2.call(jax.device_put(w, P()), jax.device_put(x, P("data")))
        check("1 call on a different 4-device mesh is correct", np.allclose(y, f(w, x), atol=1e-6))

    e1 = export.export(jax.jit(f))(w, x)                          # rule 2: single-device export
    with jax.set_mesh(mesh):
        y = e1.call(ws, xs)
        check("2 single-device artifact accepts sharded inputs", e1.nr_devices == 1 and np.allclose(y, f(w, x), atol=1e-6), y.sharding.spec)

    g = lambda v: lax.psum(jnp.sum(v), "data")                   # rule 3: shard_map
    with jax.set_mesh(mesh):
        sm = jax.jit(jax.shard_map(g, in_specs=P("data"), out_specs=P()))
        e3 = export.export(sm)(xs)
        txt = e3.mlir_module()
        check("3 shard_map exports and contains all_reduce", e3.nr_devices == 4 and "all_reduce" in txt)
        check("3 shard_map result is correct", float(e3.call(xs)) == float(x.sum()), float(e3.call(xs)))

    try:                                                          # rule 4: pmap
        export.export(jax.pmap(lambda v: v * 2))(jnp.ones((4, 3))); check("4 exporting pmap directly should fail", False)
    except ValueError as err:
        check("4 exporting pmap directly → ValueError", "must be the result of `jit`" in str(err))
    ep = export.export(jax.jit(jax.pmap(lambda v: v * 2)))(jnp.ones((4, 3)))
    check("4 jit(pmap) exports", ep.nr_devices == 4, ep.nr_devices)
"""
out = run_py(shard_src, env_extra={"XLA_FLAGS": "--xla_force_host_platform_device_count=4"})
print(out)
lines = [l for l in out.splitlines() if l.startswith(("PASS", "FAIL"))]
assert len(lines) == 11 and all(l.startswith("PASS") for l in lines), lines

ckpt_path = SCRATCH / "single_ckpt"
state0 = {"w": jnp.arange(6.).reshape(2, 3), "step": jnp.array(0)}
ckptr.save(ckpt_path, state0)
ckptr.wait_until_finished()

# ❌ Wrong: saving to the same directory again (for example, overwriting one "latest checkpoint" directory every epoch)
try:
    ckptr.save(ckpt_path, {"w": state0["w"] + 1, "step": jnp.array(1)})
    ckptr.wait_until_finished()
    print("Measured: it overwrites")
except ValueError as e:
    err_exists = str(e)
    print("Measured error: ValueError:", err_exists.replace(str(SCRATCH), "<SCRATCH>"))
assert "already exists" in err_exists

# ✅ Correct: when you really mean to overwrite, pass force=True explicitly; for multi-step saving use CheckpointManager (§3.2), one subdirectory per step
state1 = {"w": state0["w"] + 1, "step": jnp.array(1)}
ckptr.save(ckpt_path, state1, force=True)
ckptr.wait_until_finished()
restored1 = ckptr.restore(ckpt_path, sds(state1))
assert int(restored1["step"]) == 1 and max_diff(restored1, state1) == 0.0
print("force=True overwrote it, and what we read back is the new content ✓")

mgr_dir = SCRATCH / "mgr_interval"
mngr = ocp.CheckpointManager(mgr_dir, options=ocp.CheckpointManagerOptions(max_to_keep=2, save_interval_steps=2))
saved_flags = [mngr.save(step, args=ocp.args.StandardSave({"w": jnp.ones(3) * step})) for step in range(7)]
mngr.wait_until_finished()
print("Return value of save() at each step:", saved_flags)
print("Steps kept on disk:", mngr.all_steps(), "| subdirectories:", sorted(p.name for p in mgr_dir.iterdir()))
print("latest_step:", mngr.latest_step())
assert saved_flags == [True, False, True, False, True, False, True]
assert mngr.all_steps() == [4, 6] and mngr.latest_step() == 6
restored = mngr.restore(6, args=ocp.args.StandardRestore(jax.eval_shape(lambda: {"w": jnp.ones(3)})))
assert max_diff(restored, {"w": jnp.ones(3) * 6}) == 0.0
mngr.close()

losses = [5.0, 1.0, 4.0, 2.0, 9.0]          # step 4 (the latest) has the worst validation loss

# ❌ Wrong: relying only on best_fn + max_to_keep and expecting it to "keep the best 2 + the latest"
try:
    mngr = ocp.CheckpointManager(SCRATCH / "mgr_best", options=ocp.CheckpointManagerOptions(
        max_to_keep=2, best_fn=lambda m: m["loss"], best_mode="min"))
    for step, l in enumerate(losses):
        mngr.save(step, args=ocp.args.StandardSave({"w": jnp.ones(3) * step}), metrics={"loss": l})
    mngr.wait_until_finished()
    print("Measured: steps kept", mngr.all_steps(), "| best_step", mngr.best_step(), "| latest_step", mngr.latest_step())
    kept_bad, latest_bad = mngr.all_steps(), mngr.latest_step()
    mngr.close()
except Exception as e:
    print("Measured error:", type(e).__name__, e)
assert kept_bad == [1, 3] and 4 not in kept_bad          # the latest step, 4, was deleted
assert latest_bad == 3                                   # latest_step() also "went back" to step 3

# ✅ Correct: use preservation_policy to combine "the latest 1" and "the best 2"
cm = ocp.checkpoint_managers
policy = cm.AnyPreservationPolicy([
    cm.LatestN(n=1),                                                    # for resuming training
    cm.BestN(get_metric_fn=lambda m: m["loss"], reverse=True, n=2),     # reverse=True: smaller metric is better (measured)
])
mngr = ocp.CheckpointManager(SCRATCH / "mgr_policy", options=ocp.CheckpointManagerOptions(
    preservation_policy=policy, best_fn=lambda m: m["loss"], best_mode="min"))
for step, l in enumerate(losses):
    mngr.save(step, args=ocp.args.StandardSave({"w": jnp.ones(3) * step}), metrics={"loss": l})
mngr.wait_until_finished()
print("Steps kept:", mngr.all_steps(), "| best_step:", mngr.best_step(), "| latest_step:", mngr.latest_step())
assert mngr.all_steps() == [1, 3, 4] and mngr.best_step() == 1 and mngr.latest_step() == 4

best = mngr.restore(mngr.best_step(), args=ocp.args.StandardRestore(jax.eval_shape(lambda: {"w": jnp.ones(3)})))
assert max_diff(best, {"w": jnp.ones(3) * 1}) == 0.0
mngr.close()
print("The latest 1 + the best 2 are all there; restoring best_step gives the weights from step 1 ✓")

kept = {}
for rev in (False, True):
    m_ = ocp.CheckpointManager(SCRATCH / f"mgr_rev_{rev}", options=ocp.CheckpointManagerOptions(
        preservation_policy=cm.BestN(get_metric_fn=lambda m: m["loss"], reverse=rev, n=2)))
    for step, l in enumerate(losses):
        m_.save(step, args=ocp.args.StandardSave({"w": jnp.ones(3)}), metrics={"loss": l})
    m_.wait_until_finished()
    kept[rev] = m_.all_steps()
    m_.close()
print("losses =", losses, "→ reverse=False keeps", kept[False], "; reverse=True keeps", kept[True])
assert kept[False] == [0, 4] and kept[True] == [1, 3]

init_fn = lambda key: init_mlp(key)
abstract = jax.eval_shape(init_fn, jax.random.key(0))        # allocates no memory
print("Abstract template:", jax.tree.map(lambda s: f"{s.dtype}{list(s.shape)}", abstract))
ckptr.save(SCRATCH / "mlp_params", params)
ckptr.wait_until_finished()

restored_f32 = ckptr.restore(SCRATCH / "mlp_params", abstract)
assert max_diff(restored_f32, params) == 0.0

abstract_bf16 = jax.tree.map(lambda s: jax.ShapeDtypeStruct(s.shape, jnp.bfloat16), abstract)
restored_bf16 = ckptr.restore(SCRATCH / "mlp_params", abstract_bf16)
print("Restored with a bf16 template:", {k: str(v.dtype) for k, v in restored_bf16.items()})
assert all(v.dtype == jnp.bfloat16 for v in restored_bf16.values())
assert max_diff(restored_bf16, params) < 0.02
print("Restoring with the eval_shape template is bit-identical in fp32; restoring with a bf16 template converts automatically ✓")

# ❌ Wrong: the template's shapes don't match what's on disk (say, you changed the hidden width but load an old checkpoint)
wrong = dict(abstract, w1=jax.ShapeDtypeStruct((4, 16), jnp.float32))
try:
    ckptr.restore(SCRATCH / "mlp_params", wrong)
    print("Measured: it restores")
except ValueError as e:
    err_shape = str(e)
    print("Measured error: ValueError:", err_shape[:150], "...")
assert "is not compatible with the stored shape: (4, 8)" in err_shape

# ✅ Correct: look at the structure and shapes on disk with metadata() first, then decide on the template (or confirm which model version you're loading)
meta_tree = ckptr.metadata(SCRATCH / "mlp_params").item_metadata.tree
print("Shapes on disk:", {k: tuple(v.shape) for k, v in meta_tree.items()})
assert {k: tuple(v.shape) for k, v in meta_tree.items()} == {k: v.shape for k, v in params.items()}

class MLP(nnx.Module):
    def __init__(self, din, dh, dout, *, rngs: nnx.Rngs):
        self.fc1 = nnx.Linear(din, dh, rngs=rngs)
        self.fc2 = nnx.Linear(dh, dh, rngs=rngs)
        self.fc3 = nnx.Linear(dh, dout, rngs=rngs)

    def __call__(self, x):
        x = nnx.relu(self.fc1(x))
        x = nnx.relu(self.fc2(x))
        return self.fc3(x)

TEACHER = jax.random.normal(jax.random.key(100), (4, 3))      # "teacher" matrix used to generate labels

def make_batch(data_key, step, teacher, n=128):
    """The data batch for a given step: determined by (data_key, step), so restoring step restores the data position."""
    x = jax.random.normal(jax.random.fold_in(data_key, step), (n, 4))
    y = jnp.argmax(jnp.sin(2.0 * x) @ teacher, axis=-1)
    return x, y

def loss_fn(model, x, y):
    return optax.softmax_cross_entropy_with_integer_labels(model(x), y).mean()

@nnx.jit
def train_step(model, optimizer, data_key, teacher):
    x, y = make_batch(data_key, optimizer.step[...], teacher)          # data position comes from the optimizer step
    loss, grads = nnx.value_and_grad(loss_fn)(model, x, y)
    optimizer.update(model, grads)
    return loss

TX = optax.adam(1e-2)
def make_model_and_opt():
    model = MLP(4, 32, 3, rngs=nnx.Rngs(0))
    return model, nnx.Optimizer(model, TX, wrt=nnx.Param)

DATA_KEY = jax.random.key(42)
N1, N2 = 30, 30

# (a) Train N1 + N2 steps straight through, as the reference
model_ref, opt_ref = make_model_and_opt()
curve_ref = [float(train_step(model_ref, opt_ref, DATA_KEY, TEACHER)) for _ in range(N1 + N2)]

# (b) Train N1 steps, then save
model_a, opt_a = make_model_and_opt()
curve_a = [float(train_step(model_a, opt_a, DATA_KEY, TEACHER)) for _ in range(N1)]
nnx_dir = SCRATCH / "nnx_ckpt"
mngr = ocp.CheckpointManager(nnx_dir, options=ocp.CheckpointManagerOptions(max_to_keep=3))
mngr.save(int(opt_a.step[...]), args=ocp.args.StandardSave({"model": nnx.state(model_a), "opt": nnx.state(opt_a)}),
          metrics={"loss": curve_a[-1]})
mngr.wait_until_finished()
print("Saved steps:", mngr.all_steps(), "| State type:", type(nnx.state(model_a)).__name__)
assert mngr.all_steps() == [N1]
del model_a, opt_a                                   # simulate the process exiting: the in-memory model is gone