Refresh artefacts and code for the second review release (part 3)

Refresh artefacts; evaluate through the paper's harness path

Fix Orbax >=0.12 checkpoint restore (concrete sharding)

Anonymous artefact release (part 2)