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

[zhangzx-uiuc]

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!