Use setup_training_state to fix standalone checkpoint unit tests

[xibinliu]

Description

Fix the checkpoint restore issue in standalone checkpoint test:

• Delegate restoration and initialization to setup_training_state:
  Replaced manual abstract state evaluation and load_state_if_possible calls in standalone_checkpointer.py with setup_training_state(...). This aligns standalone checkpointer with standard training and ensures proper NNX state unboxing, sharding resolution, and fallback handling.
• Prevent synthetic entropy injection on restored checkpoints:
  Updated setup_initial_state and setup_training_state to return a was_restored: bool flag.
  In standalone_checkpointer.py, guarded add_entropy_to_checkpoint(state) so synthetic entropy is only injected on clean runs, preserving restored optimizer moments (mu and nu) when testing checkpoint resumption.
• Fix NNX checkpoint save layout:
  When running under pure_nnx=True, converted state using to_linen_checkpoint_dict(state.to_pure_dict()) before calling save_checkpoint, guaranteeing that saved checkpoints conform to the standard Linen disk layout and restoring cleanly without ShapeDtypeStruct type errors.

FIXES: b/532183169

Tests

standalone_dl_ckpt_test.py test pass:

mkdir -p ~/my_tpu_logs && TPU_LOG_DIR=~/my_tpu_logs GLOG_log_dir=~/my_tpu_logs PYTHONPATH=. GLOG_logtostderr=1 python -m pytest tests/integration/standalone_dl_ckpt_test.py -s -o log_cli=true -o log_cli_level=INFO

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 42.85714% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/standalone_checkpointer.py	0.00%	6 Missing ⚠️
...rc/maxtext/utils/generate_param_only_checkpoint.py	66.66%	1 Missing ⚠️
src/maxtext/utils/train_utils.py	50.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!