Error in running the example code for hierarchical ddm

[TeaforTwo1011]

Hi team,

I encountered the an error while trying to run the following hierarchical model in main tutorial:
model_reg_v_angle_hier.sample(
sampler="nuts_numpyro", chains=2, cores=1, draws=1000, tune=1000
)

**Compiling...
Traceback (most recent call last):

Cell In[8], line 2
model_reg_v_angle_hier.sample(

File C:\ProgramData\miniconda3\Lib\site-packages\hssm\hssm.py:343 in sample
self._inference_obj = self.model.fit(inference_method=sampler, **kwargs)

File C:\ProgramData\miniconda3\Lib\site-packages\bambi\models.py:325 in fit
return self.backend.run(

File C:\ProgramData\miniconda3\Lib\site-packages\bambi\backend\pymc.py:96 in run
result = self._run_mcmc(

File C:\ProgramData\miniconda3\Lib\site-packages\bambi\backend\pymc.py:211 in _run_mcmc
idata = pymc.sampling_jax.sample_numpyro_nuts(

File C:\ProgramData\miniconda3\Lib\site-packages\pymc\sampling\jax.py:630 in sample_numpyro_nuts
logp_fn = get_jaxified_logp(model, negative_logp=False)

File C:\ProgramData\miniconda3\Lib\site-packages\pymc\sampling\jax.py:136 in get_jaxified_logp
logp_fn = get_jaxified_graph(inputs=model.value_vars, outputs=[model_logp])

File C:\ProgramData\miniconda3\Lib\site-packages\pymc\sampling\jax.py:129 in get_jaxified_graph
return jax_funcify(fgraph)

File C:\ProgramData\miniconda3\Lib\functools.py:909 in wrapper
return dispatch(args[0].class)(*args, **kw)

File C:\ProgramData\miniconda3\Lib\site-packages\pytensor\link\jax\dispatch\basic.py:51 in jax_funcify_FunctionGraph
return fgraph_to_python(

File C:\ProgramData\miniconda3\Lib\site-packages\pytensor\link\utils.py:738 in fgraph_to_python
compiled_func = op_conversion_fn(

File C:\ProgramData\miniconda3\Lib\functools.py:909 in wrapper
return dispatch(args[0].class)(*args, **kw)

File C:\ProgramData\miniconda3\Lib\site-packages\pytensor\link\jax\dispatch\elemwise.py:12 in jax_funcify_Elemwise
return jax_funcify(scalar_op, **kwargs)

File C:\ProgramData\miniconda3\Lib\functools.py:909 in wrapper
return dispatch(args[0].class)(*args, **kw)

File C:\ProgramData\miniconda3\Lib\site-packages\pytensor\link\jax\dispatch\scalar.py:69 in jax_funcify_ScalarOp
jax_variadic_func = getattr(jnp, op.nfunc_variadic)

File C:\ProgramData\miniconda3\Lib\site-packages\jax_src\deprecations.py:53 in getattr
raise AttributeError(f"module {module!r} has no attribute {name!r}")

AttributeError: module 'jax.numpy' has no attribute 'product'**

I tried to debug it myself but I couldn't find a solution to this, thus, your help would be greatly appreciated! Thanks a lot!

[TeaforTwo1011]

I think I have fixed it by downgrading the jax and jaxlib version down. pip install jax==0.4.13 and pip install jaxlib==0.4.13 works for me.