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

Summary by CodeRabbit

• New Features
  • Added SBI-to-ONNX exporting for neural likelihood (NLE) and neural ratio (NRE) estimators to support HSSM’s approx_differentiable workflow via a new public exporter with mode-specific validation.
• Documentation
  • Added an “Exporting sbi Models” guide, including quick-start examples and supported/excluded architectures.
• Tests
  • Added export round-trip tests (numerical + gradient checks) across PyTorch, ONNX Runtime, and JAX, plus an HSSM MCMC integration test.
• Chores
  • Expanded optional/dev dependencies for ONNX/JAX support and updated .gitignore for SBI TensorBoard logs.

[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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: ecd3726a-9c0b-486f-a1b2-d6a49d7ffe9c

📥 Commits

Reviewing files that changed from the base of the PR and between 9540c74 and b319efe.

📒 Files selected for processing (4)

• docs/exporting_sbi_models.md
• src/lanfactory/onnx/sbi.py
• tests/test_sbi_hssm_integration.py
• tests/test_sbi_nle_export.py

✅ Files skipped from review due to trivial changes (1)

• docs/exporting_sbi_models.md

🚧 Files skipped from review as they are similar to previous changes (3)

• src/lanfactory/onnx/sbi.py
• tests/test_sbi_nle_export.py
• tests/test_sbi_hssm_integration.py

---

📝 Walkthrough

Walkthrough

Adds transform_sbi_to_onnx to lanfactory.onnx, which wraps trained SBI NLE_A/NRE_A estimators in thin nn.Module adapters and traces them to ONNX via torch.onnx.export. New test modules validate three-way numerical and gradient agreement (PyTorch/ONNX Runtime/JAX) for NLE, NRE, embedding variants, and an HSSM end-to-end MCMC integration. Documentation, mkdocs.yml navigation, and project dependencies are also updated.

Changes

SBI ONNX Exporter Feature

Layer / File(s)	Summary
Dependencies and gitignore pyproject.toml, .gitignore	sbi>=0.26, nflows>=0.14, onnxruntime>=1.17, jaxonnxruntime>=0.3 added to [all] extra and dev deps; sbi-logs/ excluded from version control.
transform_sbi_to_onnx public API and wrapper classes src/lanfactory/onnx/sbi.py, src/lanfactory/onnx/__init__.py	Adds transform_sbi_to_onnx with class-name–based unsupported-estimator rejection, mode-specific capability checks, _NLELogProbWrapper and _NRELogRatioWrapper that split a concatenated rank-1 input and reshape to scalar, and torch.onnx.export tracing. __init__.py re-exports the new function.
ONNX toolchain baseline spike tests tests/test_sbi_spike_mlp_roundtrip.py, tests/test_sbi_spike_maf_roundtrip.py	Two standalone regression tests validate torch→onnxruntime→jaxonnxruntime round-trips for a plain MLP and an nflows MAF, including jaxonnxruntime Constant-node static-arg configuration for MAF graphs.
NLE export validation tests tests/test_sbi_nle_export.py	Trains a tiny NLE_A on a 2D Gaussian toy, exports with mode="nle", asserts three-way numerical and gradient agreement, and validates that unsupported estimator types, missing .log_prob, NLE-in-NRE-mode, and non-positive example dimensions each raise the correct errors.
NRE export validation tests tests/test_sbi_nre_export.py	Trains a small NRE_A classifier on a 2D Gaussian toy, exports with mode="nre", and asserts three-way numerical and gradient agreement; includes log-ratio ordering sanity check.
NRE export with embedding networks tests/test_sbi_embeddings.py	Trains NRE_A classifiers using FCEmbedding and CNNEmbedding on a 10D embedded input, exports with mode="nre", and validates three-way numerical agreement; documents the norm_layer=nn.Identity workaround for ONNX LayerNorm incompatibility.
HSSM end-to-end integration test tests/test_sbi_hssm_integration.py	Trains NLE_A on synthetic DDM data, exports to ONNX, loads into hssm.HSSM with loglik_kind="approx_differentiable", runs short MCMC, and asserts Gelman-Rubin convergence and posterior mean recovery within ±2σ of ground truth.
User documentation and navigation docs/exporting_sbi_models.md, mkdocs.yml, README.md	Adds a full guide covering NLE/NRE quick-start workflows, supported architecture table, known constraints (LayerNorm, JAX x64, dimensionality), numerical guarantees, and float precision notes. Updates mkdocs.yml navigation and README.md with an ONNX exporter summary.

Sequence Diagram(s)

sequenceDiagram
  participant User
  participant transform_sbi_to_onnx
  participant _NLELogProbWrapper
  participant torch_onnx_export as torch.onnx.export
  participant ONNXRuntime
  participant JAXONNXRuntime as jaxonnxruntime
  participant HSSM

  User->>transform_sbi_to_onnx: estimator (NLE_A/NRE_A), path, mode, theta_dim, x_dim
  transform_sbi_to_onnx->>transform_sbi_to_onnx: validate class name + mode capabilities
  transform_sbi_to_onnx->>_NLELogProbWrapper: wrap(estimator, theta_dim, x_dim)
  _NLELogProbWrapper-->>transform_sbi_to_onnx: wrapper.forward(concat_input) → scalar
  transform_sbi_to_onnx->>torch_onnx_export: export(wrapper, dummy rank-1 input, opset=17)
  torch_onnx_export-->>User: .onnx file written to path

  User->>ONNXRuntime: InferenceSession(path).run(concat_input)
  ONNXRuntime-->>User: log_prob scalar
  User->>JAXONNXRuntime: call_onnx_model(onnx_model)(concat_input)
  JAXONNXRuntime-->>User: JAX log_prob + gradients
  User->>HSSM: HSSM(loglik=path, loglik_kind="approx_differentiable")
  HSSM-->>User: posterior samples via MCMC

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Poem

🐰 Hippity-hop, the networks take flight,
Through ONNX portals into JAX's light.
NLE and NRE, wrapped with care,
Scalar log-probs floating through the air.
Three runtimes agree — what a sight!
The bunny exports everything just right. 🎉

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.53% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The PR title clearly summarizes the main addition: introducing a new SBI-to-ONNX exporter for neural likelihood and ratio estimators, which is the primary focus of the changeset.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch sbi-connector

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 5

🧹 Nitpick comments (2)

docs/exporting_sbi_models.md (1)

96-134: 💤 Low value

Known constraints section is thorough and actionable.

The three constraints (≥2D, LayerNorm workaround, JAX x64) are well-motivated and include working code examples:

• Constraint 2 (LayerNorm) is supported by the PR objectives and aligns with the documented sbi-method incompatibility.
• Constraint 3 (JAX x64) is validated by the HSSM integration test context (Context snippet 4), which notes that HSSM's onnx2jax self-manages the jaxort_only_allow_initializers_as_static_args flag but users must opt into x64 themselves.
• Each constraint includes a code snippet or clear reproduction steps.

Minor: Consider adding a cross-reference from "Supported architectures" to this section for readers who may wonder why certain combinations are excluded.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/exporting_sbi_models.md` around lines 96 - 134, The "Known constraints"
section is well-documented, but readers consulting the "Supported architectures"
section may not understand why certain architecture combinations are excluded or
unsupported. Add a cross-reference from the "Supported architectures" section to
the "Known constraints" section, explaining that certain combinations (like 1D θ
or x, or LayerNorm in classifiers) have constraints that limit exportability.
This helps readers connect architectural choices to the documented limitations.

tests/test_sbi_spike_maf_roundtrip.py (1)

38-38: Move config mutation into a pytest fixture to ensure proper cleanup.

Module-level mutations of config.update("jaxort_only_allow_initializers_as_static_args", False) at import time (line 38) leak across test modules and never restore the previous state. This appears in at least 4 test files with no visible restoration logic, creating order-dependent test behavior. Wrap the mutation in a fixture that saves and restores the flag value:

`@pytest.fixture`(scope="function", autouse=True)
def restore_jaxort_config():
    from jaxonnxruntime import config
    prev = config._flags.get("jaxort_only_allow_initializers_as_static_args", True)
    config.update("jaxort_only_allow_initializers_as_static_args", False)
    yield
    config.update("jaxort_only_allow_initializers_as_static_args", prev)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/test_sbi_spike_maf_roundtrip.py` at line 38, The module-level
config.update call at line 38 mutates the
jaxort_only_allow_initializers_as_static_args configuration at import time
without any cleanup or restoration mechanism, causing test state to leak across
test modules. Move this configuration mutation out of module scope and into a
pytest fixture with function scope and autouse enabled, so that the fixture
saves the previous configuration value before updating it, yields control to run
the test, and then restores the original configuration value afterward to
prevent test pollution and order-dependent behavior.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/exporting_sbi_models.md`:
- Around line 101-104: The documentation at the specified location presents the
≥2D constraint ambiguously as a universal requirement when it actually applies
only to NLE with density_estimator="maf", not all density estimators or the
LANfactory exporter itself. Revise the explanation to clarify that this
constraint is specific to sbi's MAF density estimator and is enforced by sbi's
training-time behavior in nflows, not by the transform_sbi_to_onnx exporter.
Additionally, note that other density estimators such as MDN and MoG may not
have this 1D limitation, and that NRE successfully handles 1D inputs, making
clear this is a MAF-specific pre-flight limitation rather than a LANfactory
enforcement.

In `@src/lanfactory/onnx/sbi.py`:
- Around line 47-55: The transform_sbi_to_onnx function accepts
example_theta_dim and example_x_dim parameters but does not validate that they
are positive integers, which can lead to confusing errors later. Add input
validation at the start of the function to check that both example_theta_dim and
example_x_dim are positive integers (greater than zero), and raise a ValueError
with a clear message if either parameter fails validation.

In `@tests/test_sbi_embeddings.py`:
- Around line 79-85: Add a third assertion to complete the three-way pairwise
comparison between all three backends. After the existing assertions comparing
y_torch against y_ort and y_jax, add another assertion that directly compares
y_ort against y_jax using np.allclose with the same atol tolerance parameter.
Follow the same pattern as the existing assertions by including an error message
that displays the maximum absolute difference between y_ort and y_jax to ensure
full pairwise agreement among all three backends.

In `@tests/test_sbi_hssm_integration.py`:
- Around line 59-64: The _build_observed_dataframe function accepts an rng
parameter but never uses it when calling simulator, resulting in uncontrolled
random data generation. Pass the rng parameter to the simulator function call to
ensure the seeded random number generator is actually used for reproducible test
behavior.
- Around line 135-140: The hasattr call on line 135 attempts to access
hssm.utils.summary before verifying that the utils module exists on the hssm
object. If hssm lacks the utils attribute, this raises AttributeError
immediately before the hasattr check can protect against it. Guard the
hssm.utils reference by first checking that hssm has the utils attribute before
checking for the summary attribute within it. Restructure the hasattr condition
to verify hssm.utils exists as a prerequisite before attempting to access
hssm.utils.summary, ensuring the ArviZ fallback executes properly when either
module is missing.

---

Nitpick comments:
In `@docs/exporting_sbi_models.md`:
- Around line 96-134: The "Known constraints" section is well-documented, but
readers consulting the "Supported architectures" section may not understand why
certain architecture combinations are excluded or unsupported. Add a
cross-reference from the "Supported architectures" section to the "Known
constraints" section, explaining that certain combinations (like 1D θ or x, or
LayerNorm in classifiers) have constraints that limit exportability. This helps
readers connect architectural choices to the documented limitations.

In `@tests/test_sbi_spike_maf_roundtrip.py`:
- Line 38: The module-level config.update call at line 38 mutates the
jaxort_only_allow_initializers_as_static_args configuration at import time
without any cleanup or restoration mechanism, causing test state to leak across
test modules. Move this configuration mutation out of module scope and into a
pytest fixture with function scope and autouse enabled, so that the fixture
saves the previous configuration value before updating it, yields control to run
the test, and then restores the original configuration value afterward to
prevent test pollution and order-dependent behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 799ad103-eb67-453a-ae45-e8cf9f0dd877

📥 Commits

Reviewing files that changed from the base of the PR and between 357a1c3 and fe5d806.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (13)

• .gitignore
• README.md
• docs/exporting_sbi_models.md
• mkdocs.yml
• pyproject.toml
• src/lanfactory/onnx/__init__.py
• src/lanfactory/onnx/sbi.py
• tests/test_sbi_embeddings.py
• tests/test_sbi_hssm_integration.py
• tests/test_sbi_nle_export.py
• tests/test_sbi_nre_export.py
• tests/test_sbi_spike_maf_roundtrip.py
• tests/test_sbi_spike_mlp_roundtrip.py

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Add direct onnxruntime↔jaxonnxruntime assertion to complete the three-way check.

Lines 79–85 only anchor each backend against Torch. Add a direct y_ort vs y_jax assertion so the helper enforces full pairwise agreement.

Proposed patch

@@
     assert np.allclose(y_torch, y_ort, atol=atol), (
         f"torch vs onnxruntime: max |Δ| = {np.abs(y_torch - y_ort).max()}"
     )
     assert np.allclose(y_torch, y_jax, atol=atol), (
         f"torch vs jaxonnxruntime: max |Δ| = {np.abs(y_torch - y_jax).max()}"
     )
+    assert np.allclose(y_ort, y_jax, atol=atol), (
+        f"onnxruntime vs jaxonnxruntime: max |Δ| = {np.abs(y_ort - y_jax).max()}"
+    )

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	assert np.allclose(y_torch, y_ort, atol=atol), (
	f"torch vs onnxruntime: max |Δ| = {np.abs(y_torch - y_ort).max()}"
	)
	assert np.allclose(y_torch, y_jax, atol=atol), (
	f"torch vs jaxonnxruntime: max |Δ| = {np.abs(y_torch - y_jax).max()}"
	)
	assert np.allclose(y_torch, y_ort, atol=atol), (
	f"torch vs onnxruntime: max |Δ| = {np.abs(y_torch - y_ort).max()}"
	)
	assert np.allclose(y_torch, y_jax, atol=atol), (
	f"torch vs jaxonnxruntime: max |Δ| = {np.abs(y_torch - y_jax).max()}"
	)
	assert np.allclose(y_ort, y_jax, atol=atol), (
	f"onnxruntime vs jaxonnxruntime: max |Δ| = {np.abs(y_ort - y_jax).max()}"
	)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/test_sbi_embeddings.py` around lines 79 - 85, Add a third assertion to
complete the three-way pairwise comparison between all three backends. After the
existing assertions comparing y_torch against y_ort and y_jax, add another
assertion that directly compares y_ort against y_jax using np.allclose with the
same atol tolerance parameter. Follow the same pattern as the existing
assertions by including an error message that displays the maximum absolute
difference between y_ort and y_jax to ensure full pairwise agreement among all
three backends.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
python - <<'PY'
import importlib
h = importlib.import_module("hssm")
print("has hssm.utils:", hasattr(h, "utils"))
if hasattr(h, "utils"):
    print("has hssm.utils.summary:", hasattr(h.utils, "summary"))
PY

Repository: lnccbrown/LANfactory

Length of output: 601

---

🏁 Script executed:

cat -n tests/test_sbi_hssm_integration.py | sed -n '130,145p'

Repository: lnccbrown/LANfactory

Length of output: 724

---

🏁 Script executed:

# Look for setup files that might specify HSSM version constraints
find . -type f \( -name "setup.py" -o -name "setup.cfg" -o -name "pyproject.toml" -o -name "requirements*.txt" \) | head -20

Repository: lnccbrown/LANfactory

Length of output: 506

---

🏁 Script executed:

# Check if there's documentation about HSSM versions or changelog
find . -type f -name "*requirements*" -o -name "*version*" | grep -E "\.(txt|py|md)$" | head -20

Repository: lnccbrown/LANfactory

Length of output: 489

---

🏁 Script executed:

cat pyproject.toml

Repository: lnccbrown/LANfactory

Length of output: 3901

---

🏁 Script executed:

head -50 tests/test_sbi_hssm_integration.py

Repository: lnccbrown/LANfactory

Length of output: 2082

---

🏁 Script executed:

grep -n "import hssm\|from hssm" tests/test_sbi_hssm_integration.py

Repository: lnccbrown/LANfactory

Length of output: 46

---

🏁 Script executed:

# Check the actual structure around line 135 to understand the context better
sed -n '120,150p' tests/test_sbi_hssm_integration.py

Repository: lnccbrown/LANfactory

Length of output: 1024

---

🌐 Web query:

HSSM Python library utils module version history

💡 Result:

The HSSM (Hierarchical Sequential Sampling Modeling) Python library organizes utility functions across several submodules rather than a single utils module [1][2][3]. The most prominent utility-focused submodule is hssm.distribution_utils, which provides functions for building PyMC distribution classes and handling advanced model construction [2][3]. Other utility functions are found in specialized submodules, such as hssm.utils (e.g., decorate_atomic_simulator) [1] and hssm.rl.utils (introduced in version 0.3.0 for balanced-panel validation in reinforcement learning models) [4]. Regarding the version history, HSSM has undergone continuous development since its initial releases [5][6]. Significant recent updates include: Version 0.3.0: A major release featuring support for choice-only models, the Racing Diffusion Model (RDM3), Poisson Race model, BayesFlow LRE integration, and a new reinforcement learning (RLSSM) configuration system [5][7]. This version also introduced hssm.rl.utils [4]. Version 0.2.12: Introduced compatibility with Python 3.13 and Bambi 0.17.1+, added the make_distribution_for_supported_model utility function, and began private implementations of RLSSM features [5][7]. Version 0.2.x series: These versions focused on incremental improvements, including expanded model support, bug fixes, and documentation updates, building upon the foundations laid in earlier 0.1.x releases [5][6]. For the most up-to-date information on modules and release-specific changes, the official HSSM changelog and API documentation should be consulted [8][5][7].

Citations:

• 1: https://lnccbrown.github.io/HSSM/tutorials/add_custom_rlssm_model/
• 2: https://lnccbrown.github.io/HSSM/api/distribution_utils/
• 3: https://github.com/lnccbrown/HSSM/blob/main/docs/api/distribution_utils.md
• 4: Rlssm class make model dist HSSM#915
• 5: https://lnccbrown.github.io/HSSM/changelog/
• 6: https://pypi.org/project/HSSM/0.2.8/
• 7: https://github.com/lnccbrown/HSSM/blob/main/docs/changelog.md
• 8: https://github.com/lnccbrown/HSSM

---

Guard hssm.utils before dereferencing it.

Line 135 evaluates hssm.utils before the hasattr() check can protect against its absence. If a given HSSM version lacks the utils module, this raises AttributeError immediately and the ArviZ fallback never executes. Add a guard for the module's existence:

Suggested fix

-    summary = hssm.utils.summary(idata) if hasattr(hssm.utils, "summary") else None
+    summary = (
+        hssm.utils.summary(idata)
+        if hasattr(hssm, "utils") and hasattr(hssm.utils, "summary")
+        else None
+    )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/test_sbi_hssm_integration.py` around lines 135 - 140, The hasattr call
on line 135 attempts to access hssm.utils.summary before verifying that the
utils module exists on the hssm object. If hssm lacks the utils attribute, this
raises AttributeError immediately before the hasattr check can protect against
it. Guard the hssm.utils reference by first checking that hssm has the utils
attribute before checking for the summary attribute within it. Restructure the
hasattr condition to verify hssm.utils exists as a prerequisite before
attempting to access hssm.utils.summary, ensuring the ArviZ fallback executes
properly when either module is missing.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
src/lanfactory/onnx/sbi.py	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

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