Use resolved mesh size for context parallel sharding

[huytransformer]

Description

Fixes TPU Flash/Tokamax Splash context-parallel sizing when ici_context_parallelism=-1. This PR reads the context-parallel size from the mesh in the attention path and train setup.

Tests

Ran TPU repro:

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
  run_name=cp_minus_one_repro \
  model_name=llama3.1-8b \
  attention=flash \
  use_tokamax_splash=true \
  ici_fsdp_parallelism=1 \
  ici_context_parallelism=-1 \
  packing=false \
  dataset_type=synthetic \
  steps=1 \
  enable_checkpointing=false

Checklist

• 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 50.00000% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/train_utils.py	33.33%	0 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!