Tests failing due to incompatibility between TFP and JAX 0.7

[slinderman]

Our tests are failing because of a compatibility issue with TFP and the newest version of JAX. See tensorflow/probability#2009.

[PaoloRanzi81]

Thanks for such a comment of shared pain... Indeed, it is several days that I am struggling to run the two jupyter notebooks of lgssm_parallel_inference.ipynb and lgssm_hmc.ipynb. The culprit in the parallel notebook is the following line: from dynamax.linear_gaussian_ssm import lgssm_smoother, parallel_lgssm_smoother. The error message is: AttributeError: jax.interpreters.xla.pytype_aval_mappings was deprecated in JAX v0.5.0 and removed in JAX v0.7.0. jax.core.pytype_aval_mappings can be used as a replacement in most cases.

I would like to implement a parallelized Kalman filter with MCMC inference for predicting physiological time-series on a short-term horizon. I have tried with tfp.sts.fit_with_hmc in Tensorflow Probability, but it does not work unfortunately (the variational inference version does work).

I have tried with the combination of Python packages suggested in here, but it still throws an error, though.

Could you kindly share a list of Python packages (and their specific versions) which works smoothly with "dynamax"? We would like a bullet-proof requirements.txt. That would be a great workaround for now.

I attach mine (the one which throws an error) below.

requirements.txt

[PaoloRanzi81]

I am replying to my own comment above. I have taken inspiration from a requirements.txt found here.

After some trial-and-error, I have found a combination of Python packages which is WORKING with the previously mentioned notebooks from dynamax. Unfortunately it runs only on CPUs... Anyhow, better than nothing! I have uploaded the WORKING requirements.txt below since it might be useful for someone facing the same problem.

I am using an Ubuntu 24.04.2 LTS operating system on a ASUS ROG Strix with x86_64 architecture and Python 3.12.3.

requirements.txt

[bantin]

Also had this problem. Was able to temporarily get around it by installing tfp-nightly (as detailed in tensorflow/probability#2009)

[PaoloRanzi81]

Thanks for the feedback. I do confirm that installing nightly builds (wheels) both for TensorFlow and TensorFlow Probability works with GPUs, as well !!!

For newbies, you can download the nightly builds wheels from repositories here and here for TF and TFP, respectively. Please choose them according to your own Operating System and Python version.

As follows, the steps I have used on Linux terminal (Ubuntu 24.04.2) with Python 3.12:

# create virtual environment called "benchmarking"
python3 -m venv benchmarking

# enter the virtual environment
source benchmarking/bin/activate

# install dynamax, matplotlib, pandas and numpy
pip install dynamax matplotlib pandas numpy

# uninstall versions installed by default by dynamax
pip uninstall tensorflow-probability jax jaxlib 

# re-install both TF and TFP wheels from nightly builds repositories
pip install "/home/paolo/miscellaneous/downloads/tf_nightly-2.21.0.dev20250801-cp312-cp312-manylinux_2_27_x86_64.whl"
pip install "/home/paolo/miscellaneous/downloads/tfp_nightly-0.26.0.dev20250808-py2.py3-none-any.whl"

# re-install JAX with GPU support
pip install -U "jax[cuda12]"