Upgrade orbax deps?

[ihh]

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`). ***

[alvarobartt]

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? 🤗

[ihh]

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.

[VolodyaCO]

Same here. The required jax versions are very old.

[alvarobartt]

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!