Construct Custom Models from simulators and JAX callables¶

In [1]:

import os
import pickle

import arviz as az
import jax.numpy as jnp
import lanfactory as lf
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm

import hssm
from hssm.config import ModelConfig

import os import pickle import arviz as az import jax.numpy as jnp import lanfactory as lf import matplotlib.pyplot as plt import numpy as np import pymc as pm import hssm from hssm.config import ModelConfig

wandb not available
wandb not available

Simulate Data¶

As a pre-amble, we will simulate a simple dataset from the DDM model to use through the example.

In [2]:

# simulate some data from the model
SEED = 123
obs_ddm = hssm.simulate_data(
    theta=dict(v=0.40,
               a=1.25,
               t=0.2,
               z=0.5),
    model="ddm",
    size=500,
    random_state = SEED,
)

# simulate some data from the model SEED = 123 obs_ddm = hssm.simulate_data( theta=dict(v=0.40, a=1.25, t=0.2, z=0.5), model="ddm", size=500, random_state = SEED, )

Construct PyMC distribution from simulator and JAX callable¶

Create JAX Log-likelihood Callable¶

What we need is a callable that takes in a matrix that stacks model parameters and data trialwise, hence with input dimensions: $trials \times (dim-parameters + dim-data)$. The function should return a vector of length $trials$ that contains the log-likelihood for each trial.

If we have a JAX function with this signature, we will be able to proceed to create a valid PyMC distribution form helper functions provided by hssm.

In this particular example, we will reinstantiate a pretrained LAN via utilities from the lanfactory package. This is not necessary, but illustrates how to make use of the braoder ecosystem around HSSM.

Note: lanfactory is an optional dependency of hssm and can be installed using the command pip install hssm[notebook]. Alternatively, you can install it directly with pip install lanfactory.

In [3]:

# Loaded Net
jax_infer = lf.trainers.MLPJaxFactory(
    network_config=pickle.load(
        open(
            os.path.join(
                "data", "jax_models", "ddm", "ddm_jax_lan_network_config.pickle"
            ),
            "rb",
        )
    ),
    train=False,
)

jax_mlp_forward, _ = jax_infer.make_forward_partial(
    seed=42,
    input_dim=4 + 2,  # n-parameters (v,a,z,t) + n-data (rts and choices)
    state=os.path.join("data", "jax_models", "ddm", "ddm_jax_lan_train_state.jax"),
    add_jitted=True,
)

# Loaded Net jax_infer = lf.trainers.MLPJaxFactory( network_config=pickle.load( open( os.path.join( "data", "jax_models", "ddm", "ddm_jax_lan_network_config.pickle" ), "rb", ) ), train=False, ) jax_mlp_forward, _ = jax_infer.make_forward_partial( seed=42, input_dim=4 + 2, # n-parameters (v,a,z,t) + n-data (rts and choices) state=os.path.join("data", "jax_models", "ddm", "ddm_jax_lan_train_state.jax"), add_jitted=True, )

Checking the signature of the JAX callable¶

We can test the signature of the JAX callable by passing in a batch of inputs and checking the output shape.

In [4]:

# Testing the signature of the JAX function 2
n_dim_model_parameters = 4
n_dim_data = 2
n_trials = 1
in_ = jnp.zeros((n_trials,
                 n_dim_model_parameters + n_dim_data))
out = jax_mlp_forward(in_)
print(out.shape)

# Testing the signature of the JAX function 2 n_dim_model_parameters = 4 n_dim_data = 2 n_trials = 1 in_ = jnp.zeros((n_trials, n_dim_model_parameters + n_dim_data)) out = jax_mlp_forward(in_) print(out.shape)

(1, 1)

In [5]:

from hssm.distribution_utils.jax import make_jax_single_trial_logp_from_network_forward
import jax.numpy as jnp
jax_logp = make_jax_single_trial_logp_from_network_forward(jax_forward_fn = jax_mlp_forward,
                                                           params_only = False)

from hssm.distribution_utils.jax import make_jax_single_trial_logp_from_network_forward import jax.numpy as jnp jax_logp = make_jax_single_trial_logp_from_network_forward(jax_forward_fn = jax_mlp_forward, params_only = False)

In [6]:

# Test call
jax_logp(jnp.zeros((2)),
         jnp.array([1.0]),
         jnp.array([1.5]),
         jnp.array([0.5]),
         jnp.array([0.3]))

# Test call jax_logp(jnp.zeros((2)), jnp.array([1.0]), jnp.array([1.5]), jnp.array([0.5]), jnp.array([0.3]))

Out[6]:

Array(-20.01882519, dtype=float64)

Decorate and wrap a simulator¶

The simulator-signature expected by hssm is the following:

A simulator is a callable that takes in a matrix that stacks model parameters and data trialwise, hence with input dimensions: $trials \times (|paramters))$. The function should return a matrix of shape $trials \times |data|$.

We will use the decorate_atomic_simulator() utility to annotate the simulator with necessary metadata we use the hssm_sim_wrapper() function from the ssm-simulators package to make our simulator ready for usage inside a PyTensor RandomVariable later.

If you check the signature of the resulting rv_ready_simulator, you should find no problems shoe-horning your own custon simulator into the corresponding behavior.

decorate_atomic_simulator() is just a Python decorator that you can use around any function.

You can start from any simulator you like, we use the one from the ssm-simulators package for convenience.

In [7]:

from ssms.hssm_support import hssm_sim_wrapper, decorate_atomic_simulator
from functools import partial
from ssms.basic_simulators.simulator import simulator

rv_ready_simulator = partial(hssm_sim_wrapper, 
                             simulator_fun = simulator,
                             model = "ddm", 
                             n_replicas = 1) # AF-TODO: n_replicas should default to 1 instead of being required

# We decorate the simulator to attach some metadata
# that HSSM can use
decorated_simulator = decorate_atomic_simulator(
    model_name="ddm", choices=[-1, 1], obs_dim=2
)(rv_ready_simulator)

from ssms.hssm_support import hssm_sim_wrapper, decorate_atomic_simulator from functools import partial from ssms.basic_simulators.simulator import simulator rv_ready_simulator = partial(hssm_sim_wrapper, simulator_fun = simulator, model = "ddm", n_replicas = 1) # AF-TODO: n_replicas should default to 1 instead of being required # We decorate the simulator to attach some metadata # that HSSM can use decorated_simulator = decorate_atomic_simulator( model_name="ddm", choices=[-1, 1], obs_dim=2 )(rv_ready_simulator)

Checking the signature of the decorated simulator¶

We can check the signature of the decorated simulator by passing in a batch of parameters and checking the output shape.

In [8]:

decorated_simulator(
    theta=np.tile(np.array([1.0, 1.5, 0.5, 0.2]), (10, 1)), random_state=42
)

decorated_simulator( theta=np.tile(np.array([1.0, 1.5, 0.5, 0.2]), (10, 1)), random_state=42 )

Out[8]:

array([[ 2.02210951,  1.        ],
       [ 2.58906269,  1.        ],
       [ 1.1866796 ,  1.        ],
       [ 1.06390691,  1.        ],
       [ 1.32591701,  1.        ],
       [ 1.04748344,  1.        ],
       [ 1.37392318,  1.        ],
       [ 0.78049529, -1.        ],
       [ 1.6851691 ,  1.        ],
       [ 1.10256469,  1.        ]])

Create valid PyMC distribution¶

We have all ingredients to create a valid PyMC distribution, form a few helper functions provided by hssm at this point.

A valid Distribution will need two ingredients:

1. A RandomVariable (RV) that encodes the simulator and parameter names.
2. A likelihood function, which is a valid PyTensor Op that encodes the log-likelihood of the model.

We will use the make_hssm_rv() helper function to create the RandomVariable and the make_likelihood_callable() helper function to create the likelihood Op.

Finally, the make_distribution() helper function will package everything into a valid PyMC distribution.

NOTE:

There are a few helpful settings which we can use to customize our distribution. We will not cover all details here, however it is worth highlighting the params_is_reg argument in make_likelihood_callable(). This argument is used to tell PyMC whether the parameter is a regression parameter or not. Specifically, if a parameter is treated as a regression, the likelihood function is built to assume that the parameter is passed trial-wise, i.e. as a vector of length n_trials.

Note here, we set all parameters to be non-regression parameters, as we expect the parameters to be passed as single values in the simple PyMC model below.

In [9]:

from hssm.distribution_utils.dist import (
    make_distribution,
    make_hssm_rv,
    make_likelihood_callable,
)

# Step 1: Define a pytensor RandomVariable
CustomRV = make_hssm_rv(
    simulator_fun=decorated_simulator,
    list_params=["v", "a", "z", "t"]
)

# Step 2: Define a likelihood function
logp_jax_op = make_likelihood_callable(
    loglik=jax_logp,
    loglik_kind="approx_differentiable",
    backend="jax",
    params_is_reg=[True, False, False, False],
    params_only=False,
)

# Step 3: Define a distribution
CustomDistribution = make_distribution(
    rv=CustomRV,
    loglik=logp_jax_op,
    list_params=["v", "a", "z", "t"],
    bounds=dict(v=(-3, 3), a=(0.5, 3.0), z=(0.1, 0.9), t=(0, 2.0)),
)

from hssm.distribution_utils.dist import ( make_distribution, make_hssm_rv, make_likelihood_callable, ) # Step 1: Define a pytensor RandomVariable CustomRV = make_hssm_rv( simulator_fun=decorated_simulator, list_params=["v", "a", "z", "t"] ) # Step 2: Define a likelihood function logp_jax_op = make_likelihood_callable( loglik=jax_logp, loglik_kind="approx_differentiable", backend="jax", params_is_reg=[True, False, False, False], params_only=False, ) # Step 3: Define a distribution CustomDistribution = make_distribution( rv=CustomRV, loglik=logp_jax_op, list_params=["v", "a", "z", "t"], bounds=dict(v=(-3, 3), a=(0.5, 3.0), z=(0.1, 0.9), t=(0, 2.0)), )

passing here: 
params_only:  False
params_is_reg:  [True, False, False, False]

We can now test the distribution by passing it to a simple PyMC model.

In [10]:

from pytensor import tensor as pt
# Test via basic pymc model
with pm.Model() as model_pymc:
    v = pm.Normal("v", mu=0, sigma=1)
    a = pm.Weibull("a", alpha=2.0, beta=1.2)
    z = pm.Beta("z", alpha=10, beta=10)
    t = pm.Weibull("t", alpha=1.5, beta=0.5)

    # We define `v` as a vector of length `n_trials`
    # To conform to the expected signature of the likelihood function
    v_det = pm.Deterministic("v_det", v * pt.ones(obs_ddm.shape[0]))
    CustomDistribution("custom", observed=obs_ddm.values, v=v_det, a=a, z=z, t=t)

from pytensor import tensor as pt # Test via basic pymc model with pm.Model() as model_pymc: v = pm.Normal("v", mu=0, sigma=1) a = pm.Weibull("a", alpha=2.0, beta=1.2) z = pm.Beta("z", alpha=10, beta=10) t = pm.Weibull("t", alpha=1.5, beta=0.5) # We define `v` as a vector of length `n_trials` # To conform to the expected signature of the likelihood function v_det = pm.Deterministic("v_det", v * pt.ones(obs_ddm.shape[0])) CustomDistribution("custom", observed=obs_ddm.values, v=v_det, a=a, z=z, t=t)

In [11]:

pm.model_to_graphviz(model_pymc)

pm.model_to_graphviz(model_pymc)

max_shape:  (500,)
size:  (np.int64(500),)

Out[11]:

In [12]:

with model_pymc:
    idata = pm.sample(draws=500,
                      tune=500,
                      chains=2,
                      nuts_sampler="numpyro")

with model_pymc: idata = pm.sample(draws=500, tune=500, chains=2, nuts_sampler="numpyro")

/Users/afengler/Library/CloudStorage/OneDrive-Personal/proj_hssm/HSSM/.venv/lib/python3.11/site-packages/pymc/sampling/jax.py:475: UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(
sample: 100%|██████████| 1000/1000 [00:12<00:00, 82.15it/s, 3 steps of size 4.10e-01. acc. prob=0.84] 
sample: 100%|██████████| 1000/1000 [00:10<00:00, 93.81it/s, 3 steps of size 3.90e-01. acc. prob=0.89] 
There was 1 divergence after tuning. Increase `target_accept` or reparameterize.
We recommend running at least 4 chains for robust computation of convergence diagnostics

In [13]:

az.summary(idata)

az.summary(idata)

Out[13]:

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
v	0.508	0.022	0.467	0.549	0.001	0.001	506.0	457.0	1.00
a	1.275	0.028	1.221	1.325	0.001	0.001	668.0	466.0	1.00
z	0.472	0.014	0.447	0.499	0.001	0.001	426.0	327.0	1.01
t	0.184	0.021	0.145	0.222	0.001	0.001	544.0	508.0	1.00
v_det[0]	0.508	0.022	0.467	0.549	0.001	0.001	506.0	457.0	1.00
...	...	...	...	...	...	...	...	...	...
v_det[495]	0.508	0.022	0.467	0.549	0.001	0.001	506.0	457.0	1.00
v_det[496]	0.508	0.022	0.467	0.549	0.001	0.001	506.0	457.0	1.00
v_det[497]	0.508	0.022	0.467	0.549	0.001	0.001	506.0	457.0	1.00
v_det[498]	0.508	0.022	0.467	0.549	0.001	0.001	506.0	457.0	1.00
v_det[499]	0.508	0.022	0.467	0.549	0.001	0.001	506.0	457.0	1.00

504 rows × 9 columns

In [14]:

az.plot_trace(idata)
plt.tight_layout()

az.plot_trace(idata) plt.tight_layout()

Custom HSSM model from simulator and JAX callable¶

Next, we will create a custom HSSM model from the simulator and JAX callable. After the work we have done above, this is now very straightforward.

The only hssm specific extra step is to define a ModelConfig object, which bundles all information about the model.

Then we pass our ModelConfig object to the HSSM class, along with the data and the log-likelihood function, and hssm will take care of the rest.

Importantly, hssm will automatically understand how to construct the correct likelihood function for the specified model configuration (parameter-wise regression settings, etc.). A step we have to accomplish manually in the code above.

In [15]:

# Define model config
my_custom_model_config = ModelConfig(
    response=["rt", "response"],
    list_params=["v", "a", "z", "t"],
    bounds={
        "v": (-2.5, 2.5),
        "a": (1.0, 3.0),
        "z": (0.0, 0.9),
        "t": (0.001, 2),
    },
    rv=decorated_simulator,
    backend="jax",
    choices=[-1, 1],
)

# Define model config my_custom_model_config = ModelConfig( response=["rt", "response"], list_params=["v", "a", "z", "t"], bounds={ "v": (-2.5, 2.5), "a": (1.0, 3.0), "z": (0.0, 0.9), "t": (0.001, 2), }, rv=decorated_simulator, backend="jax", choices=[-1, 1], )

In [16]:

# Define the HSSM model
model = hssm.HSSM(
    data=obs_ddm,
    model="my_new_model",  # some name for the model
    model_config=my_custom_model_config,
    loglik_kind="approx_differentiable",  # use the blackbox loglik
    loglik=jax_logp,
    p_outlier=0,
)

model.graph()

# Define the HSSM model model = hssm.HSSM( data=obs_ddm, model="my_new_model", # some name for the model model_config=my_custom_model_config, loglik_kind="approx_differentiable", # use the blackbox loglik loglik=jax_logp, p_outlier=0, ) model.graph()

You have specified the `lapse` argument to include a lapse distribution, but `p_outlier` is set to either 0 or None. Your lapse distribution will be ignored.
passing here: 
params_only:  False
params_is_reg:  [True, False, False, False]
Model initialized successfully.
max_shape:  (500,)
size:  (np.int64(500),)

Out[16]:

In [17]:

# Test sampling
model.sample(draws=500,
             tune=200,
             nuts_sampler="numpyro",
             chains = 2,
             cores = 2,
             discard_tuned_samples=False)

# Test sampling model.sample(draws=500, tune=200, nuts_sampler="numpyro", chains = 2, cores = 2, discard_tuned_samples=False)

Using default initvals. 

/Users/afengler/Library/CloudStorage/OneDrive-Personal/proj_hssm/HSSM/.venv/lib/python3.11/site-packages/pymc/sampling/jax.py:475: UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(
sample: 100%|██████████| 700/700 [00:08<00:00, 84.96it/s, 7 steps of size 4.33e-01. acc. prob=0.88]  
sample: 100%|██████████| 700/700 [00:07<00:00, 97.64it/s, 3 steps of size 4.85e-01. acc. prob=0.90]  
There was 1 divergence after tuning. Increase `target_accept` or reparameterize.
We recommend running at least 4 chains for robust computation of convergence diagnostics
100%|██████████| 1000/1000 [00:00<00:00, 1287.66it/s]

Out[17]:

arviz.InferenceData

• posterior

  <xarray.Dataset> Size: 36kB
  Dimensions:  (chain: 2, draw: 500)
  Coordinates:
    * chain    (chain) int64 16B 0 1
    * draw     (draw) int64 4kB 0 1 2 3 4 5 6 7 ... 493 494 495 496 497 498 499
  Data variables:
      v        (chain, draw) float64 8kB 0.5065 0.5218 0.5144 ... 0.4876 0.5138
      z        (chain, draw) float64 8kB 0.4671 0.4671 0.4669 ... 0.4706 0.4799
      t        (chain, draw) float64 8kB 0.1846 0.1945 0.1819 ... 0.1861 0.1711
      a        (chain, draw) float64 8kB 1.203 1.215 1.282 ... 1.304 1.275 1.259
  Attributes:
      created_at:                  2025-09-29T20:06:54.606700+00:00
      arviz_version:               0.22.0
      inference_library:           numpyro
      inference_library_version:   0.19.0
      sampling_time:               15.733445
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.15.0

  xarray.Dataset

  • Dimensions:
    • chain: 2
    • draw: 500
  • Coordinates: (2)
    • chain
      (chain)
      int64
      0 1

      array([0, 1])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499], shape=(500,))

  • Data variables: (4)
    • v
      (chain, draw)
      float64
      0.5065 0.5218 ... 0.4876 0.5138

      array([[0.50652225, 0.52177754, 0.51440562, 0.50536979, 0.51314912,
              0.54001204, 0.54985591, 0.5046882 , 0.51807674, 0.51625702,
              0.54692524, 0.51421264, 0.51710899, 0.49378862, 0.4947128 ,
              0.53235415, 0.51005456, 0.53332947, 0.52632567, 0.52684755,
              0.52779519, 0.5294698 , 0.4964419 , 0.53699494, 0.51500358,
              0.5055484 , 0.52071999, 0.5045647 , 0.51804783, 0.50986454,
              0.50864139, 0.50109367, 0.4881502 , 0.48044392, 0.50691332,
              0.51683577, 0.54273262, 0.51347458, 0.49269567, 0.47884544,
              0.48654871, 0.51413014, 0.48880319, 0.53460902, 0.50877805,
              0.52585213, 0.53724231, 0.48906722, 0.53126491, 0.50541683,
              0.51971993, 0.48543819, 0.49915962, 0.50343884, 0.49394213,
              0.5238966 , 0.53831292, 0.48980574, 0.49413236, 0.48285012,
              0.52668845, 0.5065933 , 0.53272115, 0.4759908 , 0.54927688,
              0.50804959, 0.51941195, 0.51273276, 0.50634931, 0.53831005,
              0.46902599, 0.53549095, 0.51559344, 0.49458857, 0.52373304,
              0.51676471, 0.51585202, 0.4971989 , 0.50324829, 0.52659986,
              0.48776029, 0.49583361, 0.51642685, 0.50864799, 0.52048169,
              0.50799867, 0.52166282, 0.51429888, 0.5165577 , 0.5260447 ,
              0.52355933, 0.50797142, 0.49952107, 0.50913998, 0.52369284,
              0.53037348, 0.51470877, 0.49160171, 0.51312112, 0.51316648,
      ...
              0.52179346, 0.52146549, 0.4975658 , 0.51490137, 0.49408913,
              0.49357823, 0.4903704 , 0.54105208, 0.53586734, 0.52674522,
              0.52276486, 0.51645444, 0.52549249, 0.51279359, 0.49941194,
              0.50048403, 0.54235656, 0.52382056, 0.53626012, 0.50025861,
              0.50750055, 0.51241461, 0.49984952, 0.52608315, 0.51588354,
              0.50221473, 0.50796516, 0.50973547, 0.50397186, 0.50601964,
              0.51215477, 0.5056883 , 0.56627884, 0.54770328, 0.53363162,
              0.51231985, 0.52510808, 0.51069525, 0.52613202, 0.5189598 ,
              0.53080553, 0.49671126, 0.48387644, 0.54014847, 0.49591714,
              0.51315711, 0.51697004, 0.4746685 , 0.45192018, 0.48479673,
              0.47378715, 0.5070617 , 0.51941783, 0.52041026, 0.51621638,
              0.51840835, 0.5376444 , 0.5480531 , 0.47862457, 0.55005902,
              0.48254969, 0.53249948, 0.49245901, 0.52639586, 0.51866585,
              0.52917051, 0.47613987, 0.52214713, 0.50733743, 0.5031042 ,
              0.50084576, 0.50611646, 0.52997496, 0.53510888, 0.4896497 ,
              0.49678009, 0.51052401, 0.49369998, 0.47115982, 0.48818947,
              0.51175901, 0.50525546, 0.50490745, 0.52387144, 0.47844011,
              0.5524566 , 0.47731821, 0.54346315, 0.47564284, 0.52586607,
              0.47151783, 0.50987316, 0.51923989, 0.50236927, 0.53014138,
              0.50434145, 0.53301445, 0.48701781, 0.48762538, 0.51375528]])

    • z
      (chain, draw)
      float64
      0.4671 0.4671 ... 0.4706 0.4799

      array([[0.46707289, 0.46705514, 0.46692803, 0.47987911, 0.45962194,
              0.46317014, 0.45906122, 0.46282049, 0.4644378 , 0.44430251,
              0.48033406, 0.46000606, 0.48352326, 0.47605609, 0.47954661,
              0.46501301, 0.45971501, 0.4553757 , 0.45790673, 0.45624751,
              0.46850379, 0.4596723 , 0.490406  , 0.46825442, 0.46342307,
              0.47528961, 0.49005823, 0.45315351, 0.47394187, 0.47479753,
              0.46634729, 0.4548491 , 0.46783867, 0.48156271, 0.47026937,
              0.47505446, 0.47346747, 0.45719763, 0.50011816, 0.46522773,
              0.48447175, 0.47669342, 0.46732189, 0.45849133, 0.47834132,
              0.45868229, 0.48453105, 0.47224073, 0.47483256, 0.45105506,
              0.45361608, 0.49861481, 0.46291452, 0.47087862, 0.46867769,
              0.44926645, 0.45506913, 0.47721   , 0.44919516, 0.44265885,
              0.47237607, 0.47273957, 0.4595531 , 0.4970589 , 0.44651148,
              0.47823903, 0.47181496, 0.46244842, 0.47561841, 0.48189254,
              0.48110401, 0.49584675, 0.49425351, 0.46092834, 0.48234143,
              0.45123694, 0.47135874, 0.4834707 , 0.47291848, 0.47194059,
              0.47459032, 0.47560924, 0.47769324, 0.46442091, 0.47322174,
              0.48696376, 0.49363185, 0.49298928, 0.46981199, 0.47354827,
              0.45761414, 0.47514897, 0.47331598, 0.48249721, 0.47704065,
              0.4615957 , 0.48299653, 0.48384578, 0.47130569, 0.456024  ,
      ...
              0.45540446, 0.49545631, 0.49123246, 0.4697623 , 0.47844338,
              0.48505791, 0.47279213, 0.47322986, 0.47536561, 0.48097117,
              0.46548548, 0.46817435, 0.47038046, 0.476976  , 0.48116558,
              0.47031945, 0.44471027, 0.43314817, 0.46103846, 0.47088118,
              0.47016685, 0.47509805, 0.47031555, 0.47021288, 0.45339514,
              0.47489114, 0.47029873, 0.48359386, 0.48268777, 0.48784088,
              0.49077189, 0.43836746, 0.43322178, 0.43823064, 0.44973226,
              0.48783122, 0.45595254, 0.45344318, 0.45896935, 0.45753584,
              0.46321131, 0.47667537, 0.45512941, 0.46965716, 0.47920567,
              0.47711878, 0.46174016, 0.49011458, 0.50081173, 0.51565168,
              0.49756786, 0.47295238, 0.48417192, 0.47976807, 0.47978892,
              0.47472801, 0.47449449, 0.47090881, 0.46405832, 0.47460208,
              0.45743532, 0.48150159, 0.47477507, 0.49362611, 0.48366951,
              0.48506502, 0.45938746, 0.4683823 , 0.47814166, 0.47370846,
              0.47876426, 0.47953095, 0.47232635, 0.46383063, 0.44893052,
              0.47465338, 0.46168378, 0.48665151, 0.48094522, 0.48359071,
              0.47340486, 0.5037483 , 0.44489693, 0.48619151, 0.47417816,
              0.4667634 , 0.44728717, 0.47600679, 0.45439402, 0.48325295,
              0.46113067, 0.47756224, 0.46874161, 0.45013817, 0.46071212,
              0.45152349, 0.45267524, 0.47384116, 0.47057068, 0.4799169 ]])

    • t
      (chain, draw)
      float64
      0.1846 0.1945 ... 0.1861 0.1711

      array([[0.18456158, 0.19451824, 0.18187668, 0.19215992, 0.18040708,
              0.12422972, 0.12987815, 0.17223773, 0.18662542, 0.17583724,
              0.20090795, 0.18522286, 0.18899613, 0.17754081, 0.18308799,
              0.20097747, 0.16911695, 0.17705788, 0.18244548, 0.17024359,
              0.19510386, 0.17679574, 0.16724145, 0.19150161, 0.17713633,
              0.1744337 , 0.16516208, 0.18628509, 0.1881942 , 0.16294873,
              0.18868426, 0.14776777, 0.14421856, 0.14563122, 0.19885959,
              0.18502666, 0.18313126, 0.20998347, 0.17602116, 0.20520314,
              0.17027401, 0.19856656, 0.16295118, 0.16638263, 0.16790246,
              0.18074459, 0.17463519, 0.15856979, 0.21700053, 0.16139955,
              0.17217231, 0.22924561, 0.19018002, 0.17018554, 0.21247883,
              0.180579  , 0.1880965 , 0.16739468, 0.15096576, 0.15606402,
              0.18655344, 0.17892445, 0.17249962, 0.18054001, 0.19018294,
              0.19184777, 0.18263277, 0.19645286, 0.19156914, 0.20134055,
              0.17557041, 0.20403268, 0.19747363, 0.18045437, 0.17238034,
              0.16074271, 0.18261818, 0.1944816 , 0.19033238, 0.20844295,
              0.17417366, 0.19716517, 0.18083023, 0.17025333, 0.16726881,
              0.18727384, 0.18311117, 0.1833653 , 0.20115333, 0.1861249 ,
              0.19073879, 0.18817095, 0.20039477, 0.21020831, 0.21336497,
              0.2124005 , 0.18041996, 0.1952306 , 0.16834383, 0.18055734,
      ...
              0.15711347, 0.21299212, 0.19809229, 0.19636827, 0.20976128,
              0.20781172, 0.20064402, 0.16720983, 0.16467209, 0.17013383,
              0.19335899, 0.16243467, 0.1709768 , 0.17978643, 0.20401343,
              0.17196931, 0.16499063, 0.16200092, 0.14032082, 0.19433783,
              0.20421419, 0.17053999, 0.17132096, 0.17421094, 0.20411631,
              0.18975589, 0.19629324, 0.17071886, 0.21954638, 0.18791472,
              0.17990066, 0.17335143, 0.16523443, 0.15891853, 0.16167262,
              0.17424785, 0.17920048, 0.16052326, 0.15164332, 0.13682306,
              0.13114856, 0.20541622, 0.16342143, 0.20366233, 0.19828323,
              0.19262285, 0.1797596 , 0.1923774 , 0.18446856, 0.15691725,
              0.16060145, 0.1906133 , 0.19141024, 0.20975387, 0.19775923,
              0.20634758, 0.16942101, 0.18197607, 0.20001394, 0.16637003,
              0.15506095, 0.20135218, 0.1945418 , 0.19857911, 0.20059121,
              0.19653459, 0.1787149 , 0.15941823, 0.147958  , 0.15013485,
              0.1802808 , 0.17507775, 0.17481469, 0.15813237, 0.19462434,
              0.19565736, 0.1610906 , 0.20360879, 0.1998874 , 0.201282  ,
              0.19814574, 0.18558622, 0.17281304, 0.16678042, 0.16033091,
              0.18324388, 0.13495458, 0.13226886, 0.14767839, 0.20636781,
              0.16629921, 0.19763679, 0.19788505, 0.18946303, 0.18261652,
              0.1677317 , 0.15499104, 0.20004703, 0.1861162 , 0.17112339]])

    • a
      (chain, draw)
      float64
      1.203 1.215 1.282 ... 1.275 1.259

      array([[1.20326242, 1.21547486, 1.28196939, 1.26824224, 1.27868017,
              1.30775333, 1.31356188, 1.29482078, 1.27108465, 1.29994867,
              1.26005882, 1.27176033, 1.29019972, 1.30064045, 1.29125096,
              1.25399686, 1.29259317, 1.2601696 , 1.26543079, 1.29476289,
              1.28216191, 1.24519897, 1.29597018, 1.3018195 , 1.24366272,
              1.30443081, 1.31850478, 1.25406027, 1.27450707, 1.25190875,
              1.29008927, 1.29719632, 1.29591881, 1.28027005, 1.30248993,
              1.28426117, 1.27737017, 1.21993823, 1.30872926, 1.25389113,
              1.26727956, 1.26487547, 1.31002092, 1.32795031, 1.30623199,
              1.27422322, 1.31676113, 1.30595967, 1.23553664, 1.30046996,
              1.30225591, 1.26614353, 1.26839809, 1.26863994, 1.26868928,
              1.27788577, 1.26496472, 1.30559604, 1.26260209, 1.2766912 ,
              1.25675627, 1.30552605, 1.26219227, 1.29302994, 1.27901441,
              1.26583056, 1.29682506, 1.2845927 , 1.24105554, 1.25024567,
              1.31194613, 1.25328479, 1.25666717, 1.27529196, 1.29527339,
              1.25951093, 1.27436495, 1.27437558, 1.28395351, 1.25792395,
              1.2997523 , 1.28692682, 1.28204978, 1.29471661, 1.2712169 ,
              1.34678947, 1.31455788, 1.32730276, 1.29021304, 1.22595541,
              1.27639303, 1.25500251, 1.25320671, 1.2656999 , 1.27100652,
              1.23800855, 1.25475011, 1.26935107, 1.27733733, 1.27980021,
      ...
              1.27379357, 1.23233555, 1.21868056, 1.26735702, 1.21958769,
              1.24634098, 1.24813669, 1.28999096, 1.30438674, 1.31657242,
              1.2849697 , 1.27378044, 1.27769661, 1.27559349, 1.27533641,
              1.28208649, 1.28115603, 1.28604052, 1.29983571, 1.28544507,
              1.25815635, 1.29233857, 1.32336619, 1.29625323, 1.23553288,
              1.2905234 , 1.25885322, 1.29361631, 1.26489186, 1.25603513,
              1.25230589, 1.30001   , 1.29138987, 1.30278022, 1.29795554,
              1.28683047, 1.27622095, 1.28748055, 1.29358678, 1.300482  ,
              1.29138781, 1.2969677 , 1.310484  , 1.22659953, 1.26222605,
              1.26211137, 1.28587821, 1.2690419 , 1.27666632, 1.29772072,
              1.28817393, 1.30744993, 1.27408361, 1.28353668, 1.27233579,
              1.27991762, 1.25252658, 1.25243262, 1.26710836, 1.30641241,
              1.306253  , 1.28083182, 1.27312591, 1.30061871, 1.30713919,
              1.29103343, 1.28194207, 1.28620973, 1.25526126, 1.27399556,
              1.29309899, 1.28797268, 1.31404751, 1.27042055, 1.25840402,
              1.28255085, 1.27111364, 1.29644186, 1.20631779, 1.2146087 ,
              1.28895929, 1.29617733, 1.2845508 , 1.28515429, 1.29614503,
              1.30256947, 1.27369114, 1.32740797, 1.28741486, 1.29841707,
              1.27417442, 1.24022551, 1.24035251, 1.2631388 , 1.230291  ,
              1.31636072, 1.25485288, 1.30445601, 1.27485287, 1.25944683]])

  • Indexes: (2)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='draw', length=500))

  • Attributes: (8)

    created_at :

    2025-09-29T20:06:54.606700+00:00

    arviz_version :

    0.22.0

    inference_library :

    numpyro

    inference_library_version :

    0.19.0

    sampling_time :

    15.733445

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.15.0

  
  
• log_likelihood

  <xarray.Dataset> Size: 4MB
  Dimensions:      (chain: 2, draw: 500, __obs__: 500)
  Coordinates:
    * chain        (chain) int64 16B 0 1
    * draw         (draw) int64 4kB 0 1 2 3 4 5 6 ... 493 494 495 496 497 498 499
    * __obs__      (__obs__) int64 4kB 0 1 2 3 4 5 6 ... 494 495 496 497 498 499
  Data variables:
      rt,response  (chain, draw, __obs__) float64 4MB -0.6224 -2.021 ... -1.468
  Attributes:
      modeling_interface:          bambi
      modeling_interface_version:  0.15.0

  xarray.Dataset

  • Dimensions:
    • chain: 2
    • draw: 500
    • __obs__: 500
  • Coordinates: (3)
    • chain
      (chain)
      int64
      0 1

      array([0, 1])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499], shape=(500,))

    • __obs__
      (__obs__)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499], shape=(500,))

  • Data variables: (1)
    • rt,response
      (chain, draw, __obs__)
      float64
      -0.6224 -2.021 ... -2.053 -1.468

      array([[[-0.62237211, -2.02123135, -0.65461652, ..., -2.02499833,
               -1.95394559, -1.4575279 ],
              [-0.6157066 , -2.03435519, -0.64863427, ..., -1.99133337,
               -1.96757512, -1.58214591],
              [-0.69634185, -2.09189441, -0.7211599 , ..., -1.91077587,
               -2.02821286, -1.77029471],
              ...,
              [-0.77577145, -2.24709649, -0.78745507, ..., -1.86565394,
               -2.14412899, -1.62868603],
              [-0.79030502, -2.20325618, -0.80532601, ..., -1.87062328,
               -2.10726798, -1.6703428 ],
              [-0.71591563, -2.13398685, -0.74138719, ..., -1.9236206 ,
               -2.06418349, -1.55125517]],
      
             [[-0.68506456, -2.10268973, -0.71375864, ..., -2.00611148,
               -2.02687599, -1.27847818],
              [-0.7164939 , -2.1873615 , -0.7444462 , ..., -1.88696752,
               -2.11639722, -1.64060232],
              [-0.63470474, -2.0338495 , -0.66775462, ..., -1.97647035,
               -1.96845704, -1.6417243 ],
              ...,
              [-0.73608402, -2.0941854 , -0.743669  , ..., -1.881587  ,
               -2.04065418, -2.05846558],
              [-0.6961844 , -2.06422003, -0.72446526, ..., -1.92972615,
               -2.00719751, -1.76601448],
              [-0.64641981, -2.11335535, -0.67403786, ..., -1.9571649 ,
               -2.05325828, -1.46773047]]], shape=(2, 500, 500))

  • Indexes: (3)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='draw', length=500))

    • __obs__
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='__obs__', length=500))

  • Attributes: (2)

    modeling_interface :

    bambi

    modeling_interface_version :

    0.15.0

  
  
• sample_stats

  <xarray.Dataset> Size: 53kB
  Dimensions:          (chain: 2, draw: 500)
  Coordinates:
    * chain            (chain) int64 16B 0 1
    * draw             (draw) int64 4kB 0 1 2 3 4 5 6 ... 494 495 496 497 498 499
  Data variables:
      acceptance_rate  (chain, draw) float64 8kB 0.5239 0.7118 ... 1.0 0.7309
      step_size        (chain, draw) float64 8kB 0.4328 0.4328 ... 0.4848 0.4848
      diverging        (chain, draw) bool 1kB False False False ... False False
      energy           (chain, draw) float64 8kB 908.3 907.6 907.0 ... 903.9 903.8
      n_steps          (chain, draw) int64 8kB 15 7 7 7 7 7 7 7 ... 7 7 7 7 7 7 3
      tree_depth       (chain, draw) int64 8kB 4 3 3 3 3 3 3 3 ... 3 3 3 3 3 3 3 2
      lp               (chain, draw) float64 8kB 906.8 904.2 900.9 ... 901.5 903.1
  Attributes:
      created_at:                  2025-09-29T20:06:54.612280+00:00
      arviz_version:               0.22.0
      modeling_interface:          bambi
      modeling_interface_version:  0.15.0

  xarray.Dataset

  • Dimensions:
    • chain: 2
    • draw: 500
  • Coordinates: (2)
    • chain
      (chain)
      int64
      0 1

      array([0, 1])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499], shape=(500,))

  • Data variables: (7)
    • acceptance_rate
      (chain, draw)
      float64
      0.5239 0.7118 0.9705 ... 1.0 0.7309

      array([[0.52393343, 0.7118161 , 0.97046066, 0.99047327, 0.95275798,
              0.84114153, 0.98280665, 0.97711775, 0.9792688 , 0.83610171,
              1.        , 0.92483341, 0.96295498, 0.82696203, 0.99452692,
              0.86324609, 0.77493171, 0.9808109 , 1.        , 0.86130655,
              0.79840878, 0.7806092 , 1.        , 0.9628028 , 0.82035009,
              0.85751029, 0.85454662, 1.        , 0.98873153, 0.34775265,
              0.7708082 , 0.89134292, 0.99022155, 0.79857568, 1.        ,
              0.78696712, 0.79584343, 0.81427266, 0.98239214, 0.9634928 ,
              0.40036365, 0.90296017, 0.98829238, 0.86769189, 0.88399572,
              0.98852822, 0.85450386, 0.99518107, 0.98527684, 0.39970448,
              0.89543664, 0.89073655, 0.98015966, 1.        , 0.95988642,
              0.59705642, 0.93709688, 0.99689801, 0.92237694, 0.93106675,
              0.97379331, 0.99480084, 0.92797063, 0.97461149, 0.95558046,
              0.97000395, 0.86045566, 0.81264795, 0.87548354, 0.87680889,
              0.97712315, 0.88399498, 0.9743201 , 0.86514813, 0.9624663 ,
              0.94834954, 0.90691949, 0.98379004, 0.87732389, 0.96017077,
              0.71762895, 0.92300329, 0.95958699, 0.87619106, 0.71166408,
              0.76419653, 0.98761606, 0.98748984, 0.72512779, 0.49136033,
              0.96735357, 0.98173797, 0.88189485, 0.95763691, 0.62690973,
              0.5129819 , 0.40850283, 0.93313876, 0.9012123 , 0.94945998,
      ...
              0.95583748, 0.84951772, 1.        , 0.53814874, 0.83945149,
              0.8862674 , 0.97614181, 0.99198861, 1.        , 1.        ,
              0.98723618, 0.96339296, 0.98080841, 0.99811316, 0.93857612,
              0.95841922, 0.96749002, 0.87690711, 1.        , 1.        ,
              0.98337577, 0.89790626, 0.88550373, 1.        , 0.86407164,
              0.99028176, 0.9901891 , 0.9398204 , 0.99126409, 0.99424472,
              0.78186373, 0.88671442, 0.97604768, 0.99848435, 1.        ,
              0.92315637, 1.        , 0.96837409, 0.99873499, 0.97747662,
              0.92572133, 0.79628691, 0.50545273, 0.98748658, 0.8937786 ,
              0.95925382, 0.73115822, 0.77018945, 0.85131121, 0.76928097,
              0.99731288, 0.95791786, 0.99571711, 0.63428965, 0.93768797,
              1.        , 0.36302055, 0.97678579, 0.99749454, 0.97775436,
              1.        , 0.92395086, 1.        , 0.90329417, 1.        ,
              0.99085133, 1.        , 1.        , 0.76740404, 1.        ,
              0.91583443, 0.96247439, 0.96107379, 0.72816971, 0.7804471 ,
              0.9786633 , 0.95050823, 0.99389834, 0.82813873, 1.        ,
              0.99803553, 0.86036999, 0.9908861 , 0.94057096, 0.7469984 ,
              0.97565615, 0.84282805, 1.        , 0.99239291, 0.99662151,
              0.69987799, 0.98116611, 0.98346206, 0.9341981 , 0.98345025,
              0.69671508, 1.        , 0.95375352, 1.        , 0.73092882]])

    • step_size
      (chain, draw)
      float64
      0.4328 0.4328 ... 0.4848 0.4848

      array([[0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
              0.43279338, 0.43279338, 0.43279338, 0.43279338, 0.43279338,
      ...
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ,
              0.484808  , 0.484808  , 0.484808  , 0.484808  , 0.484808  ]])

    • diverging
      (chain, draw)
      bool
      False False False ... False False

      array([[False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
      ...
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False, False, False, False, False,
              False, False, False, False, False]])

    • energy
      (chain, draw)
      float64
      908.3 907.6 907.0 ... 903.9 903.8

      array([[908.34076236, 907.62818975, 907.00864962, 901.35579791,
              902.80754054, 906.28132199, 906.91650315, 908.20481244,
              902.88065434, 905.05649815, 907.26845232, 905.9481659 ,
              902.40790906, 905.44199188, 901.75302691, 902.52820246,
              903.25661073, 902.81494617, 902.27455634, 903.63143594,
              903.40827034, 903.70816852, 904.16958751, 905.81101362,
              903.74080265, 904.67782595, 904.06084592, 904.56506076,
              903.26585246, 905.91748583, 904.29799677, 905.05127697,
              904.75932698, 908.33434615, 905.35257701, 903.32336945,
              904.80198196, 906.50554129, 904.95056794, 905.08032574,
              905.83160879, 905.16468963, 903.68855069, 906.53514962,
              904.97303991, 902.63570285, 904.95474201, 905.80005178,
              903.34552045, 905.63439099, 904.17673769, 906.4278951 ,
              905.18265345, 902.31126547, 903.22169666, 908.17350038,
              904.61969002, 903.36386108, 906.5133706 , 907.58213331,
              908.29381492, 902.1674065 , 902.88090852, 904.51104245,
              908.41751648, 905.08438282, 902.3877344 , 903.92753529,
              903.14900719, 904.63486455, 906.41457957, 907.53837331,
              906.1297282 , 903.67033672, 902.82567269, 904.09195511,
              904.25854434, 902.08692238, 902.51592422, 903.37989244,
      ...
              902.18649039, 903.64735183, 905.05340912, 903.00820888,
              904.97297878, 904.27886243, 902.20837809, 903.08408358,
              904.47749201, 904.76886016, 904.16339119, 907.69907289,
              910.73706653, 908.96088915, 905.21769741, 905.09676831,
              902.77211512, 903.05040847, 902.89269536, 903.97154871,
              904.97655557, 906.4496457 , 906.43807476, 907.10981224,
              907.82776893, 902.32447187, 903.2137481 , 906.37616979,
              905.49922728, 910.55362341, 909.30308597, 906.00377764,
              903.13163867, 903.99197895, 903.16098447, 902.35344091,
              907.12841664, 906.49182135, 904.18277083, 904.58565177,
              904.88463539, 904.67628716, 902.89334654, 905.04650418,
              905.42562531, 903.99547013, 904.63364435, 903.93756645,
              906.50298206, 906.15884341, 904.722615  , 901.46418653,
              902.65072502, 905.35461497, 906.30710108, 906.00293486,
              903.1248005 , 904.07276554, 908.33580426, 906.91402628,
              904.44712909, 905.07256623, 906.30424013, 905.58901253,
              905.44267737, 904.6776849 , 909.0812336 , 909.45838514,
              907.4972217 , 908.31725365, 906.83006527, 905.75322667,
              903.41577561, 905.37028187, 905.27670818, 906.9604372 ,
              905.88113978, 905.61972084, 903.8785336 , 903.82843125]])

    • n_steps
      (chain, draw)
      int64
      15 7 7 7 7 7 7 7 ... 7 7 7 7 7 7 3

      array([[15,  7,  7,  7,  7,  7,  7,  7,  7,  7,  7,  3,  7,  7,  3,  7,
               7,  7,  7,  7,  7,  7,  7,  7, 15,  7,  3,  7,  7,  3,  7,  7,
              15, 15,  7,  3,  3,  7, 15, 15,  7,  7,  7,  7,  7,  7,  7,  7,
               7,  7,  3, 11,  7,  7,  7, 15,  7,  7,  7,  3,  7, 11,  7,  7,
              15,  7,  7,  7,  7,  7,  7,  7,  3,  7,  7,  7,  7,  7,  7,  7,
               7,  7, 15,  7,  7,  3,  7,  7,  7,  3,  3,  7,  7,  7,  7,  7,
               7,  7, 15,  7,  7,  3,  7,  7,  7,  7,  7, 15,  7, 15,  7,  7,
              11, 15,  7,  7, 15, 15,  7,  7,  3,  7,  7, 11,  7,  3,  3, 11,
               3,  7,  7,  5,  7,  3,  7, 15, 15,  7,  3,  7,  7,  7,  7, 15,
               7, 15, 15,  3, 15,  7,  7,  7,  7, 15,  7,  7,  3,  3,  7,  7,
               7,  7,  7,  7,  3,  7,  3,  3, 15,  7,  7,  7,  3, 15,  3, 15,
               7,  7, 15, 15,  7,  7,  7,  7,  3,  7,  3, 11, 15,  3, 19, 15,
               7,  7,  7,  7,  7,  3, 11, 15,  7,  7,  7,  3,  7,  7,  7,  7,
               7,  7,  7,  7,  7,  7,  7,  7,  7,  3,  7,  7,  7,  7,  3,  3,
               7,  7,  7,  7,  7,  7,  7,  7,  7,  7,  3,  7,  7, 11,  7,  7,
               7,  7,  7,  7,  7,  7,  7,  7,  3,  7,  3,  7,  7,  7,  7,  7,
               7,  7,  7, 15, 15,  3,  3,  7,  3,  7,  7,  5,  7,  3,  7,  3,
               7, 15,  7,  7,  3,  7,  7,  7,  7,  7,  3,  7,  7, 39, 15,  3,
               7,  3,  3,  7,  7, 15,  7,  7,  3,  3, 15,  7,  3,  3,  7,  7,
               3,  7,  7,  7, 11,  1, 15,  7,  3,  7,  7,  6,  3,  7,  7, 23,
      ...
               7,  7,  7,  7,  3,  7,  7, 15,  3, 15,  7, 15,  7, 15, 15, 15,
               3,  7,  7,  7,  3,  7, 15,  7, 15,  3,  7, 15,  7,  7,  7,  7,
              15,  7, 15,  3, 15, 15,  3, 11,  7,  7, 15, 15, 15, 11, 15,  3,
               7, 15,  7,  7,  7,  7, 15,  7,  7,  7, 11,  3, 11, 15, 15,  7,
               7,  3, 15,  7,  7,  7,  7, 15,  7, 15,  7,  3,  7,  7, 15,  7,
               7,  7,  7,  7,  7,  3,  7,  7,  7,  7,  7,  7,  3,  7,  3,  7,
              15, 15,  7,  3,  7,  7,  7,  7,  7,  7,  7,  7,  7,  7,  3,  7,
               7,  7,  7,  7,  7,  7, 15,  7,  7,  7,  7,  7,  7,  7,  7,  7,
               7,  7,  7,  7,  7,  7,  7,  7,  7,  7,  7,  7,  7,  7,  7,  7,
               7,  3, 15,  7,  7, 15,  7,  3,  3,  7,  7,  7,  3,  7,  7,  7,
              15,  7,  3,  7,  7,  3, 15,  7,  7,  7,  7,  7,  7,  3,  7,  7,
               7,  7,  7,  7,  7,  7,  7, 11,  7,  7,  7,  7,  7,  7,  7,  7,
               7,  7,  7,  7,  7,  7,  7, 15, 15,  7,  3,  7,  7,  7,  7,  7,
               7,  7,  7,  7,  7,  7,  7, 15,  7,  7,  7,  7,  7,  7,  7, 15,
               7,  3,  7,  7,  7, 15,  7,  3,  7,  7,  7,  7, 15,  7,  3, 15,
               7,  3,  3, 15,  7,  7,  7,  7,  3, 15,  7,  7,  7,  5, 11,  7,
               3,  7,  3,  7, 15,  7,  3,  3,  3,  7,  7, 11,  7,  7,  7,  7,
               7,  7,  7,  7,  3,  7,  7,  3,  7,  7, 15,  7,  7,  7,  7,  3,
              15,  7,  7,  7,  7,  7,  7,  7,  7, 15,  7,  7,  7,  7,  7,  7,
               7,  7,  7,  3]])

    • tree_depth
      (chain, draw)
      int64
      4 3 3 3 3 3 3 3 ... 3 3 3 3 3 3 3 2

      array([[4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 3, 3, 2, 3, 3, 3, 3, 3, 3, 3,
              3, 3, 4, 3, 2, 3, 3, 2, 3, 3, 4, 4, 3, 2, 2, 3, 4, 4, 3, 3, 3, 3,
              3, 3, 3, 3, 3, 3, 2, 4, 3, 3, 3, 4, 3, 3, 3, 2, 3, 4, 3, 3, 4, 3,
              3, 3, 3, 3, 3, 3, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 3, 3, 2, 3, 3,
              3, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4, 3, 3, 2, 3, 3, 3, 3, 3, 4, 3, 4,
              3, 3, 4, 4, 3, 3, 4, 4, 3, 3, 2, 3, 3, 4, 3, 2, 2, 4, 2, 3, 3, 3,
              3, 2, 3, 4, 4, 3, 2, 3, 3, 3, 3, 4, 3, 4, 4, 2, 4, 3, 3, 3, 3, 4,
              3, 3, 2, 2, 3, 3, 3, 3, 3, 3, 2, 3, 2, 2, 4, 3, 3, 3, 2, 4, 2, 4,
              3, 3, 4, 4, 3, 3, 3, 3, 2, 3, 2, 4, 4, 2, 5, 4, 3, 3, 3, 3, 3, 2,
              4, 4, 3, 3, 3, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 3, 3,
              3, 3, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 3, 3, 4, 3, 3, 3, 3,
              3, 3, 3, 3, 3, 3, 2, 3, 2, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 2, 2, 3,
              2, 3, 3, 3, 3, 2, 3, 2, 3, 4, 3, 3, 2, 3, 3, 3, 3, 3, 2, 3, 3, 6,
              4, 2, 3, 2, 2, 3, 3, 4, 3, 3, 2, 2, 4, 3, 2, 2, 3, 3, 2, 3, 3, 3,
              4, 1, 4, 3, 2, 3, 3, 3, 2, 3, 3, 5, 3, 3, 3, 3, 3, 3, 4, 3, 4, 3,
              3, 3, 3, 3, 3, 2, 3, 3, 3, 4, 4, 2, 3, 3, 3, 3, 2, 3, 3, 3, 3, 2,
              2, 3, 4, 3, 3, 3, 3, 3, 3, 3, 3, 2, 3, 3, 3, 3, 3, 3, 3, 3, 2, 4,
              3, 3, 3, 3, 4, 3, 5, 3, 3, 3, 2, 4, 3, 3, 3, 3, 3, 3, 3, 4, 3, 3,
              3, 3, 3, 4, 3, 3, 3, 3, 3, 2, 3, 2, 3, 4, 3, 4, 3, 2, 2, 3, 2, 3,
              3, 3, 3, 3, 3, 3, 4, 3, 3, 3, 2, 3, 3, 4, 3, 3, 3, 3, 4, 3, 3, 3,
      ...
              2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4,
              3, 4, 4, 3, 2, 3, 3, 3, 3, 4, 3, 3, 4, 3, 3, 3, 3, 3, 3, 4, 3, 2,
              2, 3, 3, 3, 2, 3, 3, 3, 3, 3, 3, 4, 3, 3, 3, 3, 3, 3, 3, 2, 3, 3,
              3, 4, 3, 3, 3, 3, 3, 4, 2, 2, 3, 3, 4, 3, 3, 3, 3, 4, 3, 3, 2, 3,
              2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 3, 3, 4, 3, 3, 3, 3, 2, 3, 3, 3,
              3, 3, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4, 3, 3, 3, 3, 3, 3, 3, 3, 2, 3,
              3, 4, 2, 4, 3, 4, 3, 4, 4, 4, 2, 3, 3, 3, 2, 3, 4, 3, 4, 2, 3, 4,
              3, 3, 3, 3, 4, 3, 4, 2, 4, 4, 2, 4, 3, 3, 4, 4, 4, 4, 4, 2, 3, 4,
              3, 3, 3, 3, 4, 3, 3, 3, 4, 2, 4, 4, 4, 3, 3, 2, 4, 3, 3, 3, 3, 4,
              3, 4, 3, 2, 3, 3, 4, 3, 3, 3, 3, 3, 3, 2, 3, 3, 3, 3, 3, 3, 2, 3,
              2, 3, 4, 4, 3, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 3, 3, 3, 3, 3,
              3, 3, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
              3, 3, 3, 3, 3, 3, 3, 2, 4, 3, 3, 4, 3, 2, 2, 3, 3, 3, 2, 3, 3, 3,
              4, 3, 2, 3, 3, 2, 4, 3, 3, 3, 3, 3, 3, 2, 3, 3, 3, 3, 3, 3, 3, 3,
              3, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 3, 2, 3,
              3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 3, 3, 3, 3, 3, 3, 3, 4, 3, 2,
              3, 3, 3, 4, 3, 2, 3, 3, 3, 3, 4, 3, 2, 4, 3, 2, 2, 4, 3, 3, 3, 3,
              2, 4, 3, 3, 3, 3, 4, 3, 2, 3, 2, 3, 4, 3, 2, 2, 2, 3, 3, 4, 3, 3,
              3, 3, 3, 3, 3, 3, 2, 3, 3, 2, 3, 3, 4, 3, 3, 3, 3, 2, 4, 3, 3, 3,
              3, 3, 3, 3, 3, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2]])

    • lp
      (chain, draw)
      float64
      906.8 904.2 900.9 ... 901.5 903.1

      array([[906.769642  , 904.23236474, 900.912389  , 901.01483087,
              901.33862262, 905.66255571, 905.64204048, 901.50760461,
              901.08629863, 904.32974108, 904.14608215, 901.3561925 ,
              901.57268067, 901.57759171, 901.2872726 , 901.86252107,
              901.54493053, 902.25950593, 901.6493578 , 901.828825  ,
              901.41447177, 902.82858124, 903.0978933 , 902.50146752,
              902.51758929, 901.39885246, 903.84967164, 902.47708703,
              900.91985112, 904.03128632, 901.34611998, 903.31081742,
              903.86524136, 905.47079198, 902.55635645, 900.9160106 ,
              902.52704921, 903.49258855, 904.08282695, 902.8293488 ,
              903.31803294, 901.15676437, 902.78414041, 903.50487934,
              901.80416287, 901.49115873, 903.73848507, 902.7683851 ,
              903.12857773, 902.99475006, 902.48545246, 905.39801252,
              901.53626592, 901.79704072, 902.87299782, 902.58055751,
              902.54727812, 902.27568028, 905.06750364, 906.27412812,
              901.69288249, 901.45597628, 902.23231573, 903.55203981,
              904.38935788, 900.98523569, 901.19965523, 901.60991687,
              901.96896919, 903.5373384 , 903.68031709, 905.27727669,
              902.99540891, 901.80496423, 902.18384642, 903.19167336,
              900.86660825, 901.32505022, 901.06458487, 901.72072188,
      ...
              901.1013042 , 901.35794805, 902.73267809, 901.34014202,
              902.92390053, 901.24445666, 901.09963531, 901.99834813,
              902.98877348, 902.24871146, 904.01122092, 906.30900509,
              907.93102397, 905.30460068, 902.72480678, 902.37620249,
              901.68143736, 902.24651562, 902.43323965, 903.59777545,
              904.87389186, 903.09313592, 904.58988073, 904.14191923,
              901.30761239, 901.08069641, 901.24521092, 902.45041688,
              905.30687745, 908.63752385, 905.61766221, 902.08934437,
              901.58340512, 902.59686668, 901.34756825, 901.81222061,
              904.6614735 , 904.10850413, 902.82907342, 904.09437442,
              904.24065091, 902.63624297, 901.29861631, 904.17865003,
              903.15176579, 902.83993826, 903.53033349, 902.08936653,
              905.99555742, 903.7742865 , 901.16508933, 901.3465417 ,
              902.10566217, 903.27606066, 904.11437364, 901.32712961,
              902.1602544 , 902.82144427, 906.23247844, 904.42359527,
              901.37972397, 904.11587076, 903.81616055, 903.46804695,
              903.25094815, 903.75040964, 908.20420352, 906.34512064,
              905.4865916 , 903.43538028, 903.93688758, 901.91201239,
              901.95622135, 902.95834748, 903.74994562, 904.222996  ,
              904.50401344, 903.59998306, 901.52795925, 903.07495726]])

  • Indexes: (2)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='draw', length=500))

  • Attributes: (4)

    created_at :

    2025-09-29T20:06:54.612280+00:00

    arviz_version :

    0.22.0

    modeling_interface :

    bambi

    modeling_interface_version :

    0.15.0

  
  
• observed_data

  <xarray.Dataset> Size: 12kB
  Dimensions:                  (__obs__: 500, rt,response_extra_dim_0: 2)
  Coordinates:
    * __obs__                  (__obs__) int64 4kB 0 1 2 3 4 ... 496 497 498 499
    * rt,response_extra_dim_0  (rt,response_extra_dim_0) int64 16B 0 1
  Data variables:
      rt,response              (__obs__, rt,response_extra_dim_0) float64 8kB 0...
  Attributes:
      created_at:                  2025-09-29T20:06:54.613262+00:00
      arviz_version:               0.22.0
      inference_library:           numpyro
      inference_library_version:   0.19.0
      sampling_time:               15.733445
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.15.0

  xarray.Dataset

  • Dimensions:
    • __obs__: 500
    • rt,response_extra_dim_0: 2
  • Coordinates: (2)
    • __obs__
      (__obs__)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499], shape=(500,))

    • rt,response_extra_dim_0
      (rt,response_extra_dim_0)
      int64
      0 1

      array([0, 1])

  • Data variables: (1)
    • rt,response
      (__obs__, rt,response_extra_dim_0)
      float64
      0.7336 1.0 1.114 ... 0.4049 1.0

      array([[ 0.73362625,  1.        ],
             [ 1.11384964, -1.        ],
             [ 0.80583185,  1.        ],
             [ 0.80243492,  1.        ],
             [ 2.19640899,  1.        ],
             [ 0.65593803,  1.        ],
             [ 0.90178442,  1.        ],
             [ 4.80741978,  1.        ],
             [ 0.59347796,  1.        ],
             [ 1.13075411,  1.        ],
             [ 1.92198789,  1.        ],
             [ 0.90965253,  1.        ],
             [ 2.11170959,  1.        ],
             [ 2.21477485,  1.        ],
             [ 2.17002368, -1.        ],
             [ 2.209167  ,  1.        ],
             [ 0.86345994,  1.        ],
             [ 1.64354908, -1.        ],
             [ 1.05286789,  1.        ],
             [ 3.11714959, -1.        ],
      ...
             [ 1.5871619 , -1.        ],
             [ 0.52945417,  1.        ],
             [ 0.6030978 ,  1.        ],
             [ 0.80487525,  1.        ],
             [ 0.8251211 ,  1.        ],
             [ 2.62559605, -1.        ],
             [ 1.10837114,  1.        ],
             [ 3.04465795,  1.        ],
             [ 1.80810845, -1.        ],
             [ 1.22626626,  1.        ],
             [ 1.82505405,  1.        ],
             [ 0.54845285,  1.        ],
             [ 1.01013482,  1.        ],
             [ 1.86507869,  1.        ],
             [ 0.94010693,  1.        ],
             [ 1.73986757,  1.        ],
             [ 2.15558338,  1.        ],
             [ 2.28776789,  1.        ],
             [ 1.02572453, -1.        ],
             [ 0.40494695,  1.        ]])

  • Indexes: (2)
    • __obs__
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='__obs__', length=500))

    • rt,response_extra_dim_0
      PandasIndex

      PandasIndex(Index([0, 1], dtype='int64', name='rt,response_extra_dim_0'))

  • Attributes: (8)

    created_at :

    2025-09-29T20:06:54.613262+00:00

    arviz_version :

    0.22.0

    inference_library :

    numpyro

    inference_library_version :

    0.19.0

    sampling_time :

    15.733445

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.15.0

  
  

In [18]:

az.plot_trace(model.traces)
plt.tight_layout()

az.plot_pair(model.traces)
plt.tight_layout()

az.plot_trace(model.traces) plt.tight_layout() az.plot_pair(model.traces) plt.tight_layout()

We hope you find it easy to use the above example to leverage hssm to fit your own custom models.