fix(base): pass sampler choice through to bambi verbatim

[AlexanderFengler]

Summary

Single-line fix to HSSMBase.sample: pass the user's sampler argument
through to bambi as inference_method=sampler instead of collapsing all
NUTS variants to inference_method="pymc".

The bug

HSSMBase.sample(sampler="numpyro") was silently routing to PyMC NUTS,
not numpyro NUTS. Same for sampler="blackjax" and sampler="nutpie".
Only sampler="pymc" worked as advertised.

Archaeology

• Aug 5, 2024 (commit aef3f9b, "Fix compatibility with Bambi Fix compatibility with Bambi #516"):
  introduced a working pattern — inference_method="mcmc" (generic
  bambi NUTS marker) plus kwargs["nuts_sampler"]="numpyro"/"blackjax"
  injected separately. Bambi's old inference_method="mcmc" was a
  generic switch that read nuts_sampler from kwargs. Correct under
  the old bambi API.
• Dec 17, 2025 (commit 20c100b, "fix: update model.sample api to
  be consistent with bambi's"): bambi renamed its inference_method
  values (mcmc → pymc, nuts_numpyro → numpyro, etc.). This commit
  mechanically updated the string list, but also deleted the
  kwargs["nuts_sampler"] injection block. The flatten-conditional
  was left in place. After this commit, all four NUTS samplers route
  to inference_method="pymc" → nuts_sampler="pymc", with no
  recourse.

The bug has been live since Dec 17, 2025 (~5 months).

The fix

self._inference_obj = self.model.fit(
    inference_method=sampler,   # bambi accepts pymc/numpyro/blackjax/nutpie directly
    ...
)

Bambi natively understands each NUTS variant via
_SUPPORTED_METHODS = {"pymc", "numpyro", "blackjax", "nutpie", "vi", "laplace"}
and routes each to a distinct pm.sample(nuts_sampler=...) call. The
collapse conditional has no purpose under the new API.

Other gates in HSSMBase.sample that read the user's sampler arg
(parallel-sampling warning, init default, jitter handling,
step-sampler check) all check sampler directly, not the post-
normalization inference_method value. None are affected.

Test plan

• Existing HSSM ONNX test suite passes
  (pytest tests/distribution_utils/test_onnx.py tests/distribution_utils/test_onnx_model.py).
• Tests using sampler="numpyro" should be re-run as part of
  review — they were silently exercising PyMC NUTS pre-fix and will
  exercise the actual numpyro path post-fix. Specifically:
  • tests/test_rlssm.py:300
  • tests/test_save_load.py:39
  • tests/slow/test_mcmc.py:100-101
• Verify that sampler="blackjax" and sampler="nutpie" now
  invoke their respective NUTS backends in the user's runtime if
  those backends are installed.

Surfaced by

The sbi NLE tutorial work on the sbi-integration branch — exported
ONNX graphs from sbi-trained classifiers tripped a cloudpickle bug
in PyMC's multiprocess sampler because sampler="numpyro" wasn't
actually invoking numpyro (numpyro doesn't fork processes). Once
fixed, the sbi tutorial runs pm.sample(nuts_sampler="numpyro") and
sidesteps the cloudpickle path entirely.

🤖 Generated with Claude Code

[digicosmos86]

LGTM!