safejax icon indicating copy to clipboard operation
safejax copied to clipboard

Published 20 hours ago • verified publisher icon alvarobartt

Upgrade orbax deps?

Open ihh opened this issue 2 years ago • 5 comments

Installing with pip3 fails with the following error

      *** Orbax is a namespace, and not a standalone package. For model checkpointing and exporting utilities, please install `orbax-checkpoint` and `orbax-export` respectively (instead of `orbax`). ***

ihh avatar Sep 06 '23 23:09 ihh

Hi @ihh thanks for reporting this issue! I'm happy to go ahead with the upgrade of orbax and/or add orbax-export, would you like to contribute with a PR for that? :hugs:

alvarobartt avatar Sep 07 '23 09:09 alvarobartt

Thanks @alvarobartt. I've given it a try (see incoming PR), but unfortunately I am still unable to use safejax myself (or to test it) since I am encountering many dependency conflicts beyond just this issue. Many of these may be because of my development environment (I am working on a Mac) or my own naivete (I'm new to Python) but I am posting the PR anyway as a starting point for discussion.

ihh avatar Sep 07 '23 23:09 ihh

Same here. The required jax versions are very old.

VolodyaCO avatar Oct 23 '24 21:10 VolodyaCO

Hi here @VolodyaCO I'm glad that someone's using safejax (or at least trying to) 🤗

Could you please let me know which are the ideal dependency versions you'd like to see pinned for JAX, Flax, etc.? Thanks in advance!

alvarobartt avatar Oct 24 '24 08:10 alvarobartt