The ONNX likelihood contract¶

This page is the canonical statement of the rules an ONNX file must follow to be used as an HSSM likelihood via loglik_kind="approx_differentiable" — whether it was trained with LANfactory, BayesFlow, sbi, or anything else that can emit ONNX.

It is also runnable: every rule below is demonstrated live against a real network artifact, so if HSSM's enforcement or the contract itself ever changes, this page fails to execute rather than silently going stale.

The contract in two sentences. Export one per-trial forward pass with every input dimension concrete — no dynamic_axes, no symbolic shapes. HSSM batches across trials itself, by wrapping your graph in jax.vmap.

Run this how-to¶

On Colab, uncomment and run the installation cell below once, then restart the runtime. For local setup, GPU extras, and troubleshooting see the Installation guide.

In [1]:

Copied!

# %pip install hssm
# %pip install hssm

What HSSM expects of the graph¶

Input: a single flat per-trial vector containing the model parameters first, then the data columns — shape (n_params + n_data_cols,) or (1, n_params + n_data_cols) (see Rank below). For a DDM-family LAN that means (v, a, z, t, rt, choice).
Output: the log-likelihood of that single trial — a scalar, or any shape that squeezes to one ((), (1,), (1, 1)).
Every dimension concrete.

Let's inspect a real, contract-compliant artifact — the 3-choice race LAN shipped with these docs (6 parameters + 2 data columns):

In [2]:

Copied!





import onnx

GOOD = "../tutorials/data/race_3_no_bias_lan_no_batch.onnx"

graph = onnx.load(GOOD).graph
inp, out = graph.input[0], graph.output[0]


def dims(value_info):
    """Return each dimension as its concrete value or its symbolic name."""
    return [
        d.dim_value if d.dim_value else d.dim_param
        for d in value_info.type.tensor_type.shape.dim
    ]


print(f"input  {inp.name!r}: dims={dims(inp)}")
print(f"output {out.name!r}: dims={dims(out)}")
print("ops:", sorted({n.op_type for n in graph.node}))
import onnx

GOOD = "../tutorials/data/race_3_no_bias_lan_no_batch.onnx"

graph = onnx.load(GOOD).graph
inp, out = graph.input[0], graph.output[0]


def dims(value_info):
    """Return each dimension as its concrete value or its symbolic name."""
    return [
        d.dim_value if d.dim_value else d.dim_param
        for d in value_info.type.tensor_type.shape.dim
    ]


print(f"input  {inp.name!r}: dims={dims(inp)}")
print(f"output {out.name!r}: dims={dims(out)}")
print("ops:", sorted({n.op_type for n in graph.node}))

input  'onnx::Gemm_0': dims=[1, 8]
output '19': dims=[1, 1]
ops: ['Gemm', 'Tanh']

Every dimension is a concrete integer — (1, 8) in, (1, 1) out. That is the whole invariant.

Why: the silent-corruption failure mode¶

HSSM converts ONNX to JAX with jaxonnxruntime, which traces your graph against its construction-time input shape and bakes the resulting shapes into the returned closure. A graph exported with a dynamic batch axis does not fail loudly when called at a different batch size — it silently returns wrong numbers for any model with a batch-dependent intermediate (a log-det accumulator in a normalizing flow, a Reshape whose -1 resolves against the batch dimension).

Single-trial export plus HSSM-side jax.vmap is mathematically equivalent — the likelihood is per-trial — has zero JIT overhead after XLA fusion, and makes the failure mode impossible. This is why the contract is enforced at load time instead of documented as a recommendation. Watch it happen: we give the same network a symbolic batch dimension and ask HSSM to load it.

In [3]:

Copied!





import os
import tempfile

from hssm.distribution_utils import make_likelihood_callable

BAD = os.path.join(tempfile.gettempdir(), "race3_dynamic.onnx")
bad_model = onnx.load(GOOD)
bad_model.graph.input[0].type.tensor_type.shape.dim[0].dim_param = "batch"
onnx.save(bad_model, BAD)

try:
    make_likelihood_callable(
        loglik=BAD,
        loglik_kind="approx_differentiable",
        backend="jax",
        params_is_reg=[False] * 6,
    )
except ValueError as err:
    print(f"ValueError: {err}")
import os
import tempfile

from hssm.distribution_utils import make_likelihood_callable

BAD = os.path.join(tempfile.gettempdir(), "race3_dynamic.onnx")
bad_model = onnx.load(GOOD)
bad_model.graph.input[0].type.tensor_type.shape.dim[0].dim_param = "batch"
onnx.save(bad_model, BAD)

try:
    make_likelihood_callable(
        loglik=BAD,
        loglik_kind="approx_differentiable",
        backend="jax",
        params_is_reg=[False] * 6,
    )
except ValueError as err:
    print(f"ValueError: {err}")

ValueError: ONNX model has dynamic (symbolic) input dimensions: onnx::Gemm_0[batch]. HSSM uses single-trial input shapes and vmaps over trials at a layer above this conversion -- re-export the model with a concrete per-trial input shape (omit `dynamic_axes` in `torch.onnx.export`, or pass a single rank-1 dummy as LANfactory.onnx.transform_sbi_to_onnx does). Dynamic dims here would cause jaxonnxruntime to silently produce wrong outputs for graphs with batch-dependent intermediates (e.g. log-det accumulators).

The compliant file, by contrast, loads into a ready-to-use pytensor Op:

In [4]:

Copied!





loglik_op = make_likelihood_callable(
    loglik=GOOD,
    loglik_kind="approx_differentiable",
    backend="jax",
    params_is_reg=[False] * 6,
)
type(loglik_op).__name__
loglik_op = make_likelihood_callable(
    loglik=GOOD,
    loglik_kind="approx_differentiable",
    backend="jax",
    params_is_reg=[False] * 6,
)
type(loglik_op).__name__

Out[4]:

'LANLogpOp'

Rank: what is and is not required¶

The invariant is concrete dimensions. Rank is not part of the contract — it follows from how your tracer lowers a dense layer, and the ecosystem legitimately contains both forms:

exporter	traced dummy	lowering
LANfactory `transform_onnx.py` (LAN/CPN/OPN, torch)	`(1, D)`	`Gemm`
LANfactory `jax_export.py` (LAN/CPN/OPN, jax2onnx)	`(1, D)`	`Gemm`
LANfactory `sbi.py`	`(D,)`	`MatMul` + `Add`
LANfactory `bayesflow.py`	`(D,)`	`MatMul` + `Add`

We can reproduce both rows with a toy network — same model, two dummies:

In [5]:

Copied!





import warnings

import torch

net = torch.nn.Sequential(
    torch.nn.Linear(8, 16), torch.nn.Tanh(), torch.nn.Linear(16, 1)
)

# dynamo=False selects the TorchScript exporter; the newer dynamo exporter
# also emits concrete dims by default but requires the onnxscript package
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    for dummy, name in [(torch.zeros(8), "rank-1"), (torch.zeros(1, 8), "rank-2")]:
        path = os.path.join(tempfile.gettempdir(), f"tiny_{name}.onnx")
        torch.onnx.export(net, (dummy,), path, dynamo=False)
        g = onnx.load(path).graph
        ops = sorted({n.op_type for n in g.node})
        # both forms must load in HSSM — make the claim executable
        make_likelihood_callable(
            loglik=path,
            loglik_kind="approx_differentiable",
            backend="jax",
            params_is_reg=[False] * 6,
        )
        print(
            f"{name} dummy -> input dims {dims(g.input[0])}, ops {ops} -> loads in HSSM"
        )
import warnings

import torch

net = torch.nn.Sequential(
    torch.nn.Linear(8, 16), torch.nn.Tanh(), torch.nn.Linear(16, 1)
)

# dynamo=False selects the TorchScript exporter; the newer dynamo exporter
# also emits concrete dims by default but requires the onnxscript package
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    for dummy, name in [(torch.zeros(8), "rank-1"), (torch.zeros(1, 8), "rank-2")]:
        path = os.path.join(tempfile.gettempdir(), f"tiny_{name}.onnx")
        torch.onnx.export(net, (dummy,), path, dynamo=False)
        g = onnx.load(path).graph
        ops = sorted({n.op_type for n in g.node})
        # both forms must load in HSSM — make the claim executable
        make_likelihood_callable(
            loglik=path,
            loglik_kind="approx_differentiable",
            backend="jax",
            params_is_reg=[False] * 6,
        )
        print(
            f"{name} dummy -> input dims {dims(g.input[0])}, ops {ops} -> loads in HSSM"
        )

rank-1 dummy -> input dims [8], ops ['Add', 'MatMul', 'Tanh'] -> loads in HSSM
rank-2 dummy -> input dims [1, 8], ops ['Gemm', 'Tanh'] -> loads in HSSM

torch.onnx.export lowers Linear to rank-agnostic MatMul+Add from a rank-1 dummy, and to Gemm (whose ONNX spec requires rank 2) from a (1, D) dummy. Both forms load in HSSM and, measured under vmap+jit, run identically. All production networks on franklab/HSSM are (1, D) Gemm.

One real rank constraint: flow-based graphs. Graphs that internally slice their input (sbi and BayesFlow flow exports split the input into theta and x) must be traced with a rank-1 dummy: a (1, D) trace emits Slice ops with axes=[1] that fail under HSSM's vmap. Pure feed-forward MLPs (LANs) are fine at either rank. When in doubt, trace rank-1 if torch.onnx.export is your tracer.

Errors you will see¶

The dynamic-dims ValueError demonstrated above — re-export without dynamic_axes.
ValueError: ... int64 constant outside the int32 range ... — you are running hssm.set_floatX("float32") with a flow-based export. Flow graphs carry an INT64_MAX open-ended-slice sentinel that JAX truncates to -1 when x64 is off, which would corrupt the likelihood; HSSM raises instead. Use the default float64 setting for flow-based ONNX likelihoods.

The end-to-end check: HSSM smoke-load¶

The final verification for any artifact: build a real hssm.HSSM model with it and confirm the initial log-probability is finite — no sampling required.

In [6]:

Copied!





import math
import warnings

import hssm

data = hssm.simulate_data(
    model="race_no_bias_3", theta=[0.8, 0.6, 0.4, 2.0, 0.5, 0.3], size=200
)

model = hssm.HSSM(
    data=data,
    model="race_no_bias_3",
    model_config={
        "list_params": ["v0", "v1", "v2", "a", "z", "t"],
        "bounds": {
            "v0": (0.0, 2.5),
            "v1": (0.0, 2.5),
            "v2": (0.0, 2.5),
            "a": (1.0, 3.0),
            "z": (0.0, 0.9),
            "t": (0.001, 2),
        },
        "backend": "jax",
    },
    loglik=GOOD,
    loglik_kind="approx_differentiable",
    choices=[0, 1, 2],
    z=0.5,
    p_outlier=0,
)

# the numba object-mode UserWarning is expected and noisy — silence it narrowly
with warnings.catch_warnings():
    warnings.simplefilter("ignore", UserWarning)
    initial_logp = float(
        model.pymc_model.compile_logp()(model.pymc_model.initial_point())
    )
print(f"initial logp: {initial_logp:.2f}")
assert math.isfinite(initial_logp), "initial log-probability must be finite"
import math
import warnings

import hssm

data = hssm.simulate_data(
    model="race_no_bias_3", theta=[0.8, 0.6, 0.4, 2.0, 0.5, 0.3], size=200
)

model = hssm.HSSM(
    data=data,
    model="race_no_bias_3",
    model_config={
        "list_params": ["v0", "v1", "v2", "a", "z", "t"],
        "bounds": {
            "v0": (0.0, 2.5),
            "v1": (0.0, 2.5),
            "v2": (0.0, 2.5),
            "a": (1.0, 3.0),
            "z": (0.0, 0.9),
            "t": (0.001, 2),
        },
        "backend": "jax",
    },
    loglik=GOOD,
    loglik_kind="approx_differentiable",
    choices=[0, 1, 2],
    z=0.5,
    p_outlier=0,
)

# the numba object-mode UserWarning is expected and noisy — silence it narrowly
with warnings.catch_warnings():
    warnings.simplefilter("ignore", UserWarning)
    initial_logp = float(
        model.pymc_model.compile_logp()(model.pymc_model.initial_point())
    )
print(f"initial logp: {initial_logp:.2f}")
assert math.isfinite(initial_logp), "initial log-probability must be finite"

Model initialized successfully.

initial logp: -10229.39

Exporting correctly¶

With raw torch.onnx.export: pass a single-trial dummy and omit dynamic_axes entirely, exactly as in the toy example above. If you trained with LANfactory, BayesFlow, or sbi, prefer LANfactory's exporters (transform_onnx, transform-jax-onnx, transform_sbi_to_onnx, and the BayesFlow export) — they follow the contract by construction. LANfactory's development version also ships an executable checker to call from your exporter's tests instead of re-deriving the rules:

from lanfactory.onnx.contract import assert_single_trial_contract

assert_single_trial_contract("my_model.onnx", expected_input_width=8)

Not to be confused with the blackbox route¶

The tutorial Custom models from ONNX files (blackbox) shows a trick that rewrites input dimensions to dynamic — the same surgery we used above to break the contract — in order to enable batched onnxruntime inference. That applies only to loglik_kind="blackbox", where the ONNX file is executed by onnxruntime inside an ordinary Python function and HSSM never converts the graph. The two routes are opposites on this point; pick the route first, then shape the file.

Verification checklist¶

Before publishing or sharing an ONNX likelihood:

Every input dimension is concrete (inspect as in the first cell above).
Round-trip parity: onnxruntime output matches your source model on ~1000 in-bounds parameter draws (atol=1e-4).
HSSM smoke-load passes (as demonstrated above): finite initial log-probability.
The input column order matches your list_params + data columns — HSSM passes parameters in list_params order, then (rt, response).