Serialization

Opake keeps training state in explicit values (clip state, noise state, optimizer state, schedules, Accountant / DpProcess, …). Use opake.serialization for a single flat dict[str, Any] suitable for torch.save / torch.load.

Restore is template-driven: pass a freshly constructed object of the same shape as at save time; keys present in the checkpoint overwrite leaves, and missing keys keep template values (forward compatibility when new fields appear).

Dispatch resolves a leaf by exact type and then by __mro__, so a subclass of a registered type — a custom torch.Tensor subclass, for example — is serialized by the nearest base-class handler rather than dropped. (torch.nn.Parameter is common enough to get its own exact-type handler that preserves the subclass and its requires_grad flag on restore, so it does not rely on the __mro__ fallback.) A leaf that is neither registered nor a generic container (dataclass / NamedTuple / tuple / list / mapping) nor a primitive raises InputTypeError on both save and restore instead of being silently skipped. Genuinely inert leaves that the template reproduces (vendor structure handles such as optree.PyTreeSpec) are declared with opake.serialization.register_template_restored; nothing is written for them and the template supplies them on load.

PerGroup checkpoints through the same API. Path keys serialize as their str() form under the structural walker — for example, groups.('a',) for a flat leaf, or groups.('layer', 'weight') for a nested path — alongside values.<group_name>. When DP bias correction is enabled on Adam-family optimizers, phi is a path-keyed dict from opt.init so from_state_dict round-trips without resetting φ. NumPy ndarray leaves are supported alongside torch.Tensor.

Domain pages with examples: Optimizers, Accounting.

Array leaf policy

torch.Tensor, torch.nn.Parameter and numpy.ndarray leaves resolve three attributes against the template, each under its own rule:

• Shape — must match the template exactly; a mismatch raises CheckpointError for every array leaf kind. Broadcast-compatible shapes (a length-1 buffer against a length-d slot) are what make the check load-bearing: they would otherwise restore without an error and keep training against state the accountant no longer prices.
• Dtype — taken from the template, and the checkpoint value is cast to it. The template carries the live compute dtype, so resuming an fp32 checkpoint in bf16 works.
• Device — taken from the template for tensors, so a checkpoint read with map_location="cpu" lands on the training device. ndarray leaves have no device.

A leaf absent from the checkpoint keeps the template's value; that is the forward-compatibility rule above, not a mismatch. An error raised while restoring a registered leaf carries the offending key as an exception note.

API surface (no per-type state_dict())

Opake centralizes (de)serialization in opake.serialization.state_dict and opake.serialization.from_state_dict. Registered types (including Accountant) are written and restored only through these module-level functions — there are no instance methods named state_dict / from_state_dict on those classes.

Minimal round trip for an Accountant:

from opake.accounting import Accountant, identity
from opake.serialization import from_state_dict, state_dict

acct = Accountant() | identity()
flat = state_dict(acct)
acct2 = from_state_dict(Accountant(), flat)

The same pattern applies to clip state, noise state, functional optimizer state, and any other value that flows through the DP training loop. Custom types may register handlers with opake.serialization.register_serializer, or — when a leaf carries no run state and the template reproduces it — declare it inert with opake.serialization.register_template_restored.

serialization

Functional serialization for Opake explicit state objects.

Flattens tensor trees, NumPy arrays, dataclasses, named tuples, sequences, and string-keyed dicts into a flat dict[str, Any] suitable for torch.save / torch.load.

Restore is template-driven: supply a freshly-initialised object of the same shape as at save time; each path present in the dict overwrites the corresponding leaf. Missing paths keep the template (forward compatibility when new fields appear).

Sub-packages may register custom serializers with :func:register_serializer. Lookup is by exact type and then by __mro__, so a subclass of a registered type (a custom torch.Tensor subclass, say) uses the nearest base class handler. torch.Tensor, numpy.ndarray, and torch.nn.Parameter handlers register automatically when opake-engine is loaded; nn.Parameter gets its own exact-type handler (preserving the subclass and requires_grad) rather than relying on the __mro__ fallback. opake-accounting registers PLD process types; stack wheels register their state objects.

A leaf that is neither registered nor a generic container nor a primitive raises TypeError on both save and restore rather than being dropped. Declare the genuinely inert ones — vendor structure handles and the like — with :func:register_template_restored.

Serializer

Bases: Protocol

Pair of save / load callables registered with the dispatcher.

Source code in packages/opake-base/src/opake/api/base/serialization/_types.py

@runtime_checkable
class Serializer(Protocol):
    """Pair of save / load callables registered with the dispatcher."""

    def state_dict(self, obj: Any) -> dict[str, Any]: ...

    def from_state_dict(self, template: Any, sd: Mapping[str, Any]) -> Any: ...

from_state_dict(template, sd)

Rebuild from sd using template for shape and omitted leaves.

For types registered with :func:register_serializer, the template argument selects the handler; handlers may ignore it when the dict is self-describing (e.g. DpProcess subclasses). Generic Python container shapes come from the template; missing leaves keep the template's value (forward compatibility when new fields appear).

Source code in packages/opake-base/src/opake/api/base/serialization/_dispatch.py

def from_state_dict(template: Any, sd: Mapping[str, Any]) -> Any:
    """Rebuild from ``sd`` using ``template`` for shape and omitted leaves.

    For types registered with :func:`register_serializer`, the
    ``template`` argument selects the handler; handlers may ignore it
    when the dict is self-describing (e.g. ``DpProcess`` subclasses).
    Generic Python container shapes come from the template; missing
    leaves keep the template's value (forward compatibility when new
    fields appear).
    """
    return _walk_load(template, sd, "")

lookup_serializer(typ)

Return the serializer pair registered for exactly typ, or None.

Source code in packages/opake-base/src/opake/api/base/serialization/_registry.py

def lookup_serializer(
    typ: type[Any],
) -> tuple[StateDictFn, FromStateDictFn] | None:
    """Return the serializer pair registered for exactly ``typ``, or ``None``."""
    return _REGISTRY.get(typ)

register_serializer(typ, state_dict_fn, from_state_dict_fn)

Register a serializer pair for typ.

Re-registering the same type silently replaces the existing pair — this lets test code stub out registrations without unregistering first. Production code should not register the same type twice.

Source code in packages/opake-base/src/opake/api/base/serialization/_registry.py

def register_serializer(
    typ: type[Any],
    state_dict_fn: StateDictFn,
    from_state_dict_fn: FromStateDictFn,
) -> None:
    """Register a serializer pair for ``typ``.

    Re-registering the same type silently replaces the existing pair —
    this lets test code stub out registrations without unregistering
    first. Production code should not register the same type twice.
    """
    _REGISTRY[typ] = (state_dict_fn, from_state_dict_fn)

register_template_restored(typ)

Declare typ inert: nothing is saved, the template supplies it.

The escape hatch for leaves that carry no run state and are fully determined by a freshly-built template — vendor structure handles such as optree.PyTreeSpec, for example. Without a registration the dispatcher raises on an unrecognised non-container leaf rather than dropping it, so this is how a wheel says "dropping this one is intentional".

Do not use it for anything a resumed run needs to reproduce: the saved checkpoint will not contain the value, and a template built with different settings restores those different settings.

Source code in packages/opake-base/src/opake/api/base/serialization/_registry.py

def register_template_restored(typ: type[Any]) -> None:
    """Declare ``typ`` inert: nothing is saved, the template supplies it.

    The escape hatch for leaves that carry no run state and are fully
    determined by a freshly-built template — vendor structure handles
    such as ``optree.PyTreeSpec``, for example. Without a registration
    the dispatcher raises on an unrecognised non-container leaf rather
    than dropping it, so this is how a wheel says "dropping this one is
    intentional".

    Do not use it for anything a resumed run needs to reproduce: the
    saved checkpoint will not contain the value, and a template built
    with different settings restores those different settings.
    """
    register_serializer(typ, lambda _obj: {}, lambda template, _sd: template)

resolve_serializer(typ)

Return the serializer pair for typ or its nearest registered base.

Exact type first, then __mro__ order, so nn.Parameter and any other torch.Tensor subclass resolve to the torch.Tensor handler. Returns None when nothing in the MRO is registered.

Source code in packages/opake-base/src/opake/api/base/serialization/_registry.py

def resolve_serializer(
    typ: type[Any],
) -> tuple[StateDictFn, FromStateDictFn] | None:
    """Return the serializer pair for ``typ`` or its nearest registered base.

    Exact type first, then ``__mro__`` order, so ``nn.Parameter`` and any
    other ``torch.Tensor`` subclass resolve to the ``torch.Tensor``
    handler.  Returns ``None`` when nothing in the MRO is registered.
    """
    pair = _REGISTRY.get(typ)
    if pair is not None:
        return pair
    for base in typ.__mro__[1:]:
        pair = _REGISTRY.get(base)
        if pair is not None:
            return pair
    return None

state_dict(obj)

Serialise obj to a flat dict[str, Any].

Registered types use their exact-type handler; everything else is walked structurally (dataclass, NamedTuple, tuple, list, mapping, primitives). Opaque non-containers are dropped — the load is template-driven and reads them back from the template.

Source code in packages/opake-base/src/opake/api/base/serialization/_dispatch.py

def state_dict(obj: Any) -> dict[str, Any]:
    """Serialise ``obj`` to a flat ``dict[str, Any]``.

    Registered types use their exact-type handler; everything else is
    walked structurally (dataclass, NamedTuple, tuple, list, mapping,
    primitives). Opaque non-containers are dropped — the load is
    template-driven and reads them back from the template.
    """
    out: dict[str, Any] = {}
    _walk_save(obj, "", out)
    return out