Support checkpointing TE quantizations with new remat policies#2773

jberchtold-nvidia wants to merge 1 commit intoAI-Hypercomputer:mainAI-Hypercomputer/maxtext:mainfrom

jberchtold-nvidia:jberchtold/checkpoint-te-quantizationsjberchtold-nvidia/maxtext:jberchtold/checkpoint-te-quantizationsCopy head branch name to clipboard