Integrate bayesflow NLE via ONNX exporter (tutorial only)#965
Open
AlexanderFengler wants to merge 4 commits into
Open
Integrate bayesflow NLE via ONNX exporter (tutorial only)#965AlexanderFengler wants to merge 4 commits into
AlexanderFengler wants to merge 4 commits into
Conversation
Sibling of sbi_nle_integration.ipynb. Demonstrates the end-to-end path: train a bayesflow CouplingFlow NLE on DDM data, export to ONNX via lanfactory.onnx.transform_bayesflow_to_onnx (companion LANfactory PR), and hand the file to HSSM via the same loglik="*.onnx" / loglik_kind="approx_differentiable" gesture that sbi exports already use. Final part overlays the bayesflow-NLE posterior against HSSM's analytical Navarro & Fuss DDM posterior as ground truth. Part 1 includes a temporary workaround line setting jaxort_only_allow_initializers_as_static_args=False directly so the notebook runs standalone on main. HSSM PR #964 sets this automatically inside hssm.distribution_utils.onnx2jax; once it merges, the manual line can be deleted. This is a docs-only addition - no HSSM code changes. Companion PRs: - LANfactory bayesflow-connector (stacked on #79 sbi-connector) - HSSMSpine: bayesflow-onnx-integration.md design doc + upstream-bugs-from-bayesflow-onnx-work.md catalog of upstream defects surfaced during this work Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Replaces the hardcoded ./bayesflow_onnx_artifacts/ path in Part 4 with a clearly-named ARTIFACT_DIR constant that defaults to ~/bayesflow_onnx_tutorial/ (outside the HSSM repo). Comment block in the cell shows two override examples: an arbitrary user path, or tempfile.mkdtemp() for an ephemeral demo run. Adds `import tempfile` in Part 1 so the override example works out of the box. Why: the previous default (./bayesflow_onnx_artifacts/) wrote into whatever directory the notebook was running from, which for typical setups meant docs/tutorials/ inside the HSSM repo. Re-running the notebook would accumulate untracked artifacts in the working tree (the same pattern that left sbi-logs/ and sbi_onnx_artifacts/ around from earlier sbi tutorial work). Moving the default outside the repo eliminates that footgun. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new documentation tutorial notebook demonstrating how to train a BayesFlow NLE (CouplingFlow), export it to ONNX via LANfactory, and use the exported file in HSSM through the existing loglik="*.onnx" / loglik_kind="approx_differentiable" pathway, with posterior comparisons against HSSM’s analytical DDM as a reference.
Changes:
- Add
bayesflow_nle_onnx_integration.ipynbtutorial covering BayesFlow NLE → ONNX export → HSSM consumption. - Include a full end-to-end example (simulation, training, export, numpyro sampling, analytical comparison plots).
Comments suppressed due to low confidence (1)
docs/tutorials/bayesflow_nle_onnx_integration.ipynb:248
- Same as above for the analytical reference model: consider setting
mp_ctx="spawn"(and/orcores=1) on thissample(...)call to avoid multiprocessing issues during notebook execution on some platforms.
"idata_analytical = model_analytical.sample(\n",
" sampler=\"numpyro\",\n",
" draws=1000,\n",
" tune=1000,\n",
" chains=2,\n",
" target_accept=0.9,\n",
")\n",
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+83
to
+90
| "## Part 3 — Train a bayesflow CouplingFlow NLE on DDM simulations\n", | ||
| "\n", | ||
| "The CouplingFlow setup below is the v1 ONNX-friendly configuration documented in [LANfactory's `exporting_bayesflow_models.md`](https://github.com/lnccbrown/LANFactory/blob/main/docs/exporting_bayesflow_models.md). Three opinionated choices that aren't bayesflow's defaults:\n", | ||
| "\n", | ||
| "- `permutation=None` — `FixedPermutation` uses `keras.ops.take`, which exports as `aten::ravel`, unsupported in ONNX opset 17/20.\n", | ||
| "- `transform=AffineTransform(clamp=False)` (explicit instance) — default `clamp=True` emits `aten::asinh`. Also, `bf.networks.CouplingFlow(..., transform_kwargs={\"clamp\": False})` silently drops the kwarg (upstream bug); pass an instance.\n", | ||
| "- `activation=\"silu\"` — default `\"hard_silu\"` exports as a single fused `HardSwish` op that jaxonnxruntime can't run. Real SiLU decomposes to `Sigmoid` + `Mul`." | ||
| ] |
Comment on lines
+152
to
+159
| "source": "## Part 4 — Export the trained approximator to ONNX\n\nOne call. The exporter raises clearly if any v1 constraint is violated (wrong `KERAS_BACKEND`, non-identity adapter, missing `inference_network`).\n\n**Where to write the file.** The cell below sets `ARTIFACT_DIR` to `~/bayesflow_onnx_tutorial/` — outside the HSSM repo, so re-running the notebook doesn't leave artifacts in your working tree. Change it to any path you want to keep the trained ONNX around for downstream work; set it to `tempfile.mkdtemp()` for an ephemeral demo run." | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": "# User-configurable: where the .onnx file lands. Default is outside the HSSM\n# repo so notebook re-runs don't pollute the working tree.\n# Override examples:\n# ARTIFACT_DIR = Path(\"/path/to/my/project/onnx\") # keep nearby\n# ARTIFACT_DIR = Path(tempfile.mkdtemp()) # ephemeral\nARTIFACT_DIR = Path.home() / \"bayesflow_onnx_tutorial\"\nARTIFACT_DIR.mkdir(parents=True, exist_ok=True)\nonnx_path = ARTIFACT_DIR / \"ddm_nle.onnx\"\n\ntransform_bayesflow_to_onnx(\n approximator,\n str(onnx_path),\n mode=\"nle\",\n example_theta_dim=len(DDM_PARAM_NAMES),\n example_x_dim=2,\n)\nprint(f\"wrote {onnx_path} ({onnx_path.stat().st_size:,} bytes)\")" |
Comment on lines
+1
to
+6
| { | ||
| "cells": [ | ||
| { | ||
| "cell_type": "markdown", | ||
| "metadata": {}, | ||
| "source": [ |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
docs/tutorials/bayesflow_nle_onnx_integration.ipynb— the bayesflow sibling ofsbi_nle_integration.ipynb(#964). Demonstrates the full pipeline:ContinuousApproximator(CouplingFlow NLE) on synthetic DDM data.lanfactory.onnx.transform_bayesflow_to_onnx(companion LANfactory PR).loglik="*.onnx"/loglik_kind="approx_differentiable"gesture that sbi exports already use.Docs-only PR — no HSSM code changes.
Branch relationship — important context
This PR's branch (
bayesflow-integration) is a sibling ofsbi-integration(#964), not a child of it. Both targetmain. They can land in any order.Decision rationale: the bayesflow tutorial doesn't depend on any HSSM-side code from #964. The two HSSM-side ONNX patches in #964 (auto-
jax_enable_x64,jaxort_only_allow_initializers_as_static_args=False) are improvements that this tutorial would benefit from — but Part 1 of this notebook sets both flags explicitly so it runs standalone against currentmain. Once #964 lands, a small follow-up commit removes the redundant workaround lines.Architecturally this matches reality: bayesflow and sbi are independent training frameworks whose ONNX exports converge on the same HSSM-side consumption path.
What's in this PR
docs/tutorials/bayesflow_nle_onnx_integration.ipynbsbi_nle_integration.ipynb's structureSection structure:
ssm-simulatorstransform_bayesflow_to_onnxCompanion PRs
bayesflow-connector(stacked onsbi-connector/ feat: updating notebooks #79) — containstransform_bayesflow_to_onnxthat this tutorial uses.bayesflow-onnx-plans(stacked onadd-meta-tutorials-to-ingest/ Add tests for WFPT Distribution for PYMC version >=4.0.0 #9) — the design doc and an upstream-bugs catalog.v1 constraints documented in the tutorial
The CouplingFlow setup in Part 3 uses non-default values for
permutation,transform,use_actnorm, andsubnet_kwargs.activation. These are the v1 ONNX-friendly knobs documented in LANfactory'sexporting_bayesflow_models.md. The Part 3 markdown explains each choice so readers know why their training config might differ from a "vanilla" bayesflow setup.Test plan
main(manual workaround lines in Part 1 should make this work standalone)🤖 Generated with Claude Code