jaxpi.checkpointing

Orbax checkpointing and resume helpers.

get_ckpt_path() [source]

get_ckpt_path(config)

Checkpoint root for a run: <saving.ckpt_path or cwd>/<wandb.name>/ckpt.

has_checkpoint_steps() [source]

has_checkpoint_steps(path)

Whether an Orbax checkpoint directory contains at least one saved step.

latest_time_window() [source]

latest_time_window(ckpt_path, pattern='time_window_(\\d+)')

Index of the last trained time window, based on checkpoint directories.

create_checkpoint_manager() [source]

create_checkpoint_manager(config, ckpt_path, suffix=None)

save_checkpoint() [source]

save_checkpoint(ckpt_mngr, state)

restore_checkpoint() [source]

restore_checkpoint(ckpt_mngr, state, step=None)

CustomJSONEncoder [source]

Bases: JSONEncoder

Extensible JSON https://json.org encoder for Python data structures.

Supports the following objects and types by default:

+-------------------+---------------+ | Python | JSON | +===================+===============+ | dict | object | +-------------------+---------------+ | list, tuple | array | +-------------------+---------------+ | str | string | +-------------------+---------------+ | int, float | number | +-------------------+---------------+ | True | true | +-------------------+---------------+ | False | false | +-------------------+---------------+ | None | null | +-------------------+---------------+

To extend this to recognize other objects, subclass and implement a .default() method with another method that returns a serializable object for o if possible, otherwise it should call the superclass implementation (to raise TypeError).

CustomJSONEncoder.default() [source]

CustomJSONEncoder.default(self, obj)

Implement this method in a subclass such that it returns a serializable object for o, or calls the base implementation (to raise a TypeError).

For example, to support arbitrary iterators, you could implement default like this::

def default(self, o):
    try:
        iterable = iter(o)
    except TypeError:
        pass
    else:
        return list(iterable)
    # Let the base class default method raise the TypeError
    return super().default(o)

save_config() [source]

save_config(config, workdir, name=None)