Error: cannot import name 'pytree' from 'jaxlib'

[lviano]

Thanks a lot for sharing the code of your interesting paper!

I am following the procedure to install the requirements but when running python run.py, I run into this error cannot import name 'pytree' from 'jaxlib'.

Does it run on your side with the specified versions of jaxlib and jax ?

Thanks in advance for your help.

Best,
Luca