Prepare for shape inference in trainer.
Prepare for shape inference in trainer.

- Make initialize_model_state takes a sample_inputs
- Make initialize_partitioned_model_state takes a sample_inputs

PiperOrigin-RevId: 463493255

google/praxis,google/praxis