feat(onnx): bayesflow → ONNX exporter (transform_bayesflow_to_onnx)

[AlexanderFengler]

Summary

Adds lanfactory.onnx.transform_bayesflow_to_onnx — the bayesflow sibling of transform_sbi_to_onnx from #79. Wraps a trained bayesflow.ContinuousApproximator (NLE) or RatioApproximator (NRE) and writes a single-trial ONNX file consumable by HSSM's loglik_kind="approx_differentiable" path.

The user gesture becomes the same regardless of training framework: hssm.HSSM(loglik="model.onnx", loglik_kind="approx_differentiable") works for LAN, sbi, and now bayesflow exports through the identical HSSM-side code path (no HSSM changes required).

Branch relationship

This branch is stacked on sbi-connector (#79). The diff in this PR is exactly the bayesflow-specific additions; #79's commits are not duplicated. When #79 merges to main, GitHub will offer to auto-retarget this PR's base to main.

What's in this PR

File	Action	Purpose
src/lanfactory/onnx/bayesflow.py	new	Exporter module (NLE + NRE wrappers + transform_bayesflow_to_onnx)
src/lanfactory/onnx/__init__.py	edit	Export the new function
pyproject.toml	edit	New [bayesflow] optional extra; refactor existing sbi+nflows pair into a new [sbi] extra (symmetric with [bayesflow]); both added to [all] and [dev]
tests/test_bayesflow_nle_export.py	new	6 tests: three-way numerical agreement (atol=1e-5), gradient agreement (atol=1e-4), log-prob ordering sanity, three guard tests
tests/test_bayesflow_nre_export.py	new	4 tests: same shape for the NRE path
tests/test_bayesflow_hssm_integration.py	new	End-to-end DDM smoke (pytest.importorskip("hssm"), mirrors test_sbi_hssm_integration.py)
docs/exporting_bayesflow_models.md	new	Sibling of exporting_sbi_models.md

Architectural contract

Same I/O contract as the sbi exporter:

• Input: rank-1 tensor of shape (theta_dim + x_dim,), parameters first then observations.
• Output: rank-0 scalar log-likelihood.
• Opset: pinned to 17 for jaxonnxruntime reproducibility.

Key implementation choice: the wrapper bakes the bayesflow Standardize layer's accumulated moving_mean / moving_std as torch buffer constants at construction time. This sidesteps the dynamic-shape ops (If, Size, Tile) that the live Keras layer would emit at trace time — jaxonnxruntime doesn't have a Size handler. The constants are correct because training is complete by export time.

v1 constraints (documented, enforced where introspectable)

User must train with:

• permutation=None (FixedPermutation → aten::ravel, unsupported in opset 17/20)
• use_actnorm=False (untested in v1)
• transform=AffineTransform(clamp=False) as an explicit instance (find_transform("affine") silently drops transform_kwargs — bayesflow upstream bug, catalogued in companion HSSMSpine PR)
• subnet_kwargs.activation="silu" (default hard_silu exports as the fused ONNX op HardSwish, no jaxonnxruntime handler; silu decomposes to Sigmoid + Mul)
• Identity Adapter (numpy-only Adapter ops can't be baked into ONNX)

Each violation produces an actionable error message at export time.

Test status

• bayesflow NLE: 6/6 passing
• bayesflow NRE: 4/4 passing
• bayesflow ↔ HSSM integration: skip-on-missing-HSSM in this env (designed for the coordinated cross-repo CI matrix)
• sbi regression check: existing sbi tests still pass (no regression). Note: the bayesflow test modules call torch.set_grad_enabled(True) after importing bayesflow to undo the global autograd disable that bayesflow's torch backend does at import time. This is a known upstream issue documented in the companion HSSMSpine PR's upstream-bugs catalog.

Companion PRs

• HSSM: bayesflow-integration branch (new), adds docs/tutorials/bayesflow_nle_onnx_integration.ipynb. Sibling, not child, of sbi-integration (#964) — works against HSSM main standalone (Part 1 includes the manual jaxort_only_allow_initializers_as_static_args=False workaround that #964 plans to auto-handle inside onnx2jax).
• HSSMSpine: bayesflow-onnx-plans branch (stacked on NameError when wandb is not found #9 for the cross-reference to sbi-onnx-integration.md). Adds the design doc and an upstream-bugs catalog covering the five real upstream defects surfaced during this work.

Test plan

• bayesflow NLE tests pass (6/6)
• bayesflow NRE tests pass (4/4)
• No regression on existing sbi tests
• Cross-repo HSSM integration test (pytest.importorskip("hssm")) once both packages can be installed in the same env

🤖 Generated with Claude Code