mmm_roas.ipynb model fitting fails with runtime exception for me under Python 3.10.0

[scottlyden]

Apart from adding this code block ...

%pip install -U arviz
%pip install -U graphviz
%pip install -U jax
%pip install -U numpyro
%pip install -U pymc
%pip install -U pymc-marketing

... I haven't made any changes to the notebook.

The exception message is:

XlaRuntimeError: FAILED_PRECONDITION: Buffer Definition Event: Error dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sCpuCallback error: Traceback (most recent call last):
  File "c:\users\xxxx yyyy\pycharmprojects\p3_10\venv\lib\site-packages\jax\_src\interpreters\mlir.py", line 2781, in _wrapped_callback
RuntimeError: Incorrect output dtype for return value #0: Expected: int64, Actual: int32

The problem appears to be connected with the numpyro sampler; when I change the sampler to the default (NUTS), the fitting runs fine.

Any ideas much appreciated.

If anyone is able to run the notebook as-is, would you please post the version of Python and above packages that you used?

I am using:

pymc==5.18.2
pymc-marketing==0.10.0
pymc3==3.11.4
jax==0.4.35
jaxlib==0.4.35
numpyro==0.15.3

with Python 3.10.0

Thanks!

[AnantChandra]

This should do it :)

import jax
jax.config.update("jax_enable_x64", True)

[williambdean]

These might also be related: #1515 #1427