Incompatible jax or jaxlib version in tools/inference/inference_pipeline.ipynb

In the 3rd code block of the inference pipeline notebook, it has this error

RuntimeError Traceback (most recent call last)

in <cell line: 1>() ----> 1 import jax 2 import jax.numpy as jnp 3 4 # check how many devices are available 5 jax.local_device_count()

4 frames

/usr/local/lib/python3.10/dist-packages/jax/_src/lib/init.py in check_jaxlib_version(jax_version, jaxlib_version, minimum_jaxlib_version) 67 f'incompatible with jax version {jax_version}. Please ' 68 'update your jax and/or jaxlib packages.') ---> 69 raise RuntimeError(msg) 70 71 return _jaxlib_version

RuntimeError: jaxlib version 0.4.7 is newer than and incompatible with jax version 0.3.25. Please update your jax and/or jaxlib packages.