tree_multimap deprecated in favor of tree_map

google-deepmind / rlax

https://rlax.readthedocs.io

Apache License 2.0

1.24k stars 85 forks source link

tree_multimap deprecated in favor of tree_map #100

Closed rohan-mehta-1024 closed 2 years ago

rohan-mehta-1024 commented 2 years ago

I get this error when trying to import rlax, I think because tree_multimap has been deprecated?

ImportError: cannot import name 'tree_multimap' from 'jax.tree_util' (/home/rohanmehta/anaconda3/envs/py39/lib/python3.9/site-packages/jax/tree_util.py)

katebaumli commented 2 years ago

Hi Rohan, are you using the latest version of rlax? This should be fixed in version 0.1.4

rohan-mehta-1024 commented 2 years ago

Yes, I didn't realize I was using an older version. Sorry about that, and thank you :)