convert maxtext trained orbax checkpoint to HF checkpoint

[jwyang-google]

hi @khatwanimohit and @gobbleturk

sorry I messed up with the previous PR and closed that one. Please review this new one instead. The changes are all same.