def new_opt_state_and_model_state(rng):
      """Returns optimizer and model states suitable for training a model."""
      weights, state = self._model_with_loss.init(input_signature, rng=rng)
      (slots, opt_params) = opt.tree_init(weights)
      return (OptState(weights, slots, opt_params), state)
Returns optimizer and model states suitable for training a model.
trax/supervised/trainer_lib.py