fix(jax): honor training.stat_file for model statistics #5870

Summary

JAX training carries training.stat_file through configuration, but the entrypoint currently invokes statistics computation/loading without passing the configured file path. The option is therefore accepted but appears to be ignored.

Scope

• Resolve the configured statistics path consistently with the training working directory.
• Pass it into JAX model statistics loading/computation.
• Reuse existing statistics when compatible and write newly computed statistics safely.
• Validate model/type-map/descriptor compatibility before accepting a file.
• Define chief-only file I/O and synchronization for distributed training.
• Close HDF5/file resources on success and failure.

Acceptance criteria

• A configured compatible statistics file is loaded instead of recomputing statistics.
• A missing target file is created at the configured location when writing is allowed.
• Incompatible or malformed files produce actionable errors.
• Only the designated rank performs mutable file I/O in distributed runs, and all ranks receive identical statistics.
• Tests cover load, create, incompatibility, and distributed ownership paths.

Refs #5756.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh