Serialization

jax_morph.TrajectoryRecord

TrajectoryRecord(
    history: BaseState,
    dt: Array,
    provenance: Mapping[str, object],
)

A loaded complete trajectory and the metadata needed to interpret it.

Attributes:

• history –

  Stacked complete BaseState history with leading frame axis n_steps + 1.

• dt –

  Scalar JAX array containing the macro-step duration.

• provenance –

  JSON-safe mapping stored alongside the history.

---

jax_morph.save_model

save_model(
    path: str | PathLike[str], model: object
) -> None

Save a model's supported numeric PyTree leaves into a versioned artifact.

The artifact stores numeric leaves only. Static model structure is recorded for validation and a caller-supplied template provides the executable Python objects at load time.

Parameters:

• path (str | PathLike[str]) –

  Destination artifact path.

• model (object) –

  Model or other supported Equinox/JAX PyTree to serialize.

Raises:

• PersistenceError –

  If the artifact cannot be written.

• TypeError –

  If the model contains unsupported payload leaves or static structure.

---

jax_morph.load_model

load_model(
    path: str | PathLike[str], like: object
) -> object

Restore a model artifact's numeric leaves into an exactly matching like template.

Parameters:

• path (str | PathLike[str]) –

  Source artifact path.

• like (object) –

  Matching live model template that supplies executable static objects.

Returns:

• object –

  A copy of like with the artifact's numeric leaves restored.

Raises:

• PersistenceError –

  If the artifact is malformed, incompatible, or does not match like.

• TypeError –

  If like contains unsupported static structure.

---

jax_morph.save_state

save_state(
    path: str | PathLike[str], state: BaseState
) -> None

Save one self-describing, unstacked generated state snapshot.

Parameters:

• path (str | PathLike[str]) –

  Destination artifact path.

• state (BaseState) –

  Unstacked generated state to serialize.

Raises:

• PersistenceError –

  If the artifact cannot be written.

• TypeError –

  If the state or its schema is invalid or cannot be serialized.

• ValueError –

  If state is a stacked history.

---

jax_morph.load_state

load_state(
    path: str | PathLike[str],
    *,
    space: object | None = None,
) -> BaseState

Load one self-describing state, reconstructing a built-in saved space when possible.

Parameters:

• path (str | PathLike[str]) –

  Source artifact path.

• space (object | None, default: None ) –

  Replacement (displacement, shift) pair for a custom saved space. Defaults to None.

Returns:

• BaseState –

  Reconstructed unstacked BaseState instance.

Raises:

• PersistenceError –

  If the artifact is malformed, incompatible, or needs an unavailable custom space.

---

jax_morph.save_trajectory

save_trajectory(
    path: str | PathLike[str],
    history: BaseState,
    dt: object,
    *,
    provenance: Mapping[str, object] | None = None,
) -> None

Save one complete adjacent-macro-step history and its scalar dt metadata.

Parameters:

• path (str | PathLike[str]) –

  Destination artifact path.

• history (BaseState) –

  Complete stacked BaseState history with leading frame axis n_steps + 1.

• dt (object) –

  Finite real macro-step duration matching adjacent history.t frames.

• provenance (Mapping[str, object] | None, default: None ) –

  JSON-safe metadata to store with the trajectory. Defaults to None.

Raises:

• PersistenceError –

  If the artifact cannot be written.

• TypeError –

  If history is invalid or provenance is not JSON-safe.

• ValueError –

  If history is not a complete stacked trajectory or dt is inconsistent.

---

jax_morph.load_trajectory

load_trajectory(
    path: str | PathLike[str],
    *,
    space: object | None = None,
) -> TrajectoryRecord

Load a complete self-describing history, scalar JAX dt, and provenance mapping.

Parameters:

• path (str | PathLike[str]) –

  Source artifact path.

• space (object | None, default: None ) –

  Replacement (displacement, shift) pair for a custom saved space. Defaults to None.

Returns:

• TrajectoryRecord –

  TrajectoryRecord containing the stacked history, macro-step duration, and provenance.

Raises:

• PersistenceError –

  If the artifact is malformed, incompatible, or needs an unavailable custom space.