def model(self) -> ModelPyTree:
    """The full model, including parameters and nonlearnable parts."""
    return (
        pz.select(self.model_without_params)
        .at_instances_of(pz.nn.Parameter)
        .apply(lambda p: dataclasses.replace(p, value=self.params[p.name]))
    )
The full model, including parameters and nonlearnable parts.
penzai/deprecated/v1/toolshed/basic_training.py