diff --git a/docs/exporting_bayesflow_models.md b/docs/exporting_bayesflow_models.md index 22b7eae..dd8b3ed 100644 --- a/docs/exporting_bayesflow_models.md +++ b/docs/exporting_bayesflow_models.md @@ -1,5 +1,8 @@ # Exporting bayesflow-trained networks to ONNX +> The artifact rules every exporter here satisfies are collected in +> [The ONNX likelihood contract](https://lnccbrown.github.io/HSSM/how_to/custom_onnx_likelihoods/). + LANfactory's [`transform_bayesflow_to_onnx`](api/onnx.md) is the bayesflow sibling of [`transform_sbi_to_onnx`](exporting_sbi_models.md). It wraps a trained [`bayesflow`](https://github.com/bayesflow-org/bayesflow) @@ -220,7 +223,7 @@ the caller. | Non-identity adapters | Numpy-only operations can't be baked into ONNX; see Constraint 4 above. Pointwise tensor adapter ops (log, sqrt, scale) are a candidate for v1.x. | | Transformer / attention summary networks | Contain `LayerNormalization` (no jaxonnxruntime handler) and dynamic-shape attention. | | FlowMatching, DiffusionModel, ConsistencyModel inference networks | `log_prob` requires ODE integration, not ONNX-exportable. | -| `KERAS_BACKEND=jax` workflows | Use the bayesflow LRE-style in-memory JAX callable path (see [`bayesflow_lre_integration.ipynb`](https://github.com/lnccbrown/HSSM/blob/main/docs/tutorials/bayesflow_lre_integration.ipynb) in HSSM). | +| `KERAS_BACKEND=jax` workflows | Use the bayesflow LRE-style in-memory JAX callable path (see [`bayesflow_lre_integration.ipynb`](https://lnccbrown.github.io/HSSM/tutorials/bayesflow_lre_integration/) in HSSM). | ## Numerical guarantees diff --git a/docs/exporting_sbi_models.md b/docs/exporting_sbi_models.md index 051b90e..c45e539 100644 --- a/docs/exporting_sbi_models.md +++ b/docs/exporting_sbi_models.md @@ -3,7 +3,8 @@ LANfactory's [`transform_sbi_to_onnx`](api/onnx.md) wraps a trained [`sbi`](https://github.com/sbi-dev/sbi) estimator and writes a single-trial ONNX file that HSSM's `loglik_kind="approx_differentiable"` path can consume -exactly like a LAN export. Use it to bring sbi-trained NLE density estimators +exactly like a LAN export (the artifact rules are collected in +[The ONNX likelihood contract](https://lnccbrown.github.io/HSSM/how_to/custom_onnx_likelihoods/)). Use it to bring sbi-trained NLE density estimators or NRE ratio classifiers into a [HSSM](https://github.com/lnccbrown/HSSM) model. ## Installation diff --git a/docs/index.md b/docs/index.md index 1a63fba..d8b34a8 100755 --- a/docs/index.md +++ b/docs/index.md @@ -93,7 +93,8 @@ transform-onnx --network-config-file \ ``` The resulting file can be used directly with -[HSSM](https://lnccbrown.github.io/HSSM/). Networks trained outside LANfactory +[HSSM](https://lnccbrown.github.io/HSSM/) — see +[The ONNX likelihood contract](https://lnccbrown.github.io/HSSM/how_to/custom_onnx_likelihoods/) for the artifact rules. Networks trained outside LANfactory can be exported the same way — see the [sbi](exporting_sbi_models.md) and [bayesflow](exporting_bayesflow_models.md) export guides. diff --git a/docs/tutorials/exporting_bayesflow_to_onnx.ipynb b/docs/tutorials/exporting_bayesflow_to_onnx.ipynb index 90e0818..80e96e9 100644 --- a/docs/tutorials/exporting_bayesflow_to_onnx.ipynb +++ b/docs/tutorials/exporting_bayesflow_to_onnx.ipynb @@ -202,8 +202,11 @@ "\n", "`transform_bayesflow_to_onnx` bakes the standardizer's accumulated mean/std\n", "into the graph as constants (sidestepping the dynamic-shape ops the live Keras\n", - "layer would emit) and writes a **rank-1** single-trial graph, opset 17 — the\n", - "same contract as the sbi and LAN exporters, so HSSM consumes it identically." + "layer would emit) and writes a **rank-1** single-trial graph, opset 17. It satisfies the same\n", + "contract as the sbi and LAN exporters — **every input dim concrete, no dynamic\n", + "axes** ([canonical statement](https://lnccbrown.github.io/HSSM/how_to/custom_onnx_likelihoods/)) — though rank differs by tracer:\n", + "the LAN exporters emit rank-2 `(1, D)` Gemm graphs; rank-1 is required here\n", + "because the flow graph slices its input." ] }, { @@ -394,7 +397,8 @@ "source": [ "## 4. Consume it in HSSM\n", "\n", - "The `.onnx` file drops into HSSM exactly like a LAN or sbi export:\n", + "The `.onnx` file drops into HSSM exactly like a LAN or sbi export\n", + "(full rules: [The ONNX likelihood contract](https://lnccbrown.github.io/HSSM/how_to/custom_onnx_likelihoods/)):\n", "\n", "```python\n", "import jax\n", @@ -412,7 +416,7 @@ "```\n", "\n", "For the consumption side end to end, see HSSM's\n", - "[Build HSSM models starting from ONNX files](https://github.com/lnccbrown/HSSM/blob/main/docs/tutorials/blackbox_contribution_onnx_example.ipynb)\n", + "[Build HSSM models starting from ONNX files](https://lnccbrown.github.io/HSSM/tutorials/blackbox_contribution_onnx_example/)\n", "tutorial. The **NRE** path is analogous: train a `RatioApproximator` and pass\n", "`mode=\"nre\"` to `transform_bayesflow_to_onnx`." ] diff --git a/docs/tutorials/exporting_sbi_to_onnx.ipynb b/docs/tutorials/exporting_sbi_to_onnx.ipynb index 40c6f3c..49aff29 100644 --- a/docs/tutorials/exporting_sbi_to_onnx.ipynb +++ b/docs/tutorials/exporting_sbi_to_onnx.ipynb @@ -512,8 +512,10 @@ "## 2. Export to ONNX\n", "\n", "`transform_sbi_to_onnx` wraps the trained estimator into a **rank-1**\n", - "single-trial graph (parameters first, observations second; opset 17). The\n", - "rank-1 contract is what lets HSSM `vmap` the graph over trials." + "single-trial graph (parameters first, observations second; opset 17). The graph is\n", + "rank-1 because flow/ratio graphs *slice* their input — a `(1, D)` trace bakes\n", + "`Slice` axes that fail under HSSM's `vmap`. The contract invariant itself is\n", + "**concrete dims, no dynamic axes** ([canonical statement](https://lnccbrown.github.io/HSSM/how_to/custom_onnx_likelihoods/))." ] }, { @@ -699,7 +701,8 @@ "## 4. Consume it in HSSM\n", "\n", "The `.onnx` file drops into HSSM exactly like a LAN export — HSSM handles\n", - "the `vmap` over trials and (recent versions) the x64 flag:\n", + "the `vmap` over trials and (recent versions) the x64 flag —\n", + "the full rules are in [The ONNX likelihood contract](https://lnccbrown.github.io/HSSM/how_to/custom_onnx_likelihoods/):\n", "\n", "```python\n", "import jax\n", diff --git a/src/lanfactory/onnx/jax_export.py b/src/lanfactory/onnx/jax_export.py index 2b4f397..bfcbbda 100644 --- a/src/lanfactory/onnx/jax_export.py +++ b/src/lanfactory/onnx/jax_export.py @@ -5,8 +5,10 @@ convert — ``transform-onnx`` only reads torch state dicts, and HSSM consumes ONNX exclusively. -The export follows the ecosystem's single-trial ONNX contract (see HSSMSpine -CLAUDE.md): the graph is traced with a **concrete** ``(1, input_dim)`` dummy +The export follows the ecosystem's single-trial ONNX contract; the canonical +statement lives in HSSM's docs: +https://lnccbrown.github.io/HSSM/how_to/custom_onnx_likelihoods/ +The graph is traced with a **concrete** ``(1, input_dim)`` dummy and no dynamic axes, exactly like the torch MLP exporter and the production networks on franklab/HSSM. Every dim is static, so HSSM's load-time check passes, and HSSM's rank-1-per-trial + ``jax.vmap`` consumption works because diff --git a/src/lanfactory/onnx/sbi.py b/src/lanfactory/onnx/sbi.py index 236c8ac..3e9bfb3 100644 --- a/src/lanfactory/onnx/sbi.py +++ b/src/lanfactory/onnx/sbi.py @@ -9,9 +9,10 @@ exporter): "train a network and emit an ONNX HSSM can read" stays a single conceptual home in LANfactory regardless of which library trained the network. -The exported graph follows the LAN-and-HSSM convention: a single concatenated -input of **rank 1, shape ``(theta_dim + x_dim,)``**. Inside the graph the -input is split into ``theta`` and ``x``, upranked to ``(1, …)`` to satisfy +The exported graph takes a single concatenated single-trial input of +**rank 1, shape ``(theta_dim + x_dim,)``**. The contract invariant is concrete +dims; rank is tracer-specific (the LAN exporters emit rank-2 ``(1, D)`` +``Gemm`` graphs). Inside the graph the input is split into ``theta`` and ``x``, upranked to ``(1, …)`` to satisfy sbi's batched ``log_prob`` API, and routed through the trained estimator. HSSM vmaps this graph over trials, so the per-call input rank from HSSM is 1 — matching the export. Tracing with a 2D ``(1, D)`` dummy would emit ``Slice``