Universal checkpoint for zero stage 1 (#2284)
* Refactor universal checkpointing and tensor fragments
* Formatting
* Support zero stage1; Expand TP dim
* Remove debug prints
* Detect sharded optimizer state
* Format fixes
* Encode reshaping guide
* More symbolic constants
Co-authored-by: Michael Wyatt <michaelwyatt@microsoft.com>
