Use the low-level API with PyMC¶

This is a tutorial for advanced users who prefer to use the convenience functions and classes in HSSM to create models and sample in PyMC without using bambi. We assume that the readers of this tutorial are familiar with the internals of PyMC, pytensor, and/or JAX.

In addition to the high-level API that relies on bambi for model creation, HSSM also features a low-level API that it internally calls for creating pytensor Ops and pm.Distributions. Experienced users can use the low-level API directly with PyMC to create even more customized models. This guide shows how to use the low-level API and convenience functions HSSM provides to interface with PyMC directly.

When to drop to the low-level API¶

Three common reasons to work at this level:

1. Build custom PyMC models around HSSM-constructed observation models — to break out of the hierarchical-regression corset imposed by the bambi interface: free-form hierarchies, non-linear transformations on parameters, or time-series processes over parameters.
2. Work with custom models that HSSM does not provide out of the box (and consider contributing them to the ecosystem!).
3. Leave PyMC/HSSM at the tail end to apply your own samplers — see Using compiled log-likelihood functions.

(This framing is adapted from the MathPsych 2025 workshop, archived here.)

Run this tutorial¶

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.

# %pip install hssm

# %pip install hssm

Load Modules¶

import arviz as az
import pymc as pm
from matplotlib import pyplot as plt

import hssm

import arviz as az import pymc as pm from matplotlib import pyplot as plt import hssm

Build a model directly from PyMC¶

With DDM and DDM_SDV which are pm.Distributions, we can create models directly with PyMC. We first simulate some data using ssm_simulators and then use PyMC to specify the model and sample from the it.

# Simulate some data
v_true, a_true, z_true, t_true, sv_true = [0.5, 1.5, 0.5, 0.5, 0.1]
dataset = hssm.simulate_data(
    model="ddm_sdv",
    theta=[v_true, a_true, z_true, t_true, sv_true],
    size=1000,
)

dataset

# Simulate some data v_true, a_true, z_true, t_true, sv_true = [0.5, 1.5, 0.5, 0.5, 0.1] dataset = hssm.simulate_data( model="ddm_sdv", theta=[v_true, a_true, z_true, t_true, sv_true], size=1000, ) dataset

	rt	response
0	1.982656	1.0
1	2.063723	1.0
2	0.766794	1.0
3	3.054746	1.0
4	2.746834	1.0
...	...	...
995	3.696956	1.0
996	1.404581	1.0
997	1.731077	1.0
998	1.145393	1.0
999	8.625175	1.0

1000 rows × 2 columns

# This is a pm.Distribution available in HSSM
# There is also a DDM_SDV class for DDMs with the sv parameter
from hssm.likelihoods import DDM

with pm.Model() as ddm_pymc:
    v = pm.Uniform("v", lower=-10.0, upper=10.0)
    a = pm.HalfNormal("a", sigma=2.0)
    z = pm.Uniform("z", lower=0.01, upper=0.99)
    t = pm.Uniform("t", lower=0.0, upper=0.6, initval=0.1)

    ddm = DDM("ddm", v=v, a=a, z=z, t=t, observed=dataset.values)

    ddm_pymc_trace = pm.sample(mp_ctx="spawn", tune=200, draws=200)

az.plot_trace(ddm_pymc_trace)
plt.tight_layout()

# This is a pm.Distribution available in HSSM # There is also a DDM_SDV class for DDMs with the sv parameter from hssm.likelihoods import DDM with pm.Model() as ddm_pymc: v = pm.Uniform("v", lower=-10.0, upper=10.0) a = pm.HalfNormal("a", sigma=2.0) z = pm.Uniform("z", lower=0.01, upper=0.99) t = pm.Uniform("t", lower=0.0, upper=0.6, initval=0.1) ddm = DDM("ddm", v=v, a=a, z=z, t=t, observed=dataset.values) ddm_pymc_trace = pm.sample(mp_ctx="spawn", tune=200, draws=200) az.plot_trace(ddm_pymc_trace) plt.tight_layout()

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [v, a, z, t]

Sampling 4 chains for 200 tune and 200 draw iterations (800 + 800 draws total) took 9 seconds.

Building top-level distributions with make_distribution_for_supported_model¶

Note: this feature was added in HSSM 0.2.12.

HSSM supports building a pm.Distribution class for any supported model with the make_distribution_for_supported_model convenience function. All you need is a string specifying which model you would like and the kind of likelihood.

from hssm.distribution_utils import make_distribution_for_supported_model

ANGLE = make_distribution_for_supported_model(
    "angle", loglik_kind="approx_differentiable"
)

with pm.Model() as angle_pymc:
    v = pm.Uniform("v", lower=-10.0, upper=10.0)
    a = pm.HalfNormal("a", sigma=2.0)
    z = pm.Uniform("z", lower=0.01, upper=0.99)
    t = pm.Uniform("t", lower=0.0, upper=0.6, initval=0.1)
    theta = pm.Uniform("theta", lower=0, upper=0.5)

    angle = ANGLE(
        "angle",
        v=v,
        a=a,
        z=z,
        t=t,
        theta=theta,
        observed=dataset.values,
    )

    angle_pymc_trace = pm.sample(mp_ctx="spawn", tune=200, draws=200)

from hssm.distribution_utils import make_distribution_for_supported_model ANGLE = make_distribution_for_supported_model( "angle", loglik_kind="approx_differentiable" ) with pm.Model() as angle_pymc: v = pm.Uniform("v", lower=-10.0, upper=10.0) a = pm.HalfNormal("a", sigma=2.0) z = pm.Uniform("z", lower=0.01, upper=0.99) t = pm.Uniform("t", lower=0.0, upper=0.6, initval=0.1) theta = pm.Uniform("theta", lower=0, upper=0.5) angle = ANGLE( "angle", v=v, a=a, z=z, t=t, theta=theta, observed=dataset.values, ) angle_pymc_trace = pm.sample(mp_ctx="spawn", tune=200, draws=200)

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [v, a, z, t, theta]

Sampling 4 chains for 200 tune and 200 draw iterations (800 + 800 draws total) took 36 seconds.

The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details

The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

Building custom top-level distributions with make_distribution¶

Note: This tutorial has undergone major updates in HSSM 0.2.1+ following breaking changes in the distribution_utils api. Please follow this tutorial closely if your previous code no longer works.

What a random variable minimally needs¶

In probabilistic programming, the minimal specification of a random variable that lets you sample every quantity of interest in a standard Bayesian analysis is:

• a likelihood function, and
• a valid simulator.

With both in hand you can sample from the posterior, the prior predictive, and the posterior predictive — essentially everything. You can run MCMC without ever supplying a valid simulator, but you are then on your own for prior/posterior predictive sampling. This is why make_distribution asks for an rv (the simulator side) alongside loglik (the likelihood side) — and why an unknown rv string only warns at build time but errors the moment you sample from the random variable.

make_distribution¶

The above example shows that, as long as the top-level distribution is known, modeling can be done in PyMC as well without using Bambi. However, as this official PyMC tutorial shows, creating a distribution in PyMC can be a time consuming-task. You will need to create a RandomVariable first and then define your custom Distribution by extending pm.Distribution class. From PyMC 5.0.0 on, pm.CustomDist simplifies this process, but the use case is not applicable to complex likelihoods in HSSM. Therefore, HSSM provides convenience functions in its distribution_utils submodule that make this process easy. Next, we use another example to show how we can use these functions to create custom pm.Distributions to be used with PyMC.

Use Case: Suppose we have a likelihood function for DDM models with standard deviations for v written. This model has 5 parameters: v, a, z, t, sv, and we want to use this function as the likelihood to create a pm.Distribution for modeling with PyMC. We can use make_distribution for this purpose.

Note: This distribution is already available in HSSM at hssm.likelihoods.DDM_SDV. For illustration purposes, we go through the same process in which this distribution is created. We can use the same procedure for other distributions not currently available in HSSM.

# This is a likelihood function for the DDM with SDV
# Different from DDM which we imported in the previous example,
# which is a pm.Distribution
from hssm.distribution_utils import make_distribution
from hssm.likelihoods import logp_ddm_sdv

# We use `make_distribution` to wrap the likelihood function into a pm.Distribution
DDM_SDV = make_distribution(
    rv="ddm_sdv",
    loglik=logp_ddm_sdv,
    list_params=["v", "a", "z", "t", "sv"],
    bounds={"t": (0, 1)},
)

with pm.Model() as ddm_sdv_model:
    v = pm.Uniform("v", lower=-10.0, upper=10.0)
    a = pm.HalfNormal("a", sigma=2.0)
    z = pm.Uniform("z", lower=0.01, upper=0.99)
    t = pm.Uniform("t", lower=0.0, upper=0.6, initval=0.1)
    sv = pm.HalfNormal("sv", sigma=2.0)

    ddm = DDM_SDV("ddm", v=v, a=a, z=z, t=t, sv=sv, observed=dataset.values)

    ddm_sdv_trace = pm.sample(mp_ctx="spawn", tune=200, draws=200)

az.plot_trace(ddm_sdv_trace)
plt.tight_layout()

# This is a likelihood function for the DDM with SDV # Different from DDM which we imported in the previous example, # which is a pm.Distribution from hssm.distribution_utils import make_distribution from hssm.likelihoods import logp_ddm_sdv # We use `make_distribution` to wrap the likelihood function into a pm.Distribution DDM_SDV = make_distribution( rv="ddm_sdv", loglik=logp_ddm_sdv, list_params=["v", "a", "z", "t", "sv"], bounds={"t": (0, 1)}, ) with pm.Model() as ddm_sdv_model: v = pm.Uniform("v", lower=-10.0, upper=10.0) a = pm.HalfNormal("a", sigma=2.0) z = pm.Uniform("z", lower=0.01, upper=0.99) t = pm.Uniform("t", lower=0.0, upper=0.6, initval=0.1) sv = pm.HalfNormal("sv", sigma=2.0) ddm = DDM_SDV("ddm", v=v, a=a, z=z, t=t, sv=sv, observed=dataset.values) ddm_sdv_trace = pm.sample(mp_ctx="spawn", tune=200, draws=200) az.plot_trace(ddm_sdv_trace) plt.tight_layout()

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [v, a, z, t, sv]

Sampling 4 chains for 200 tune and 200 draw iterations (800 + 800 draws total) took 12 seconds.

The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details

The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

The above example shows how to use the make_distribution factory function to create a pm.Distribution class that can be used for modeling in PyMC in the code above. All you need is a likelihood function that can be directly used with pytensor. This mostly is a Python function written in pytensor, as is the case with logp_ddm_sdv. If you have a likelihood function written with JAX or Cython, or if you have an approximate differentiable likelihood saved in an onnx file, you need to wrap this likelihood in a pytensor Op before using the make_distribution function. The next section describes how to do so. Please refer to the next section to create a pytensor-compatible likelihood function before coming back to this section to create a PyMC Distribution class.

Detailed usage of the make_distribution factory function is as follows:

• rv: a str or a RandomVariable. If a str is provided, a RandomVariable class will be created automatically. This RandomVariable will use the str to identify a simulator provided in the ssm_simulators package as its rng_fn (sampling) function. If this str is not one of the entries to the model_config dict specified here, then the Distribution will still be created but with a warning that any attempt to sample from the RandomVariable will result in an Error. That includes sampling from the posterior distribution. The user could create his/her own RandomVariable class and define its rng_fn class method for sampling.

• loglik: an instance of a pytensor Op object or a Python Callable. This is where the likelihood function is passed in. make_distribution assumes that the likelihood function is directly usable as part of a pytensor graph, such as a function written in pytensor or an instance of a pytensor Op object. Otherwise, please refer to the next section to create a pytensor-compatible likelihood function.

  The signature for the likelihood also has to follow a specific pattern. Please refer to this section for more details.

• list_params: a list of strs specifying the parameters used in the model. Note the order in which the parameters are defined in this list is very important! In principle it has to be consistent with the order in which parameters are passed to the Op or Callable specified with loglik.

• bounds: a dict of param : (lower, upper) specifying the range of values that each parameter can take. HSSM currently does not support specifying inf or -inf as bounds but will in an update in the near future.

Once the distribution is created, you can use it with PyMC just as the DDM class created above.

Creating pytensor-compatible likelihood functions with make_likelihood_callable¶

Update notice: Up until HSSM 0.2.0, there used to be a make_distribution_from_onnx function for creating pm.Distributions directly from onnx files. That function is now deprecated in favor of a more general and customizable process described below.

Sometimes users might have written a custom likelihood in JAX, Cython, or numba that they would like to incorporate into their PyMC modeling workflows. They might also want to incorporate a neural network as an approximate likelihood function. In all of these cases, they can use make_likelihood_callable to wrap these non-pytensor-compatible likelihoods into a pytensor Op, so that it can be used in the sampling process. The full map of routes from a likelihood to a PyMC-ready distribution:

1. Any valid Python function (Cython, numba, pure Python, or a jitted JAX function treated as opaque) → the blackbox regime, non-differentiable.
2. An .onnx file with a network signature → approx_differentiable.
3. A JAX function with a network signature → approx_differentiable — see Custom models from JAX callables.
4. A pytensor function with a likelihood signature → analytical.

For the conceptual background on these likelihood kinds, see Understanding likelihood functions in HSSM. We now describe use cases 1 and 2:

1. I have a function written in JAX, Cython, numba, or even pure Python:

In this case, we assume that these are "black box" likelihoods that are not differentiable. A JAX function passed this way is treated as an opaque black box (route 1 above). To use a differentiable JAX network forward function instead, pass it via the approx_differentiable route (route 3).

from hssm.distribution_utils import make_likelihood_callable

# Suppose you have a non-differentiable function called my_awesome_likelihood_func
loglik_op = make_likelihood_callable(
    loglik=my_awesome_likelihood_func,
    loglik_kind="blackbox",
)

MY_DIST = make_distribution(
    loglik=loglik_op,
    ...
)

What happens under the hood is that make_likelihood_callable will construct an Op and wrap this Op around the function. Please note that this does NOT make the function differentiable. When using the MY_DIST distribution created above, please use a sampler that does not make use of the gradients such as the slice sampler.

2. I have a neural network saved in an onnx file that I want to use as the likelihood function.

This use case is very similar to the one above. All you need to do is to provide a str or a Path to the onnx file. make_likelihood_callable will look for the file locally first. If the file is not available locally, it will also check the HSSM hugging face model repository to look for the model.

HSSM translates the onnx file into a JAX or pytensor function depending on the backend you specify (details in Understanding likelihood functions in HSSM; the rules the onnx file itself must satisfy are in The ONNX likelihood contract).

params_is_reg semantics. With the jax backend, also pass params_is_reg: a list of bools, one per parameter in list_params order, marking whether that parameter is trial-wise (a vector — e.g. the target of a regression) or a scalar. This tells JAX how to vectorize the computation. When a parameter is marked True, your PyMC model must supply a trial-length vector for it (the param * pt.ones(n_trials) idiom). For a fully worked trial-wise regression built this way, see the archived MathPsych 2025 example.

from hssm.distribution_utils import (
    make_distribution,
    make_likelihood_callable,
)

loglik_op = make_likelihood_callable(
    loglik="ddm.onnx",  # will be downloaded from huggingface
    loglik_kind="approx_differentiable",
    backend="jax",  # the onnx will be translated to JAX
    params_is_reg=[False] * 4,  # required if backend is JAX.
    # Since the model below has no regression, we provide a list of 4 `False`s
)

DDM_JAX = make_distribution(
    rv="ddm",
    loglik=loglik_op,
    list_params=["v", "a", "z", "t"],
    bounds={"t": (0, 2), "a": (0, 2.5)},
)

with pm.Model() as ddm_jax_model:
    v = pm.Uniform("v", lower=-3, upper=3)
    a = pm.Gamma("a", mu=0.5, sigma=1.0)
    z = pm.Uniform("z", lower=0.1, upper=0.9)
    t = pm.Uniform("t", lower=0.01, upper=1.0, initval=0.1)

    ddm = DDM_JAX("ddm", v=v, a=a, z=z, t=t, observed=dataset.values)

    ddm_jax_trace = pm.sample(mp_ctx="spawn", tune=500, draws=200)

az.plot_trace(ddm_jax_trace)
plt.tight_layout()

from hssm.distribution_utils import ( make_distribution, make_likelihood_callable, ) loglik_op = make_likelihood_callable( loglik="ddm.onnx", # will be downloaded from huggingface loglik_kind="approx_differentiable", backend="jax", # the onnx will be translated to JAX params_is_reg=[False] * 4, # required if backend is JAX. # Since the model below has no regression, we provide a list of 4 `False`s ) DDM_JAX = make_distribution( rv="ddm", loglik=loglik_op, list_params=["v", "a", "z", "t"], bounds={"t": (0, 2), "a": (0, 2.5)}, ) with pm.Model() as ddm_jax_model: v = pm.Uniform("v", lower=-3, upper=3) a = pm.Gamma("a", mu=0.5, sigma=1.0) z = pm.Uniform("z", lower=0.1, upper=0.9) t = pm.Uniform("t", lower=0.01, upper=1.0, initval=0.1) ddm = DDM_JAX("ddm", v=v, a=a, z=z, t=t, observed=dataset.values) ddm_jax_trace = pm.sample(mp_ctx="spawn", tune=500, draws=200) az.plot_trace(ddm_jax_trace) plt.tight_layout()

<environment>/site-packages/pytensor/gradient.py:1327: FutureWarning: LANLogpOp should implement `pullback` instead of `L_op`/`grad`. Direct `L_op`/`grad` implementations are deprecated and will stop being called in a future version.
  input_grads = node.op.pullback(inputs, node.outputs, new_output_grads)
Initializing NUTS using jitter+adapt_diag...

<environment>/site-packages/pytensor/gradient.py:1327: FutureWarning: LANLogpOp should implement `pullback` instead of `L_op`/`grad`. Direct `L_op`/`grad` implementations are deprecated and will stop being called in a future version.
  input_grads = node.op.pullback(inputs, node.outputs, new_output_grads)

<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(
<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpVJPOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [v, a, z, t]

<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(
<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpVJPOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(
<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpVJPOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(
<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpVJPOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(
<environment>/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpVJPOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

Sampling 4 chains for 500 tune and 200 draw iterations (2_000 + 800 draws total) took 25 seconds.