research: final compile-time census + close-out (#77)

jax_compile/README.md

@@ -208,3 +208,55 @@ artifact from job 330596 exists if ever needed).
 
 Not indicated: source restructuring, jit boundaries inside likelihoods,
 replacing `lax.map` in MultiStartAdam, autotune flags.
+
+## Final census — the defaults-live user experience (issue #77, 2026-07-17)
+
+Measured through the merged wrapper defaults (#128 cache + #132 autotune-off),
+no manual flags; cold = fresh cache dir, warm = same dir, fresh process.
+
+**A100 (dedicated node)** — trace + XLA compile, seconds:
+
+| cell | cold | warm | steady eval |
+|---|---|---|---|
+| mge / jit | 4.7 + 5.7 | 4.8 + 0.3 | 0.0042 s |
+| mge / vag | 6.8 + 27.9 | 6.8 + 2.4 | 0.0098 s |
+| pix / jit | 4.0 + 5.7 | 4.2 + 0.4 | 0.0574 s |
+| pix / vag | 4.7 + 20.5 | 4.7 + 1.8 | 0.0911 s |
+
+Worst cold cell ≈ **35 s** (was ~70 min in the worst pre-#128/#132 case);
+warm ≈ **5–9 s**, almost entirely tracing.
+
+**Local CPU (idle laptop; cross-day compile variance up to 2× — treat the A100
+table as the reference):**
+
+| cell | cold | warm | 
+|---|---|---|
+| mge / jit | 21.0 + 16.9 | 10.9 + 0.5 |
+| mge / vag | 22.1 + 229.4 | 16.3 + 2.9 |
+| mge / lax.map∘vag | 31.4 + 369.9 | 29.2 + 6.2 |
+| pix / jit | 6.7 + 7.7 | 7.6 + 0.3 |
+| pix / vag | 7.3 + 34.6 | 8.2 + 1.6 |
+
+## Where any remaining speedup would come (final)
+
+1. **`jax.export` — RULED OUT** (jax 0.10.2): deserialize is 4 ms (vs 16–22 s
+   tracing) but the exported module recompiles **~156 s in every process** —
+   its compile bypasses the persistent compilation cache (`export_probe.py`,
+   two independent load processes). The standard warm path (trace + cached
+   compile) wins by an order of magnitude. Re-evaluate only if a future jax
+   caches exported-call compiles.
+2. **XLA compile-parallelism flags — inconclusive, low ceiling.** CPU
+   same-day: 137.7 s vs 229.4 s cold vag compile with
+   `--xla_cpu_parallel_codegen_split_count=32`, but cross-day variance on the
+   same cell (117 ↔ 229 s) swamps the signal. On the A100 the residual
+   no-autotune compile is 20–28 s, so even a 2× flag win saves ~10 s once per
+   machine — not worth productizing.
+3. **The tracing floor (warm cost) is jax-internal** (finding 8: 58 % jax /
+   34 % stdlib / 7 % autoarray). 4–7 s per transform on the A100 host, 11–29 s
+   on the laptop, every process. Movable only by upstream jax tracing speed or
+   by emitting fewer ops.
+
+**Close-out:** with cache + autotune-off shipped, compile cost is seconds at
+every point in the lifecycle. Any further reduction would come from upstream
+JAX (tracing speed, compile speed), not from this stack — no further
+engineering is warranted here. The compile-time arc (#71 → #74 → #77) is done.

jax_compile/export_probe.py

@@ -0,0 +1,116 @@
+"""
+Can ``jax.export`` beat the tracing floor? (final compile-time census, issue #77)
+
+After the cache (#128) and autotune-off (#132) defaults, the residual warm-run
+cost is Python *tracing* (~5-17 s per transform per process), which the
+compilation cache cannot remove. ``jax.export`` serializes the traced/lowered
+StableHLO to disk, so a repeat process could deserialize instead of retracing —
+this probe measures whether that wins.
+
+Two modes, run as separate processes so nothing is warm in-process::
+
+    python jax_compile/export_probe.py --mode save   # trace + export + serialize to disk
+    python jax_compile/export_probe.py --mode load   # deserialize + first call + steady
+
+Compare `load`'s deserialize+first against the standard warm path's
+trace+compile+first (probe.py census-warm rows). Uses the MGE ``vag`` cell.
+Appends records to ``jax_compile/results/<hardware>/export_probe.json``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+_WORKSPACE_ROOT = Path(__file__).resolve().parents[1]
+if str(_WORKSPACE_ROOT) not in sys.path:
+    sys.path.insert(0, str(_WORKSPACE_ROOT))
+
+from jax_compile.probe import build_objective, hardware_label
+
+EXPORT_PATH = Path("/tmp/claude-1000/census_export_mge_vag.bin")
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description=__doc__.split("\n")[1])
+    p.add_argument("--mode", choices=("save", "load"), required=True)
+    p.add_argument("--dataset-class", default="imaging")
+    p.add_argument("--model-type", default="mge")
+    p.add_argument("--instrument", default="hst")
+    p.add_argument("--mixed-precision", action="store_true")
+    p.add_argument("--tag", default="")
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    import jax
+    from jax import export as jax_export
+
+    record = {"mode": args.mode, "model_type": args.model_type, "transform": "vag"}
+
+    if args.mode == "save":
+        f, x0, ndim = build_objective(args)
+        vag = jax.value_and_grad(f)
+
+        t0 = time.perf_counter()
+        exported = jax_export.export(jax.jit(vag))(
+            jax.ShapeDtypeStruct(x0.shape, x0.dtype)
+        )
+        record["trace_export_s"] = round(time.perf_counter() - t0, 3)
+
+        t0 = time.perf_counter()
+        blob = exported.serialize()
+        EXPORT_PATH.write_bytes(blob)
+        record["serialize_s"] = round(time.perf_counter() - t0, 3)
+        record["blob_mb"] = round(len(blob) / 1e6, 3)
+        record["ndim"] = ndim
+        print(
+            f"[export_probe] save: trace+export {record['trace_export_s']}s  "
+            f"serialize {record['serialize_s']}s  blob {record['blob_mb']}MB"
+        )
+    else:
+        # The load path deliberately does NOT build the objective — the whole
+        # point is skipping the model/dataset tracing. It only needs an input
+        # array of the right shape.
+        import jax.numpy as jnp
+
+        t0 = time.perf_counter()
+        exported = jax_export.deserialize(EXPORT_PATH.read_bytes())
+        record["deserialize_s"] = round(time.perf_counter() - t0, 3)
+
+        in_shape = exported.in_avals[0].shape
+        x = jnp.full(in_shape, 0.5, dtype=exported.in_avals[0].dtype)
+
+        t0 = time.perf_counter()
+        out = exported.call(x)
+        jax.block_until_ready(out)
+        record["first_s"] = round(time.perf_counter() - t0, 3)
+
+        t0 = time.perf_counter()
+        for _ in range(3):
+            out = exported.call(x)
+        jax.block_until_ready(out)
+        record["steady_s"] = round((time.perf_counter() - t0) / 3, 4)
+        print(
+            f"[export_probe] load: deserialize {record['deserialize_s']}s  "
+            f"first {record['first_s']}s  steady {record['steady_s']}s"
+        )
+
+    record["jax_version"] = jax.__version__
+    record["tag"] = args.tag
+    record["timestamp"] = time.strftime("%Y-%m-%dT%H:%M:%S")
+    out_dir = _WORKSPACE_ROOT / "jax_compile" / "results" / hardware_label(jax)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out_path = out_dir / "export_probe.json"
+    existing = json.loads(out_path.read_text()) if out_path.exists() else []
+    existing.append(record)
+    out_path.write_text(json.dumps(existing, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()

jax_compile/results/local_cpu/export_probe.json

@@ -0,0 +1,36 @@
+[
+  {
+    "mode": "save",
+    "model_type": "mge",
+    "transform": "vag",
+    "trace_export_s": 22.311,
+    "serialize_s": 0.006,
+    "blob_mb": 1.936,
+    "ndim": 15,
+    "jax_version": "0.10.2",
+    "tag": "census",
+    "timestamp": "2026-07-17T09:59:45"
+  },
+  {
+    "mode": "load",
+    "model_type": "mge",
+    "transform": "vag",
+    "deserialize_s": 0.004,
+    "first_s": 159.5,
+    "steady_s": 0.2963,
+    "jax_version": "0.10.2",
+    "tag": "census",
+    "timestamp": "2026-07-17T10:02:27"
+  },
+  {
+    "mode": "load",
+    "model_type": "mge",
+    "transform": "vag",
+    "deserialize_s": 0.004,
+    "first_s": 155.995,
+    "steady_s": 0.3651,
+    "jax_version": "0.10.2",
+    "tag": "census-warm",
+    "timestamp": "2026-07-17T10:08:18"
+  }
+]