Serialise: Save & Load Runs

save writes a HybridModel or a FitResult to a directory, with the solver, and optionally the training config and objective. load reads it back into a SavedRun.

Loading needs a template. Equinox cannot rebuild module shapes from the leaf file, so load(path, like=model) takes a model built the same way and fills its leaves. The template also supplies simulate_fn and state_to_output, which are code and are not saved. save_predictors / load_predictors do the same for a bare predictor tree.

Quick links

• save
• load
• SavedRun
• save_predictors
• load_predictors

---

save()

from jaxhybridmodels.serialise import save  ·  also re-exported as jaxhybridmodels.save

save(
    path: 'str | Path',
    run: 'HybridModel | FitResult',
    config: 'OptaxTrainingConfig | EvosaxTrainingConfig | None' = None,
    objective: 'Objective | None' = None,
    extras: 'dict[str, Any] | None' = None,
) -> None

Write a model, or a fit with its history, to the directory path.

run is a HybridModel or the FitResult from fit; a FitResult also saves its history, best_step and stop_reason. config and objective record how the model was trained, and extras is any JSON-encodable dict of your own. Read it back with load.

Layout written

path/predictors.eqx Binary leaves of the predictors pytree (save_predictors output). path/metadata.jsontimestamp (UTC, ISO 8601), version, predictors (the structural fingerprint load checks like against), output_channel_names, solver (SolverConfig.to_dict()), config ({"kind": "optax" | "evosax", "fields": {...}}), objective, history (the four FitHistory series), best_step, stop_reason and extras. Callables are described, not pickled (see the module docstring).

The directory is created, parents included. Existing files are overwritten.

Source

---

load()

from jaxhybridmodels.serialise import load  ·  also re-exported as jaxhybridmodels.load

load(path: 'str | Path', *, like: 'HybridModel') -> 'SavedRun'

Read back a run written by save.

like is a model built the same way as the saved one: same predictor architecture and static configuration, same output_channel_names. It supplies simulate_fn and state_to_output, which are not saved, and the solver. The returned SavedRun.model is like with the saved predictors. The saved solver record holds class names only (constructor arguments such as RecursiveCheckpointAdjoint(checkpoints=7) are not kept), so like.solver is used as is, with a RuntimeWarning when it differs from the record.

Raises ValueError when the saved package version differs from the installed one, when like does not match the saved predictors or channel names, or when a saved config has a field the current class does not know.

Source

---

SavedRun

from jaxhybridmodels.results import SavedRun  ·  also re-exported as jaxhybridmodels.SavedRun

SavedRun(
    model: 'HybridModel',
    history: 'FitHistory',
    best_step: 'int | None',
    stop_reason: 'StopReason | None',
    config: 'OptaxTrainingConfig | EvosaxTrainingConfig | dict[str, Any] | None',
    objective: 'Objective | dict[str, Any] | None',
    extras: 'dict[str, Any]',
) -> None

A run read back by load.

model is the like= model with the saved predictors; it keeps like's solver. history, best_step and stop_reason are the saved FitResult's; a saved bare model leaves them as an empty FitHistory and None.

config and objective come back as typed objects. A section that names a callable (a custom loss, a penalty_fn, an optimizer factory, an lr schedule) cannot be rebuilt, so it comes back as the raw dict with each callable as {"__callable__": "module.qualname"}; rebind those fields yourself. extras is the dict passed to save.

Source

---

save_predictors()

from jaxhybridmodels.serialise import save_predictors

save_predictors(path: 'str | Path', predictors: 'Any') -> 'None'

Write predictors to path via eqx.tree_serialise_leaves.

eqx.tree_serialise_leaves walks the leaves the same way whatever the container shape, writing a flat binary stream of np.save-encoded leaves. The caller picks the file extension; save uses .eqx.

Source

---

load_predictors()

from jaxhybridmodels.serialise import load_predictors

load_predictors(path: 'str | Path', *, like: 'Any') -> 'Any'

Restore a predictors pytree from path, using like as the skeleton.

like must share the saved pytree's container shape and per-leaf static configuration, such as the same in_size and depth on every MLPPredictor leaf. The file's values overwrite its array leaves; its static fields supply the structure equinox needs. like is not mutated.

Source