The helper _write_post_init_model_hparams_file should take vars_weight_params instead of recomputing it.
The helper _write_post_init_model_hparams_file should take vars_weight_params instead of recomputing it.

PiperOrigin-RevId: 455704876

google/paxml,google/paxml