Integrate a BayesFlow NRE (ONNX)¶

This notebook covers the BayesFlow likelihood-ratio route into HSSM through ONNX. A DDM trial combines a continuous reaction time with a discrete choice, so the supported estimator here is a BayesFlow RatioApproximator (NRE), not a continuous-density NLE. Once trained, the estimator is exported to a single ONNX file and handed to HSSM with the same loglik="file.onnx" gesture used for LAN and sbi artifacts.

Two paths into HSSM exist, side by side:

Path	Source	Mechanism	When to use
loglik="file.onnx"	sbi or BayesFlow	ONNX file, framework-agnostic	Portability, sharing trained surrogates
loglik=<jax_callable>	bayesflow (this tutorial's sibling)	In-memory JAX callable	Fast iteration during model development

See bayesflow_lre_integration.ipynb for the JAX-callable path. This notebook covers the ONNX path.

For the map of all external-trainer routes, see Bring your own likelihood; the exported file's rules are in The ONNX likelihood contract.

Part 1 — Setup¶

Critical: KERAS_BACKEND=torch must be set before importing keras or bayesflow. torch.onnx.export cannot trace a JAX-backed Keras model. The reproducible documentation run keeps both Keras/Torch and JAX on CPU; this avoids unsupported accelerator initialization and machine-specific backend-probe logs.

Note on the two ONNX-related setup lines below (jax_enable_x64 and the jaxonnxruntime strict-mode relax): current HSSM applies both automatically when hssm.distribution_utils.onnx_utils.onnx2jax is imported (added in #964). The explicit lines below are therefore redundant on a current install — but harmless, and they keep the notebook standalone on older HSSM versions.

import logging
import os
import warnings
from contextlib import contextmanager

os.environ["KERAS_BACKEND"] = "torch"
os.environ["KERAS_TORCH_DEVICE"] = "cpu"
os.environ["JAX_PLATFORMS"] = "cpu"
logging.getLogger("pytensor.link.c.cmodule").setLevel(logging.WARNING)
logging.getLogger("pytensor.link.c.basic").setLevel(logging.WARNING)
warnings.filterwarnings(
    "ignore",
    message="Converting a tensor to a Python (integer|boolean).*",
    category=Warning,
)
warnings.filterwarnings(
    "ignore",
    message="There are not enough devices.*",
    category=UserWarning,
)
warnings.filterwarnings(
    "ignore",
    message="`init=.* is ignored by `nuts_sampler=.*",
    category=UserWarning,
)

import jax

# x64 BEFORE other JAX-touching imports; ONNX graphs from torch.onnx.export
# carry int64 shape/index tensors that get silently truncated under JAX's
# default int32 mode, producing wrong log-prob values inside HSSM.
jax.config.update("jax_enable_x64", True)

# Relax jaxonnxruntime's default strict mode on Reshape shape arguments —
# torch.onnx.export emits them as Constant nodes, not initializers, which the
# strict default rejects. Safe because the shapes are genuinely constant.
# HSSM PR #964 sets this automatically inside hssm.distribution_utils.onnx2jax;
# until that PR lands on main, we set it here.
from jaxonnxruntime import config as _jaxonnx_config

_jaxonnx_config.update("jaxort_only_allow_initializers_as_static_args", False)

import tempfile
from pathlib import Path

import arviz as az
import bayesflow as bf
import keras
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from bayesflow.datasets import OfflineDataset
from lanfactory.onnx import transform_bayesflow_to_onnx
from matplotlib.lines import Line2D
from ssms.basic_simulators.simulator import simulator

import hssm


@contextmanager
def suppress_info_logs():
    """Temporarily silence dependency INFO logs that expose local paths."""
    previous_disable = logging.root.manager.disable
    logging.disable(logging.INFO)
    try:
        yield
    finally:
        logging.disable(previous_disable)


print("keras backend:", keras.backend.backend())
print("bayesflow:    ", bf.__version__)
print("hssm:         ", hssm.__version__)

import logging import os import warnings from contextlib import contextmanager os.environ["KERAS_BACKEND"] = "torch" os.environ["KERAS_TORCH_DEVICE"] = "cpu" os.environ["JAX_PLATFORMS"] = "cpu" logging.getLogger("pytensor.link.c.cmodule").setLevel(logging.WARNING) logging.getLogger("pytensor.link.c.basic").setLevel(logging.WARNING) warnings.filterwarnings( "ignore", message="Converting a tensor to a Python (integer|boolean).*", category=Warning, ) warnings.filterwarnings( "ignore", message="There are not enough devices.*", category=UserWarning, ) warnings.filterwarnings( "ignore", message="`init=.* is ignored by `nuts_sampler=.*", category=UserWarning, ) import jax # x64 BEFORE other JAX-touching imports; ONNX graphs from torch.onnx.export # carry int64 shape/index tensors that get silently truncated under JAX's # default int32 mode, producing wrong log-prob values inside HSSM. jax.config.update("jax_enable_x64", True) # Relax jaxonnxruntime's default strict mode on Reshape shape arguments — # torch.onnx.export emits them as Constant nodes, not initializers, which the # strict default rejects. Safe because the shapes are genuinely constant. # HSSM PR #964 sets this automatically inside hssm.distribution_utils.onnx2jax; # until that PR lands on main, we set it here. from jaxonnxruntime import config as _jaxonnx_config _jaxonnx_config.update("jaxort_only_allow_initializers_as_static_args", False) import tempfile from pathlib import Path import arviz as az import bayesflow as bf import keras import matplotlib.pyplot as plt import numpy as np import pandas as pd from bayesflow.datasets import OfflineDataset from lanfactory.onnx import transform_bayesflow_to_onnx from matplotlib.lines import Line2D from ssms.basic_simulators.simulator import simulator import hssm @contextmanager def suppress_info_logs(): """Temporarily silence dependency INFO logs that expose local paths.""" previous_disable = logging.root.manager.disable logging.disable(logging.INFO) try: yield finally: logging.disable(previous_disable) print("keras backend:", keras.backend.backend()) print("bayesflow: ", bf.__version__) print("hssm: ", hssm.__version__)

INFO:bayesflow:Using backend 'torch'

WARNING:bayesflow:
When using torch backend, we need to disable autograd by default to avoid excessive memory usage. Use

with torch.enable_grad():
    ...

in contexts where you need gradients (e.g. custom training loops).

mlflow not available
mlflow not available

INFO:pytensor.configparser:Suppressed KeyError in PyTensorConfigParser.add for parameter 'cxx'!

INFO:pytensor.configparser:Suppressed KeyError in PyTensorConfigParser.add for parameter 'gcc_version_str'!

INFO:pytensor.configparser:Suppressed KeyError in PyTensorConfigParser.add for parameter 'compile__timeout'!

INFO:pytensor.configparser:Suppressed KeyError in PyTensorConfigParser.add for parameter 'DebugMode__check_c'!

INFO:pytensor.configparser:Suppressed KeyError in PyTensorConfigParser.add for parameter 'compiledir'!

INFO:pytensor.configparser:Suppressed KeyError in PyTensorConfigParser.add for parameter 'blas__ldflags'!

keras backend: torch
bayesflow:     2.0.12
hssm:          0.4.0

Part 2 — Simulate observed DDM data¶

We use ssm-simulators for ground-truth DDM samples at a known parameter vector. HSSM consumes a DataFrame with rt and response columns. To keep this an honest integration test rather than an under-calibrated four-parameter analysis, the executable example estimates only drift rate v from 100 observed trials and fixes a, z, and t at their generating values. Production multi-parameter use requires problem-specific simulation-based calibration and a larger validated training budget.

DDM_PARAM_NAMES = ["v", "a", "z", "t"]
DDM_PARAM_LOW = np.array([-2.0, 0.6, 0.3, 0.1], dtype=np.float32)
DDM_PARAM_HIGH = np.array([2.0, 1.8, 0.7, 0.5], dtype=np.float32)
TRUE_THETA = np.array([0.5, 1.2, 0.5, 0.25], dtype=np.float32)
TRUE_THETA_BY_NAME = dict(zip(DDM_PARAM_NAMES, TRUE_THETA))
INFERRED_PARAM_NAMES = ["v"]
FIXED_PARAMS = {name: float(TRUE_THETA_BY_NAME[name]) for name in ("a", "z", "t")}
N_OBS = 100
OBSERVED_DATA_SEED = 11
NRE_SAMPLING_SEED = 101
ANALYTICAL_SAMPLING_SEED = 102
SAMPLING_CHAINS = 4

TRAINING_BOUNDS = {
    name: (float(lower), float(upper))
    for name, lower, upper in zip(DDM_PARAM_NAMES, DDM_PARAM_LOW, DDM_PARAM_HIGH)
}
TRAINING_PRIORS = {
    name: {"name": "Uniform", "lower": lower, "upper": upper}
    for name, (lower, upper) in TRAINING_BOUNDS.items()
}
TRAINING_MODEL_CONFIG = {
    "bounds": TRAINING_BOUNDS,
    "default_priors": TRAINING_PRIORS,
}
SAMPLER_INIT = {"v": 0.5 * sum(TRAINING_BOUNDS["v"])}

out = simulator(
    theta=TRUE_THETA[None, :],
    model="ddm",
    n_samples=N_OBS,
    random_state=OBSERVED_DATA_SEED,
)
obs_data = pd.DataFrame(
    {
        "rt": out["rts"].squeeze().astype(np.float32),
        "response": out["choices"].squeeze().astype(np.float32),
    }
)
obs_data.head()

DDM_PARAM_NAMES = ["v", "a", "z", "t"] DDM_PARAM_LOW = np.array([-2.0, 0.6, 0.3, 0.1], dtype=np.float32) DDM_PARAM_HIGH = np.array([2.0, 1.8, 0.7, 0.5], dtype=np.float32) TRUE_THETA = np.array([0.5, 1.2, 0.5, 0.25], dtype=np.float32) TRUE_THETA_BY_NAME = dict(zip(DDM_PARAM_NAMES, TRUE_THETA)) INFERRED_PARAM_NAMES = ["v"] FIXED_PARAMS = {name: float(TRUE_THETA_BY_NAME[name]) for name in ("a", "z", "t")} N_OBS = 100 OBSERVED_DATA_SEED = 11 NRE_SAMPLING_SEED = 101 ANALYTICAL_SAMPLING_SEED = 102 SAMPLING_CHAINS = 4 TRAINING_BOUNDS = { name: (float(lower), float(upper)) for name, lower, upper in zip(DDM_PARAM_NAMES, DDM_PARAM_LOW, DDM_PARAM_HIGH) } TRAINING_PRIORS = { name: {"name": "Uniform", "lower": lower, "upper": upper} for name, (lower, upper) in TRAINING_BOUNDS.items() } TRAINING_MODEL_CONFIG = { "bounds": TRAINING_BOUNDS, "default_priors": TRAINING_PRIORS, } SAMPLER_INIT = {"v": 0.5 * sum(TRAINING_BOUNDS["v"])} out = simulator( theta=TRUE_THETA[None, :], model="ddm", n_samples=N_OBS, random_state=OBSERVED_DATA_SEED, ) obs_data = pd.DataFrame( { "rt": out["rts"].squeeze().astype(np.float32), "response": out["choices"].squeeze().astype(np.float32), } ) obs_data.head()

	rt	response
0	0.833728	1.0
1	1.076240	1.0
2	0.602407	1.0
3	2.786506	1.0
4	0.597780	1.0

fig, ax = plt.subplots(1, 1, figsize=(8, 4))
for choice, label in [(1.0, "choice +1"), (-1.0, "choice -1")]:
    rts = obs_data.loc[obs_data["response"] == choice, "rt"]
    ax.hist(rts, bins=40, alpha=0.5, label=label)
ax.set_xlabel("reaction time")
ax.set_ylabel("count")
ax.legend()
ax.set_title(f"Observed DDM data at θ={dict(zip(DDM_PARAM_NAMES, TRUE_THETA))}")
plt.tight_layout()

fig, ax = plt.subplots(1, 1, figsize=(8, 4)) for choice, label in [(1.0, "choice +1"), (-1.0, "choice -1")]: rts = obs_data.loc[obs_data["response"] == choice, "rt"] ax.hist(rts, bins=40, alpha=0.5, label=label) ax.set_xlabel("reaction time") ax.set_ylabel("count") ax.legend() ax.set_title(f"Observed DDM data at θ={dict(zip(DDM_PARAM_NAMES, TRUE_THETA))}") plt.tight_layout()

Part 3 — Train a BayesFlow ratio estimator on DDM simulations¶

The MLP setup below follows the ONNX-friendly NRE configuration documented in LANfactory's BayesFlow export guide. Ratio estimation is appropriate for the DDM's mixed continuous/discrete observation because it learns a classifier logit rather than a continuous density over [rt, choice]. We use silu, disable residual connections and dropout, and keep the adapter empty so the trained estimator can be traced into the portable ONNX graph. standardize="all" remains inside the tensor graph and is baked into the export. The NRE-C contrast set uses K=32 candidates with batches of 64, following BayesFlow's recommendation to keep K near half the batch size rather than relying on a very small contrast set.

keras.utils.set_random_seed(0)
rng = np.random.default_rng(0)

# Keep the training count divisible by the batch size: every NRE-C batch
# must contain more samples than K so it can draw K non-self contrasts.
N_TRAIN = 100_032
TRAINING_DATA_SEED = 23
theta_train = rng.uniform(
    DDM_PARAM_LOW, DDM_PARAM_HIGH, size=(N_TRAIN, len(DDM_PARAM_NAMES))
).astype(np.float32)

# One trial per θ_i. NRE convention stores θ as inference variables and
# the mixed [rt, choice] observation as inference conditions.
training_out = simulator(
    theta=theta_train,
    model="ddm",
    n_samples=1,
    random_state=TRAINING_DATA_SEED,
)
x_train = np.column_stack(
    (training_out["rts"].squeeze(), training_out["choices"].squeeze())
).astype(np.float32)

print("theta_train:", theta_train.shape, "  x_train:", x_train.shape)

keras.utils.set_random_seed(0) rng = np.random.default_rng(0) # Keep the training count divisible by the batch size: every NRE-C batch # must contain more samples than K so it can draw K non-self contrasts. N_TRAIN = 100_032 TRAINING_DATA_SEED = 23 theta_train = rng.uniform( DDM_PARAM_LOW, DDM_PARAM_HIGH, size=(N_TRAIN, len(DDM_PARAM_NAMES)) ).astype(np.float32) # One trial per θ_i. NRE convention stores θ as inference variables and # the mixed [rt, choice] observation as inference conditions. training_out = simulator( theta=theta_train, model="ddm", n_samples=1, random_state=TRAINING_DATA_SEED, ) x_train = np.column_stack( (training_out["rts"].squeeze(), training_out["choices"].squeeze()) ).astype(np.float32) print("theta_train:", theta_train.shape, " x_train:", x_train.shape)

theta_train: (100032, 4)   x_train: (100032, 2)

approximator = bf.RatioApproximator(
    inference_network=bf.networks.MLP(
        widths=(256, 256, 256),
        activation="silu",
        residual=False,
        dropout=None,
    ),
    standardize="all",
    K=32,
)
approximator.build(
    {
        "inference_variables": (None, len(DDM_PARAM_NAMES)),
        "inference_conditions": (None, 2),
    }
)
approximator.compile(optimizer=keras.optimizers.Adam(learning_rate=5e-4))

dataset = OfflineDataset(
    data={
        "inference_variables": theta_train,
        "inference_conditions": x_train,
    },
    batch_size=64,
    # MUST be identity for ONNX export; we use only the in-network
    # Standardize layer
    adapter=None,
)
history = approximator.fit(dataset=dataset, epochs=50, verbose=0)
print(f"trained BayesFlow NRE; final loss={history.history['loss'][-1]:.4f}")

approximator = bf.RatioApproximator( inference_network=bf.networks.MLP( widths=(256, 256, 256), activation="silu", residual=False, dropout=None, ), standardize="all", K=32, ) approximator.build( { "inference_variables": (None, len(DDM_PARAM_NAMES)), "inference_conditions": (None, 2), } ) approximator.compile(optimizer=keras.optimizers.Adam(learning_rate=5e-4)) dataset = OfflineDataset( data={ "inference_variables": theta_train, "inference_conditions": x_train, }, batch_size=64, # MUST be identity for ONNX export; we use only the in-network # Standardize layer adapter=None, ) history = approximator.fit(dataset=dataset, epochs=50, verbose=0) print(f"trained BayesFlow NRE; final loss={history.history['loss'][-1]:.4f}")

INFO:bayesflow:Fitting on dataset instance of OfflineDataset.

trained BayesFlow NRE; final loss=2.0777

Part 4 — Export the trained approximator to ONNX¶

One call. The exporter raises clearly if any v1 constraint is violated (wrong KERAS_BACKEND, non-identity adapter, missing inference_network).

Where to write the file. The reproducible demo uses an ephemeral temporary directory, so re-running the notebook does not leave artifacts in your working tree or expose machine-specific paths in its saved output. Replace ARTIFACT_DIR with a project directory when you want to keep the trained ONNX for downstream work.

# Keep the documentation run ephemeral. Replace this with a project directory
# when the exported model should persist after the Python process exits.
ARTIFACT_DIR = Path(tempfile.mkdtemp(prefix="hssm-bayesflow-onnx-"))
onnx_path = ARTIFACT_DIR / "ddm_nre.onnx"

transform_bayesflow_to_onnx(
    approximator,
    str(onnx_path),
    mode="nre",
    example_theta_dim=len(DDM_PARAM_NAMES),
    example_x_dim=2,
)
print(f"wrote {onnx_path.name}  ({onnx_path.stat().st_size:,} bytes)")

# Keep the documentation run ephemeral. Replace this with a project directory # when the exported model should persist after the Python process exits. ARTIFACT_DIR = Path(tempfile.mkdtemp(prefix="hssm-bayesflow-onnx-")) onnx_path = ARTIFACT_DIR / "ddm_nre.onnx" transform_bayesflow_to_onnx( approximator, str(onnx_path), mode="nre", example_theta_dim=len(DDM_PARAM_NAMES), example_x_dim=2, ) print(f"wrote {onnx_path.name} ({onnx_path.stat().st_size:,} bytes)")

wrote ddm_nre.onnx  (539,562 bytes)

Part 5 — Hand the ONNX to HSSM¶

The user gesture is identical to the sbi path and the LAN-MLP path. HSSM detects the .onnx extension, loads it via jaxonnxruntime, vmaps over trials, and wires it into a PyMC Distribution.

A learned likelihood ratio is only validated inside the parameter domain used to train it. TRAINING_MODEL_CONFIG therefore aligns HSSM's bounds and default priors with DDM_PARAM_LOW and DDM_PARAM_HIGH; allowing a free parameter to leave that domain would extrapolate the surrogate and can produce misleading posteriors. This focused integration check estimates v and passes the known generating values for a, z, and t as fixed HSSM parameters. NumPyro starts v at the interior midpoint.

The executable validation uses four chains and fails on any divergence, a maximum split-R̂ above 1.01, or a minimum bulk effective sample size below 400. Four chains make the convergence diagnostic robust enough for this strict threshold; a completed-but-invalid sampling run is never published as a successful tutorial.

with suppress_info_logs():
    model_nre = hssm.HSSM(
        data=obs_data,
        model="ddm",
        model_config=TRAINING_MODEL_CONFIG,
        loglik_kind="approx_differentiable",
        loglik=str(onnx_path),
        p_outlier=0,
        process_initvals=False,
        **FIXED_PARAMS,
    )
model_nre

with suppress_info_logs(): model_nre = hssm.HSSM( data=obs_data, model="ddm", model_config=TRAINING_MODEL_CONFIG, loglik_kind="approx_differentiable", loglik=str(onnx_path), p_outlier=0, process_initvals=False, **FIXED_PARAMS, ) model_nre

Hierarchical Sequential Sampling Model
Model: ddm

Response variable: rt,response
Likelihood: approx_differentiable
Observations: 100

Parameters:

v:
    Prior: Uniform(lower: -2.0, upper: 2.0)
    Explicit bounds: (-2.0, 2.0)

a:
    Prior: 1.2000000476837158
    Explicit bounds: (0.6000000238418579, 1.7999999523162842)

z:
    Prior: 0.5
    Explicit bounds: (0.30000001192092896, 0.699999988079071)

t:
    Prior: 0.25
    Explicit bounds: (0.10000000149011612, 0.5)

idata_nre = model_nre.sample(
    sampler="numpyro",
    draws=1200,
    tune=1000,
    chains=SAMPLING_CHAINS,
    target_accept=0.9,
    initvals=SAMPLER_INIT,
    random_seed=NRE_SAMPLING_SEED,
    progressbar=False,
    mp_ctx="spawn",
)

idata_nre = model_nre.sample( sampler="numpyro", draws=1200, tune=1000, chains=SAMPLING_CHAINS, target_accept=0.9, initvals=SAMPLER_INIT, random_seed=NRE_SAMPLING_SEED, progressbar=False, mp_ctx="spawn", )

INFO:pymc.sampling.mcmc:NUTS[numpyro]: [v]

def validate_sampling_health(idata, summary, label):
    """Fail the executable tutorial on divergent or poorly mixed chains."""
    sample_stats = idata["sample_stats"].ds
    divergence_name = next(
        name for name in ("diverging", "divergences") if name in sample_stats
    )
    divergences = int(sample_stats[divergence_name].sum().item())
    max_r_hat = float(summary["r_hat"].max())
    min_ess_bulk = float(summary["ess_bulk"].min())

    assert divergences == 0, f"{label}: {divergences} divergences"
    assert np.isfinite(max_r_hat) and max_r_hat <= 1.01, (
        f"{label}: max r_hat={max_r_hat:.3f}"
    )
    assert np.isfinite(min_ess_bulk) and min_ess_bulk >= 400, (
        f"{label}: min bulk ESS={min_ess_bulk:.0f}"
    )
    print(
        f"{label}: divergences=0, max r_hat={max_r_hat:.3f}, "
        f"min bulk ESS={min_ess_bulk:.0f}"
    )


summary_nre = az.summary(idata_nre, var_names=INFERRED_PARAM_NAMES)
validate_sampling_health(idata_nre, summary_nre, "BayesFlow NRE")
summary_nre

def validate_sampling_health(idata, summary, label): """Fail the executable tutorial on divergent or poorly mixed chains.""" sample_stats = idata["sample_stats"].ds divergence_name = next( name for name in ("diverging", "divergences") if name in sample_stats ) divergences = int(sample_stats[divergence_name].sum().item()) max_r_hat = float(summary["r_hat"].max()) min_ess_bulk = float(summary["ess_bulk"].min()) assert divergences == 0, f"{label}: {divergences} divergences" assert np.isfinite(max_r_hat) and max_r_hat <= 1.01, ( f"{label}: max r_hat={max_r_hat:.3f}" ) assert np.isfinite(min_ess_bulk) and min_ess_bulk >= 400, ( f"{label}: min bulk ESS={min_ess_bulk:.0f}" ) print( f"{label}: divergences=0, max r_hat={max_r_hat:.3f}, " f"min bulk ESS={min_ess_bulk:.0f}" ) summary_nre = az.summary(idata_nre, var_names=INFERRED_PARAM_NAMES) validate_sampling_health(idata_nre, summary_nre, "BayesFlow NRE") summary_nre

BayesFlow NRE: divergences=0, max r_hat=1.003, min bulk ESS=953

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
v	0.327	0.103	0.19	0.53	952	1112	1.00	0.0035	0.0028

az.plot_trace(idata_nre, var_names=INFERRED_PARAM_NAMES)
plt.tight_layout()

az.plot_trace(idata_nre, var_names=INFERRED_PARAM_NAMES) plt.tight_layout()

Part 6 — Analytical reference posterior via HSSM's DDM¶

DDM has a closed-form likelihood (Navarro & Fuss). We fit the same one-dimensional v problem with the same fixed nuisance parameters and prior. Any posterior drift therefore comes from the neural approximation, not from a different inference domain.

with suppress_info_logs():
    model_analytical = hssm.HSSM(
        data=obs_data,
        model="ddm",
        model_config=TRAINING_MODEL_CONFIG,
        p_outlier=0,
        process_initvals=False,
        **FIXED_PARAMS,
        # default loglik_kind here is the analytical Navarro & Fuss path
    )
idata_analytical = model_analytical.sample(
    sampler="numpyro",
    draws=1000,
    tune=1000,
    chains=SAMPLING_CHAINS,
    target_accept=0.9,
    initvals=SAMPLER_INIT,
    random_seed=ANALYTICAL_SAMPLING_SEED,
    progressbar=False,
    mp_ctx="spawn",
)
summary_analytical = az.summary(idata_analytical, var_names=INFERRED_PARAM_NAMES)
validate_sampling_health(idata_analytical, summary_analytical, "analytical DDM")
posterior_mean_gap = (summary_nre["mean"] - summary_analytical["mean"]).abs()
pooled_posterior_sd = np.hypot(summary_nre["sd"], summary_analytical["sd"])
standardized_mean_gap = posterior_mean_gap / pooled_posterior_sd
max_standardized_mean_gap = float(standardized_mean_gap.max())
comparison_summary = pd.concat(
    {
        "BayesFlow NRE": summary_nre[["mean", "sd"]],
        "analytical DDM": summary_analytical[["mean", "sd"]],
    },
    axis=1,
)
print(comparison_summary)
assert np.isfinite(max_standardized_mean_gap) and max_standardized_mean_gap <= 1.5, (
    "BayesFlow NRE posterior disagrees with the analytical DDM: "
    f"max standardized mean gap={max_standardized_mean_gap:.2f}"
)
print(f"posterior fidelity: max standardized mean gap={max_standardized_mean_gap:.2f}")
summary_analytical

with suppress_info_logs(): model_analytical = hssm.HSSM( data=obs_data, model="ddm", model_config=TRAINING_MODEL_CONFIG, p_outlier=0, process_initvals=False, **FIXED_PARAMS, # default loglik_kind here is the analytical Navarro & Fuss path ) idata_analytical = model_analytical.sample( sampler="numpyro", draws=1000, tune=1000, chains=SAMPLING_CHAINS, target_accept=0.9, initvals=SAMPLER_INIT, random_seed=ANALYTICAL_SAMPLING_SEED, progressbar=False, mp_ctx="spawn", ) summary_analytical = az.summary(idata_analytical, var_names=INFERRED_PARAM_NAMES) validate_sampling_health(idata_analytical, summary_analytical, "analytical DDM") posterior_mean_gap = (summary_nre["mean"] - summary_analytical["mean"]).abs() pooled_posterior_sd = np.hypot(summary_nre["sd"], summary_analytical["sd"]) standardized_mean_gap = posterior_mean_gap / pooled_posterior_sd max_standardized_mean_gap = float(standardized_mean_gap.max()) comparison_summary = pd.concat( { "BayesFlow NRE": summary_nre[["mean", "sd"]], "analytical DDM": summary_analytical[["mean", "sd"]], }, axis=1, ) print(comparison_summary) assert np.isfinite(max_standardized_mean_gap) and max_standardized_mean_gap <= 1.5, ( "BayesFlow NRE posterior disagrees with the analytical DDM: " f"max standardized mean gap={max_standardized_mean_gap:.2f}" ) print(f"posterior fidelity: max standardized mean gap={max_standardized_mean_gap:.2f}") summary_analytical

INFO:pymc.sampling.mcmc:NUTS[numpyro]: [v]

analytical DDM: divergences=0, max r_hat=1.003, min bulk ESS=1495
  BayesFlow NRE           analytical DDM          
           mean        sd           mean        sd
v      0.326728  0.103181       0.445819  0.084809
posterior fidelity: max standardized mean gap=0.89

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
v	0.446	0.085	0.31	0.58	1495	1393	1.00	0.0022	0.0016

Part 7 — Posterior comparison: BayesFlow NRE vs analytical¶

Overlay the v marginals and mark the true value. Both models use the same prior and fixed nuisance parameters, so disagreement reflects the learned likelihood ratio rather than a different inference domain. The executable fidelity check requires the NRE posterior mean to remain within 1.5 pooled posterior standard deviations of its analytical counterpart before this tutorial can be published as successful. This is an integration check, not a substitute for simulation-based calibration of a production estimator.

_dist_visuals = {
    "credible_interval": False,
    "point_estimate": False,
    "point_estimate_text": False,
}
comparison = az.plot_dist(
    idata_analytical,
    var_names=INFERRED_PARAM_NAMES,
    kind="kde",
    visuals={
        **_dist_visuals,
        "dist": {"color": "black", "linestyle": "--"},
    },
)
az.plot_dist(
    idata_nre,
    var_names=INFERRED_PARAM_NAMES,
    kind="kde",
    plot_collection=comparison,
    visuals={**_dist_visuals, "dist": {"color": "tab:blue"}},
)
for name in INFERRED_PARAM_NAMES:
    true_val = TRUE_THETA_BY_NAME[name]
    ax = comparison.get_viz("plot", name)
    ax.axvline(true_val, color="red", lw=1)
    ax.set_title(name)
    ax.legend(
        handles=[
            Line2D([0], [0], color="black", linestyle="--", label="analytical"),
            Line2D([0], [0], color="tab:blue", label="BayesFlow NRE"),
            Line2D([0], [0], color="red", label="truth"),
        ],
        fontsize=8,
    )
comparison

_dist_visuals = { "credible_interval": False, "point_estimate": False, "point_estimate_text": False, } comparison = az.plot_dist( idata_analytical, var_names=INFERRED_PARAM_NAMES, kind="kde", visuals={ **_dist_visuals, "dist": {"color": "black", "linestyle": "--"}, }, ) az.plot_dist( idata_nre, var_names=INFERRED_PARAM_NAMES, kind="kde", plot_collection=comparison, visuals={**_dist_visuals, "dist": {"color": "tab:blue"}}, ) for name in INFERRED_PARAM_NAMES: true_val = TRUE_THETA_BY_NAME[name] ax = comparison.get_viz("plot", name) ax.axvline(true_val, color="red", lw=1) ax.set_title(name) ax.legend( handles=[ Line2D([0], [0], color="black", linestyle="--", label="analytical"), Line2D([0], [0], color="tab:blue", label="BayesFlow NRE"), Line2D([0], [0], color="red", label="truth"), ], fontsize=8, ) comparison

<arviz_plots.plot_collection.PlotCollection at 0x14c1a2490>

Summary¶

• User gesture is the same as sbi or LAN-MLP: hssm.HSSM(loglik="file.onnx", loglik_kind="approx_differentiable"). HSSM doesn't need to know which framework trained the surrogate.
• Inference must remain inside the surrogate's training domain: use priors no wider than the parameter support used to simulate the training set, and inspect sampler diagnostics before interpreting the posterior.
• This executable example validates one free parameter: a, z, and t are fixed at their known generating values so the tutorial can test the BayesFlow → ONNX → HSSM contract without presenting an uncalibrated multi-parameter surrogate as production-ready.
• The v1 constraints on the BayesFlow side (ONNX-friendly MLP, identity adapter, tensor-based standardization, and KERAS_BACKEND=torch at export time) are enforced or documented by lanfactory.onnx.transform_bayesflow_to_onnx. The exporter raises clearly when a constraint is violated.
• For an in-memory JAX-callable alternative (no ONNX file, faster iteration during model development), see bayesflow_lre_integration.ipynb.

Out of v1 scope (tracked as future work):

• Continuous-density NLE export for mixed discrete + continuous observations; use the NRE route shown here instead
• Non-identity bayesflow Adapters (would require either baking the tensor-able subset into the ONNX graph, or shipping the adapter spec alongside the ONNX file)
• Transformer / attention summary networks (LayerNorm + dynamic-shape ops)
• FlowMatching / DiffusionModel / ConsistencyModel inference networks (log_prob requires ODE integration)