Bug when running end-to-end case study

[mjtriggs]

I'm trying to replicate this case study on my own machine. I believe I've followed all installation instructions correctly. I did have to manually install JAX, which could be a potential issue, but I'm not sure.

When I run the following cell, I get the below errors:

mmm.fit(
    X=X_train,
    y=y_train,
    target_accept=0.9,
    chains=4,
    draws=3_000,
    nuts_sampler="numpyro",
    random_seed=rng,
)

Error:

XlaRuntimeError: FAILED_PRECONDITION: Buffer Definition Event: Error dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sError dispatching computation: %sCpuCallback error: Traceback (most recent call last):
  File "c:\Users\matt.triggs\Anaconda3\envs\pymc_marketing_env\lib\site-packages\jax\_src\interpreters\mlir.py", line 2806, in _wrapped_callback
RuntimeError: Incorrect output dtype for return value #0: Expected: int64, Actual: int32

If anyone could provide any insight, that would be great. Really looking forward to making more use of this package.

[wd60622]

Hi @mjtriggs, did you encounter this when running the notebook as is or did you make any modifications, including switching out the dataset? Could you also give some details on your machine.

[mjtriggs]

Hi @wd60622 ,

I encountered this when running the cell specified - all previous ones work up until that point. I've made no modifications to the code, other than copying and pasting it from the web notebook into my own jupyter notebook running in VS code.

Machine Details: Windows 11 Pro running on Lenovo Laptop. Intel i7 CPU with 16GB RAM.

I'm utilising a Conda virtual environment - requirements file below:

requirements.txt