Issue with mmm.fit() in Version 0.11.0: XlaRuntimeError

Questions v5
,
1

Dear PyMC-Marketing Team,

I recently upgraded to version 0.11.0 and encountered an issue when running code that worked flawlessly in version 0.10.0. Specifically, the following code now fails:

mmm.fit(X=X, y=y, chains=4, target_accept=0.85, nuts_sampler="blackjax", random_seed=rng)

Even when changing the nuts_sampler parameter to "numpyro", the issue persists. Here’s the error message I receive:

XlaRuntimeError: INTERNAL: Compute error: Error dispatching computation: %sCpuCallback error: Traceback (most recent call last):
  File "C:\Users\terier99\Anaconda3\envs\marketing_env\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

I double-checked the input data (X and y) and confirmed that none of the columns are int32—everything is either float64 or int64. Additionally, as mentioned, the same data worked perfectly fine in version 0.10.0.

Could this be a bug related to the changes in the recent version? Or is there something I should modify in my setup to accommodate updates in version 0.11.0?

Thank you in advance for your assistance! Please let me know if you need any additional information to troubleshoot this issue.

1 Like
2

Hi can you raise the issue on GitHub for better visibility. It could be with the latest PyMC version being used in the latest release

3

Fortunately, I no longer need to raise this issue on GitHub. I was able to resolve the problem by adding the following code to my script:

import pytensor  
pytensor.config.floatX = "float32"

This adjustment resolved the XlaRuntimeError I encountered. I hope this might help others facing a similar issue.

4

I have opened an issue here.