error running the sample function

[TabithaSteendam]

Hi all,

I am doing the "HSSM model with regression: Case 1: One parameter is a Regression Target" part of the main tutorial and run into an error when running the .sample function. AttributeError: module 'jax.numpy' has no attribute 'product

This is the code I copied:

`# Set up trial by trial parameters
intercept = 0.3
x = np.random.uniform(-1, 1, size=1000)
y = np.random.uniform(-1, 1, size=1000)
v = intercept + (0.8 * x) + (0.3 * y)

true_values = np.column_stack(
    [v, np.repeat([[1.5, 0.5, 0.5, 0.0]], axis=0, repeats=1000)]
)

# Get mode simulations
obs_ddm_reg_v = simulator(true_values, model="ddm", n_samples=1)

dataset_reg_v = pd.DataFrame(
    {
        "rt": obs_ddm_reg_v["rts"].flatten(),
        "response": obs_ddm_reg_v["choices"].flatten(),
        "x": x,
        "y": y,
    }
)

dataset_reg_v

model_reg_v = hssm.HSSM(
    data=dataset_reg_v,
    include=[
        {
            "name": "v",
            "prior": {
                "Intercept": {"name": "Uniform", "lower": -3.0, "upper": 3.0},
                "x": {"name": "Uniform", "lower": -1.0, "upper": 1.0},
                "y": {"name": "Uniform", "lower": -1.0, "upper": 1.0},
            },
            "formula": "v ~ 1 + x + y",
            "link": "identity",
        }
    ],
)

infer_data_reg_v = model_reg_v.sample(
    sampler="nuts_numpyro", chains=1, cores=1, draws=500, tune=500
)`

When running:

`infer_data_reg_v = model_reg_v.sample(
    sampler="nuts_numpyro", chains=1, cores=1, draws=500, tune=500
)`

I get the following error message:Traceback (most recent call last):

`Cell In[3], line 1
   infer_data_reg_v = model_reg_v.sample(

 File ~\.conda\envs\HSSM\lib\site-packages\hssm\hssm.py:343 in sample
   self._inference_obj = self.model.fit(inference_method=sampler, **kwargs)

 File ~\.conda\envs\HSSM\lib\site-packages\bambi\models.py:325 in fit
   return self.backend.run(

 File ~\.conda\envs\HSSM\lib\site-packages\bambi\backend\pymc.py:96 in run
   result = self._run_mcmc(

 File ~\.conda\envs\HSSM\lib\site-packages\bambi\backend\pymc.py:211 in _run_mcmc
   idata = pymc.sampling_jax.sample_numpyro_nuts(

 File ~\.conda\envs\HSSM\lib\site-packages\pymc\sampling\jax.py:630 in sample_numpyro_nuts
   logp_fn = get_jaxified_logp(model, negative_logp=False)

 File ~\.conda\envs\HSSM\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 ~\.conda\envs\HSSM\lib\site-packages\pymc\sampling\jax.py:129 in get_jaxified_graph
   return jax_funcify(fgraph)

 File ~\.conda\envs\HSSM\lib\functools.py:889 in wrapper
   return dispatch(args[0].__class__)(*args, **kw)

 File ~\.conda\envs\HSSM\lib\site-packages\pytensor\link\jax\dispatch\basic.py:51 in jax_funcify_FunctionGraph
   return fgraph_to_python(

 File ~\.conda\envs\HSSM\lib\site-packages\pytensor\link\utils.py:738 in fgraph_to_python
   compiled_func = op_conversion_fn(

 File ~\.conda\envs\HSSM\lib\functools.py:889 in wrapper
   return dispatch(args[0].__class__)(*args, **kw)

 File ~\.conda\envs\HSSM\lib\site-packages\pytensor\link\jax\dispatch\elemwise.py:12 in jax_funcify_Elemwise
   return jax_funcify(scalar_op, **kwargs)

 File ~\.conda\envs\HSSM\lib\functools.py:889 in wrapper
   return dispatch(args[0].__class__)(*args, **kw)

 File ~\.conda\envs\HSSM\lib\site-packages\pytensor\link\jax\dispatch\scalar.py:69 in jax_funcify_ScalarOp
   jax_variadic_func = getattr(jnp, op.nfunc_variadic)

 File ~\.conda\envs\HSSM\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'`

Does anyone know what causes this?
If it helps, I installed hssm in a new environment and these are my package versions:

`arviz==0.14.0

asttokens @ file:///home/conda/feedstock_root/build_artifacts/asttokens_1698341106958/work

backports.functools-lru-cache @ 
file:///home/conda/feedstock_root/build_artifacts/backports.functools_lru_cache_1687772187254/work

bambi==0.12.0

cachetools==5.3.2

certifi==2023.7.22

cftime==1.6.3

charset-normalizer==3.3.2

cloudpickle @ file:///home/conda/feedstock_root/build_artifacts/cloudpickle_1697464713350/work

colorama @ file:///home/conda/feedstock_root/build_artifacts/colorama_1666700638685/work

coloredlogs==15.0.1

comm @ file:///home/conda/feedstock_root/build_artifacts/comm_1691044910542/work

cons==0.4.6

contourpy==1.2.0

cycler==0.12.1

Cython==3.0.5

debugpy @ file:///D:/bld/debugpy_1660619096890/work

decorator @ file:///home/conda/feedstock_root/build_artifacts/decorator_1641555617451/work

etuples==0.3.9

exceptiongroup @ file:///home/conda/feedstock_root/build_artifacts/exceptiongroup_1692026125334/work

executing @ file:///home/conda/feedstock_root/build_artifacts/executing_1698579936712/work

fastprogress==1.0.3

filelock==3.13.1

flatbuffers==23.5.26

fonttools==4.44.0

formulae==0.5.1

fsspec==2023.10.0

graphviz==0.20.1

hddm-wfpt==0.1.2

HSSM==0.1.5

huggingface-hub==0.15.1

humanfriendly==10.0

idna==3.4

importlib-metadata @ file:///home/conda/feedstock_root/build_artifacts/importlib-metadata_1688754491823/work

ipykernel @ file:///D:/bld/ipykernel_1673894762113/work

ipython @ file:///D:/bld/ipython_1698846796959/work

jax==0.4.20

jaxlib==0.4.20

jedi @ file:///home/conda/feedstock_root/build_artifacts/jedi_1696326070614/work

joblib==1.3.2

jupyter_client @ file:///home/conda/feedstock_root/build_artifacts/jupyter_client_1699283905679/work

jupyter_core @ file:///D:/bld/jupyter_core_1698673856358/work

kiwisolver==1.4.5

logical-unification==0.4.6

matplotlib==3.8.1

matplotlib-inline @ file:///home/conda/feedstock_root/build_artifacts/matplotlib-inline_1660814786464/work

miniKanren==1.0.3

ml-dtypes==0.3.1

mpmath==1.3.0

multipledispatch==1.0.0

nest-asyncio @ file:///home/conda/feedstock_root/build_artifacts/nest-asyncio_1697083700168/work

netCDF4==1.6.5

numpy==1.23.4

numpyro==0.12.1

onnx==1.15.0

onnxruntime==1.16.2

opt-einsum==3.3.0

packaging @ file:///home/conda/feedstock_root/build_artifacts/packaging_1696202382185/work

pandas==2.1.3

parso @ file:///home/conda/feedstock_root/build_artifacts/parso_1638334955874/work

pickleshare @ file:///home/conda/feedstock_root/build_artifacts/pickleshare_1602536217715/work

Pillow==10.1.0

platformdirs @ file:///home/conda/feedstock_root/build_artifacts/platformdirs_1699715570510/work

prompt-toolkit @ file:///home/conda/feedstock_root/build_artifacts/prompt-toolkit_1699963054032/work

protobuf==4.25.0

psutil @ file:///D:/bld/psutil_1653089388778/work

pure-eval @ file:///home/conda/feedstock_root/build_artifacts/pure_eval_1642875951954/work

Pygments @ file:///home/conda/feedstock_root/build_artifacts/pygments_1691408637400/work

pymc==5.6.1

pyparsing==3.1.1

pyreadline3==3.4.1

pytensor==2.12.3

python-dateutil @ file:///home/conda/feedstock_root/build_artifacts/python-dateutil_1626286286081/work

pytz==2023.3.post1

pywin32==303

PyYAML==6.0.1

pyzmq @ file:///D:/bld/pyzmq_1660329059232/work

requests==2.31.0

scikit-learn==1.3.2

scipy==1.10.1

six @ file:///home/conda/feedstock_root/build_artifacts/six_1620240208055/work

spyder-kernels @ file:///D:/bld/spyder-kernels_1642109532429/work

ssm-simulators==0.3.2

stack-data @ file:///home/conda/feedstock_root/build_artifacts/stack_data_1669632077133/work

sympy==1.12

threadpoolctl==3.2.0

toolz==0.12.0

tornado @ file:///D:/bld/tornado_1656937966227/work

tqdm==4.66.1

traitlets @ file:///home/conda/feedstock_root/build_artifacts/traitlets_1698671135544/work

typing_extensions @ file:///home/conda/feedstock_root/build_artifacts/typing_extensions_1695040754690/work

tzdata==2023.3

urllib3==2.1.0

wcwidth @ file:///home/conda/feedstock_root/build_artifacts/wcwidth_1699959196938/work

xarray==2023.10.1

xarray-einstats==0.6.0

zipp @ file:///home/conda/feedstock_root/build_artifacts/zipp_1695255097490/work`

[digicosmos86]

Hi,

It could be that lower versions of PyMC are not compatible with higher versions of Jax. Can you try downgrading Jax and Jaxlib?

Thanks!
Paul

[TabithaSteendam]

That worked, thank you!