Orbax updates causing issues with an old version of jax

[AdrienBolling]

Got the following error while training an agent when it came to saving the agents training.

WARNING:absl:SaveArgs.aggregate is deprecated, please use custom TypeHandler (https://orbax.readthedocs.io/en/latest/custom_handlers.html#typehandler) or contact Orbax team to migrate before August 1st, 2024.
Traceback (most recent call last):
  File "/Users/adrien.bolling/Applications/PyCharm Professional Edition.app/Contents/plugins/python/helpers-pro/pydevd_asyncio/pydevd_nest_asyncio.py", line 138, in run
    return loop.run_until_complete(task)
  File "/Users/adrien.bolling/Applications/PyCharm Professional Edition.app/Contents/plugins/python/helpers-pro/pydevd_asyncio/pydevd_nest_asyncio.py", line 243, in run_until_complete
    return f.result()
  File "/Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.9/lib/python3.9/asyncio/futures.py", line 201, in result
    raise self._exception
  File "/Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.9/lib/python3.9/asyncio/tasks.py", line 256, in __step
    result = coro.send(None)
  File "/Users/adrien.bolling/PycharmProjects/technician-assignment/.venv/lib/python3.9/site-packages/orbax/checkpoint/base_pytree_checkpoint_handler.py", line 620, in async_save
    commit_futures = await self.async_save(*args, **kwargs)  # pytype: disable=bad-return-type
  File "/Users/adrien.bolling/PycharmProjects/technician-assignment/.venv/lib/python3.9/site-packages/orbax/checkpoint/base_pytree_checkpoint_handler.py", line 528, in async_save
    save_args = jax.tree.map(
  File "/Users/adrien.bolling/PycharmProjects/technician-assignment/.venv/lib/python3.9/site-packages/jax/_src/deprecations.py", line 53, in getattr
    raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax' has no attribute 'tree'
python-BaseException

This is very likely due to a mismatch between jax and orbax versions, especially considering the recent changes to orbax-checkpoints.

I temporarily fixed it by upgrading jax to 0.4.30 and keeping orbax to 0.5.16, then forcibly downgrading numpy to 1.21.0 and scipy to 1.10.1

It seems to me like this will need a refactor to either use legacy functionnalities that will soon be deprecated, or a global version bump for jax, orbax, numpy, and scipy (the former 2 which in turn break other functions due to behaviour changes, I think it stems from PettingZoo requirements)

Also be advised that such problems will only increase in the future due to the following :

WARNING:absl:SaveArgs.aggregate is deprecated, please use custom TypeHandler (https://orbax.readthedocs.io/en/latest/custom_handlers.html#typehandler) or contact Orbax team to migrate before August 1st, 2024.

[ffelten]

Hi @AdrienBolling, would you mind submitting a PR as you seem to have fixed it on your side? :)

[AdrienBolling]

@ffelten As mentionned I fixed it by forcibly downgrading certain packages after install, I didn't adapt the code to the new API, I can look into it if needed however

[ffelten]

@ffelten As mentionned I fixed it by forcibly downgrading certain packages after install, I didn't adapt the code to the new API, I can look into it if needed however

That would be cool to fix if you have time :-)

[AdrienBolling]

Will do