Announcement: big_vision is transitioning from jax.pmap to jax.jit.

[akolesnikoff]

In the coming 1-2 weeks big_vision is expected to transition from pmap-based parallelism to jit-parallelism. This will enable more flexible parallelisation strategies, including, but not limited to, ZERO* and fully-sharded data-parallel (aka fsdp) training.

This transition may temporarily break project-specific code (or we will just remove such code). If you want to read/run the old code, please see the table at the end of the README for the project-specific commits to sync to.

[akolesnikoff]

Transition is done