def initial_state(
      cls,
      model: ModelPyTree,
      optimizer_def: optax.GradientTransformation,
      root_rng: PRNGKeyArray,
      loss_fn_state: LossStatePyTree = None,
  ):
    """Constructs the initial training state.

    Args:
      model: The model being optimized.
      optimizer_def: The optax optimizer to use.
      root_rng: Base random number generator; used in combination with step to
        derive per-step random numbers.
      loss_fn_state: Optional initial state for the loss function.

    Returns:
      An initial training state.
    """
    # Extract the parameters.
    pz.nn.check_no_duplicated_parameters(model)
    param_selection = pz.select(model).at_instances_of(pz.nn.Parameter)
    params = {
        param.name: param.value for param in param_selection.get_sequence()
    }
    model_without_params = param_selection.at(lambda p: p.value).set(
        pz.NotInThisPartition()
    )
    # Derive opt state from parameters.
    opt_state = optimizer_def.init(params)
    return cls(
        step=0,
        root_rng=root_rng,
        params=params,
        model_without_params=model_without_params,
        opt_state=opt_state,
        loss_fn_state=loss_fn_state,
        optimizer_def=optimizer_def,
    )
Constructs the initial training state.

    Args:
      model: The model being optimized.
      optimizer_def: The optax optimizer to use.
      root_rng: Base random number generator; used in combination with step to
        derive per-step random numbers.
      loss_fn_state: Optional initial state for the loss function.

    Returns:
      An initial training state.
    
penzai/deprecated/v1/toolshed/basic_training.py