Numba TypingError in custom statespace AR(3) model

Questions v5
1

Hi Everyone. I have been working through Making a custom StateSpace Model.ipynb present in pymc-extras/notebook and ran into

TypingError: No implementation of function elemwise(float64, float64, array(float64, 0d, C))

and message Vectorized inputs must be arrays. I initialized the model with mode="JAX" but passed the argument compile_kwargs={"mode": "NUMBA"} to pm.sample() according to the reference notebook. Is this the root cause for this error? My sampling works fine when I do not pass any argument for it and it uses the default mode.

Is mixing JAX and Numba causing this failure? Also, what is the recommended backend for custom StateSpace models?

Here is the link for the notebook that I have been working on to provide the context notebookStateSpace

Any help would be appreciated!