[Python] How to use PyMC, PyMC-Marketing, and NumPyro together when they require conflicting...

I’m trying to build a probabilistic model in Python that requires using PyMC, PyMC-Marketing, and NumPyro in the same script.

• 
  My environment: Python 3.12, Windows 64-bit
  
  
• 
  Installed packages:
  
  • pymc (latest)
  • pymc-marketing (latest)
  • numpyro 0.18.0
  • jax (latest)

When I try to import NumPyro after installing everything, I get:

ImportError: cannot import name 'pjit_p' from 'jax.experimental.pjit'


This happens because:

• PyMC and PyMC-Marketing require a newer JAX (≥ 0.4.20+).
• NumPyro 0.18.0 still imports pjit_p, which was removed after JAX 0.4.14.
• Downgrading JAX to 0.4.14 fixes NumPyro but breaks PyMC and PyMC-Marketing.

What I’ve Tried:

• Installing JAX ≤ 0.4.14 → NumPyro works, PyMC breaks.
• Installing latest JAX (0.4.30+) → PyMC works, NumPyro import fails.
• Attempted to patch numpyro/ops/provenance.py manually, but concerned it could break inference.
• Checked for compatibility notes but couldn’t find guidance on using all three together.

Question:


Is there a known way to make PyMC, PyMC-Marketing, and NumPyro work together in one environment with Python 3.12?

• Does a newer (pre-release) NumPyro version support current JAX?
• Are there recommended workarounds (e.g., monkey-patching, subprocess isolation)?
• Or is it fundamentally impossible to use these three libraries simultaneously until NumPyro updates?

Continue reading...