feat(onnx): sbi → ONNX exporter (transform_sbi_to_onnx) for NLE + NRE

[AlexanderFengler]

Summary

Adds lanfactory.onnx.transform_sbi_to_onnx, an exporter that converts a
trained sbi neural likelihood estimator
(NLE) or neural ratio estimator (NRE) to an ONNX file consumable by HSSM's
loglik_kind="approx_differentiable" path. The exported file behaves
identically to a LAN-trained ONNX from transform_to_onnx, so HSSM
consumes it with zero HSSM-side code changes for the basic case.

Lives as a sibling module to the existing LAN exporter
(lanfactory/onnx/transform_onnx.py ↔ lanfactory/onnx/sbi.py). The new
deps (sbi>=0.26, nflows>=0.14) are gated under the existing [all]
extra; no impact on the default install.

What's in this branch

#	Commit	What
C1	ac16eda	Scaffolding: stub module, [all] extra extended with sbi + nflows
C2	f7c93c8	Spike tests — torch.onnx.export → {onnxruntime, jaxonnxruntime} round-trip on a vanilla MLP and an nflows MAF. Kept as permanent regression guards.
C3	bdcabd3	Core exporter — NLE path. Wraps estimator.log_prob(x, condition=θ) as a torch.nn.Module that takes a 1D concatenated (theta, x) input. Three-way numerical agreement (torch / ORT / jaxonnxruntime) at atol=1e-5, gradient agreement at atol=1e-4.
C4	f4a54fe	NRE path — wraps the classifier logit as the log-ratio. Same numerical-agreement contract.
C5	b1fd188	Embedding-net coverage tests for FCEmbedding and CNNEmbedding.
C6	87704bb	Documentation: docs/exporting_sbi_models.md integration guide, README mention, nav update.
C7b	4990e85	End-to-end HSSM integration test (pytest.importorskip("hssm") skip-on-missing).
C9	222adf5	Critical fix — export rank-1 input contract for HSSM's vmap-over-trials. Previous 2D (1, D) dummy caused Slice axes=[1] which fails under vmap with rank-1 per-trial input.

Architectural contract

• Single-trial graph: the exported ONNX takes a 1D vector of length
  theta_dim + x_dim, matching HSSM's make_jax_logp_funcs_from_onnx
  per-trial vmap contract.
• Inside the wrapper: splits on axis 0, unsqueeze(0) to satisfy
  sbi's batched log_prob/forward API, returns scalar via reshape(()).
• Opset pinned to 17 for reproducibility against jaxonnxruntime.

Test status

• 13 LANfactory sbi tests passing locally
  (pytest tests/test_sbi_spike_*.py tests/test_sbi_nle_export.py tests/test_sbi_nre_export.py tests/test_sbi_embeddings.py).
• C7b integration test is pytest.importorskip("hssm"); clean skip
  in the current LANfactory env. Designed to run in coordinated cross-
  repo CI where both LANfactory + HSSM are installed together.

Known limitations (documented in C6 guide)

• NSF flows blocked by missing SearchSorted op in jaxonnxruntime
  (planned upstream PR, ~50 lines). Tracked in
  HSSMSpine/plans/sbi-onnx-integration.md.
• MNLE same SearchSorted blocker (categorical lookup uses
  torch.searchsorted). Same upstream PR unlocks both.
• NLE-MAF on DDM produces qualitatively wrong posteriors (rt is
  continuous but choice is discrete; MAF can't represent the mixed
  structure). The exporter still supports mode="nle" correctly —
  this is a sbi-method/data-shape issue, not an exporter bug. MNLE is
  the correct sbi method for DDM-like data; deferred until the
  upstream PR lands.

Companion PRs

Consumed by the HSSM-side keystone tutorial in
lnccbrown/HSSM#sbi-integration.
The tutorial uses an importlib fallback to load this exporter even
without LANfactory installed, so this PR can be merged independently.

Test plan

• LANfactory sbi tests pass (13/13)
• Cross-repo CI exercise of the integration test (test_sbi_hssm_integration.py)
  once both packages can be installed in the same env

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Adds an ONNX exporter for trained sbi NLE/NRE estimators so they can be consumed by HSSM’s differentiable likelihood path, alongside documentation and regression coverage for ONNX/JAX round-trips.

Changes:

• Introduces lanfactory.onnx.transform_sbi_to_onnx for rank-1 single-trial NLE/NRE ONNX export.
• Adds sbi/nflows-related optional/dev dependencies and lockfile updates.
• Adds docs and tests covering MLP/MAF round-trips, NLE/NRE exports, embeddings, and optional HSSM integration.

Reviewed changes

Copilot reviewed 12 out of 14 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
src/lanfactory/onnx/sbi.py	Implements the new sbi-to-ONNX exporter and wrappers.
src/lanfactory/onnx/__init__.py	Exposes transform_sbi_to_onnx from the ONNX package.
pyproject.toml	Adds optional and dev dependencies for sbi ONNX export/testing.
uv.lock	Locks new transitive dependencies.
tests/test_sbi_spike_mlp_roundtrip.py	Adds baseline MLP ONNX/JAX/ORT round-trip test.
tests/test_sbi_spike_maf_roundtrip.py	Adds nflows MAF ONNX/JAX/ORT round-trip test.
tests/test_sbi_nle_export.py	Adds NLE exporter numerical, gradient, and validation tests.
tests/test_sbi_nre_export.py	Adds NRE exporter numerical, gradient, and sanity tests.
tests/test_sbi_embeddings.py	Adds NRE embedding-net export round-trip tests.
tests/test_sbi_hssm_integration.py	Adds optional end-to-end HSSM integration tests.
docs/exporting_sbi_models.md	Documents installation, usage, supported architectures, and constraints.
README.md	Mentions the new sbi ONNX exporter.
mkdocs.yml	Adds the sbi exporter guide to docs navigation.
.gitignore	Ignores sbi-generated TensorBoard logs.

Comments suppressed due to low confidence (1)

src/lanfactory/onnx/sbi.py:107

• NLE mode accepts any module with .log_prob, which also includes posterior-shaped sbi estimators such as NPE/SNPE. Those estimators model p(theta | x) with the opposite input/condition semantics, so this wrapper can silently export a posterior density as if it were a likelihood (especially when theta_dim == x_dim). Add an explicit estimator-family check for true likelihood estimators before building the NLE wrapper, or require callers to pass a likelihood-specific estimator type that can be validated.

    if mode == "nle":
        if not hasattr(estimator, "log_prob"):
            raise TypeError(
                f"NLE mode requires an estimator with "
                f".log_prob(input, condition); got {estimator_cls} which lacks "
                f"it. If this is an NRE ratio classifier, use mode='nre' "
                f"instead."
            )
        wrapper: nn.Module = _NLELogProbWrapper(

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.