jaxpi.utils
Small utilities: pytree flattening, update schedules, schedule-free eval params.
flatten_pytree() [source]
python
flatten_pytree(pytree)get_eval_params() [source]
python
get_eval_params(state, schedule_free)Parameters to use for inference.
With the schedule-free optimizer wrapper, the parameters held by the train state are the training iterates; evaluation (error logging, IC propagation between time windows, final prediction) should use the schedule-free averaged parameters instead.
create_update_scheduler() [source]
python
create_update_scheduler(every: int = 100, start: int = 0) -> Callable[[int], bool]Build and return a step-checker for the given schedule.