tlm_adjoint.checkpoint_schedules.schedule

Module Contents

class tlm_adjoint.checkpoint_schedules.schedule.CheckpointAction(*args)

A checkpointing action.

property args

Action parameters.

class tlm_adjoint.checkpoint_schedules.schedule.Clear(clear_ics, clear_data)

A checkpointing action which clears the intermediate storage.

Parameters:

  • clear_ics – Whether to clear stored forward restart data.

  • clear_data – Whether to clear stored non-linear dependency data.

property clear_ics

Whether to clear stored forward restart data.

property clear_data

Whether to clear stored non-linear dependency data.

class tlm_adjoint.checkpoint_schedules.schedule.Configure(store_ics, store_data)

A checkpointing action which configures the intermediate storage.

Parameters:

  • store_ics – Whether to store forward restart data.

  • store_data – Whether to store non-linear dependency data.

property store_ics

Whether to store forward restart data.

property store_data

Whether to store non-linear dependency data.

class tlm_adjoint.checkpoint_schedules.schedule.Forward(n0, n1)

A checkpointing action which indicates forward advancement.

Parameters:

  • n0 – The forward should advance from the start of this step.

  • n1 – The forward should advance to the start of this step.

property n0

The forward should advance from the start of this step.

property n1

The forward should advance to the start of this step.

class tlm_adjoint.checkpoint_schedules.schedule.Reverse(n1, n0)

A checkpointing action which indicates adjoint advancement.

Parameters:

  • n1 – The adjoint should advance from the start of this step.

  • n0 – The adjoint should advance to the start of this step.

property n0

The adjoint should advance to the start of this step.

property n1

The adjoint should advance from the start of this step.

class tlm_adjoint.checkpoint_schedules.schedule.Read(n, storage, delete)

A checkpointing action which indicates loading of data from a checkpointing unit.

Parameters:

  • n – The step with which the loaded data is associated.

  • storage – The storage from which the data should be loaded. Either ‘RAM’ or ‘disk’.

  • delete – Whether the data should be deleted from the indicated storage after it has been loaded.

property n

The step with which the loaded data is associated.

property storage

The storage from which the data should be loaded. Either ‘RAM’ or ‘disk’.

property delete

Whether the data should be deleted from the indicated storage after it has been loaded.

class tlm_adjoint.checkpoint_schedules.schedule.Write(n, storage)

A checkpointing action which indicates saving of data to a checkpointing unit.

Parameters:

  • n – The step with which the saved data is associated.

  • storage – The storage to which the data should be saved. Either ‘RAM’ or ‘disk’.

property n

The step with which the saved data is associated.

property storage

The storage to which the data should be saved. Either ‘RAM’ or ‘disk’.

class tlm_adjoint.checkpoint_schedules.schedule.EndForward

A checkpointing action which indicates the end of the initial forward calculation.

class tlm_adjoint.checkpoint_schedules.schedule.EndReverse(exhausted)

A checkpointing action which indicates the end of an adjoint calculation.

Parameters:

exhausted – Indicates whether the schedule has concluded. If True then this action should be the last action in the schedule.

property exhausted

Indicates whether the schedule has concluded. If True then this action should be the last action in the schedule.

class tlm_adjoint.checkpoint_schedules.schedule.CheckpointSchedule(max_n=None)

A checkpointing schedule.

Actions in the schedule are accessed by iterating over elements, and actions may be implemented using single-dispatch functions. e.g.

@functools.singledispatch
def action(cp_action):
    raise TypeError(f"Unexpected checkpointing action: {cp_action}")

@action.register(Forward)
def action_forward(cp_action):
    logger.debug(f"forward: forward advance to {cp_action.n1:d}")

# ...

for cp_action in cp_schedule:
    action(cp_action)
    if isinstance(cp_action, EndReverse):
        break

Schedules control an intermediate storage, which buffers forward restart data for forward restart checkpoints, and which stores non-linear dependency data either for storage in checkpointing units or for immediate use by the adjoint. For details see

  • James R. Maddison, ‘Step-based checkpointing with high-level algorithmic differentiation’, Journal of Computational Science 82, 102405, 2024, doi: 10.1016/j.jocs.2024.102405

In ‘offline’ schedules, where the number of steps in the forward calculation is initially known, this should be provided using the max_n argument on instantiation. In ‘online’ schedules, where the number of steps in the forward calculation is initially unknown, the number of forward steps should later be provided using the CheckpointSchedule.finalize() method.

Parameters:

max_n – The number of steps in the initial forward calculation. If not supplied then this should later be provided by calling the CheckpointSchedule.finalize() method.

abstract property is_exhausted

Whether the schedule has concluded. Note that some schedules permit multiple adjoint calculation, and may never conclude.

abstract property uses_disk_storage

Whether the schedule may use disk storage. If False then no disk storage is required.

property n

The forward location. After executing all actions defined so far in the schedule the forward is at the start of this step.

property r

The number of adjoint steps advanced in the current adjoint calculation after executing all actions defined so far in the schedule.

property max_n

The number of forward steps in the initial forward calculation. May return None if this has not yet been provided to the scheduler.

property is_running

Whether the schedule is ‘running’ – i.e. at least one action has been defined so far in the schedule.

abstract iter()

A generator which should be overridden in derived classes in order to define a checkpointing schedule.

finalize(n)

Indicate the number of forward steps in the initial forward calculation.

Parameters:

n – The number of steps in the initial forward calculation.