How to initialize a train state from a PyTree of parameters?

Hi,

I am just wondering is there an efficient and straightforward function/method to initialize a train state from a PyTree of model parameters? Currently I can find `Checkpointer.restore` to restore a train state from checkpoints `TrainStateInitializer.from_scratch` to intialize a train state from scratch.

Currently I am trying to load a dictionary of model parameters (without training steps and optimizer states, etc.) and would like to do continual training on these parameters.

Thanks!

Start with TrainStateInitializer.from_scratch and Checkpointer.restore to understand the existing initialization and restoration paths. Determine whether the requested PyTree-only initialization is supported or needs a new entry point, with completion defined as continuing training from supplied parameters without checkpoint-only state.

Written by the indexing model from the issue text.