[JAX] Upstream checkpointing of quantizations in TE/JAX to MaxText

[jberchtold-nvidia]

MaxText upstream already supports TE quantization, but for optimal performance for NVFP4 and MXFP8, quantization checkpointing support should be upstreamed to MaxText to prevent 2x quantization in fwd and rematerialized in backward,

This work was initially started here but needs continuing: AI-Hypercomputer/maxtext#2773

[jberchtold-nvidia]

Haixin Liu @hx89 has offered to help upstream this. Adding this comment since his username isn't showing up in the assignee list option and I don't have access to the repo settings to add him manually