msgpack_load

Contents

msgpack_load#

class braintools.file.msgpack_load(filename, target=None, parallel=True, mismatch='error', max_size=None, verbose=False)[source]#

Load the checkpoint from the given path using the msgpack library.

This function is rewritten from the Flax APIs (google/flax).

Parameters:

  • filename (str | PathLike) – checkpoint file to restore from.

  • target (Any | None) – the object to restore the state into. If None, the state is returned as a nested dict (with all dict keys as strings; see Notes).

  • parallel (bool) – whether to read large seekable checkpoints with multiple threads.

  • mismatch (Literal['error', 'warn', 'ignore']) – How to handle mismatches between target and state dict (dict keys, list/tuple length, namedtuple fields, unit, and array shape). ‘error’ (default): raise ValueError on mismatch ‘warn’: issue warning and keep the target’s value for the mismatch ‘ignore’: silently keep the target’s value for the mismatch

  • max_size (int | None) – Maximum decoded checkpoint size in bytes (None = no limit).

  • verbose (bool) – Whether to print progress information (default: False).

Returns:

out – If target is given, a structure isomorphic to target with the restored leaf data (State leaves are restored in place). If target is None, the raw nested state dict read from the file.

Return type:

PyTree

Raises:

  • ValueError – If filename does not exist, or on a mismatch when mismatch='error'.

  • InvalidCheckpointPath – If the checkpoint data is corrupt or cannot be decoded.

Notes

Without a target, dict keys come back as strings and dict subclasses (e.g. OrderedDict) become plain dicts, because keys are stringified on save. Provide a target to recover the original key types and structure.

Contents