Fixing a parameter with backend='jax' fails with a bare AssertionError in specifyshape
- 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
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