Problem with jax tree_multimap when importing pybamm package

Question

I am trying to import pybamm and its throwing the following error at meenter image description here

Following the suggestion of a previous post I have tried

!pip install "jax<=0.3.16" "jaxlib<=0.3.16"

but it remains the same error.

Answer 1

It looks like tree_multimap was removed in 0.3.16 (see CHANGELOG), so I believe you'd need to install <=0.3.15 for that import to work