lnccbrown / lnccbrown/HSSM

Fixing a parameter with backend='jax' fails with a bare AssertionError in specifyshape

Open
#1,092 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
124
Forks
24
Avg merge
19h 32m
Merged PRs (30d)
60

Description

Fixing a parameter to a constant works on the `"jax"` backend up to the point of
sampling, then fails inside PyTensor's JAX dispatch with a **bare
`AssertionError` carrying no message**.

The parameter-fixing logic itself is correct — numpyro reports sampling
`[a, z, v]`, correctly excluding the fixed `t`. The failure happens when the
compiled graph runs.

### Reproducer

```python
import jax.numpy as jnp
import hssm
from hssm.config import ModelConfig

def logp(data, v, a, z, t):
rt = data[0]
mu = jnp.log(a) - jnp.log(v**2 + 0.25) + t
lr = jnp.log(jnp.maximum(rt, 1e-6))
return -lr - 0.5 * ((lr - mu) / 0.45) ** 2 - 0.8

df = hssm.simulate_data(model="ddm", theta=[0.9, 1.3, 0.5, 0.3], size=200, random_state=1)
cfg = ModelConfig(
response=["rt", "response"], list_params=["v", "a", "z", "t"], choices=(-1, 1),
bounds={"v": (-3., 3.), "a": (0.3, 2.5), "z": (0.1, 0.9), "t": (0., 1.)},
backend="jax",
)
common = dict(data=df, model="probe", model_config=cfg, loglik=logp,
loglik_kind="approx_differentiable", p_outlier=0)

# works
hssm.HSSM(**common).sample(sampler="numpyro", draws=20, tune=20, chains=1,
cores=1, progressbar=False, random_seed=1)

# raises
hssm.HSSM(**common, t=0.3).sample(sampler="numpyro", draws=20, tune=20, chains=1,
cores=1, progressbar=False, random_seed=1)
```

```
File ".../pymc/sampling/jax.py", line 123, in logp_fn_wrap
return logp_fn(*x)[0]
File "/var/folders/.../tmpjy4go8od", line 61, in jax_funcified_fgraph
tensor_variable_29 = specifyshape(tensor_variable_28, tensor_constant_10)
File ".../pytensor/link/jax/dispatch/shape.py", line 98, in specifyshape
assert x.ndim == len(shape)
AssertionError
```

### Two separate problems

1. Fixing a parameter does not work on this path.
2. The error is a bare `AssertionError` with an empty message, so there is
nothing to diagnose from. Even if the underlying limitation stays, a clear
error at model-build time would save a lot of time.

### Scope I could not determine

I could not compare against `backend="pytensor"`, because
`backend="pytensor"` + `loglik_kind="approx_differentiable"` + a callable
`loglik` is rejected by design with a `ValueError`. So I cannot say whether this
is specific to the JAX backend or affects fixed parameters with custom log
likelihoods more generally.

**Versions:** hssm 0.4.0, pymc 6.2.0, pytensor 3.2.3, jax 0.11.0,
numpyro 0.21.0, ssm-simulators 0.13.2, Python 3.12

Contributor guide

Open the contributing guide

Research direction

Reproduce the fixed-parameter case using HSSM.HSSM(..., t=0.3).sample with the shown JAX configuration, then inspect the JAX sampling path and the failing specifyshape call in pytensor/link/jax/dispatch/shape.py. Check whether the issue is specific to custom approximate-differentiable likelihoods; done means fixed parameters either sample successfully or produce a clear model-build error instead of a bare AssertionError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.