:orphan: :py:mod:`tlm_adjoint.checkpoint_schedules.schedule` =================================================== .. py:module:: tlm_adjoint.checkpoint_schedules.schedule Module Contents --------------- .. py:class:: CheckpointAction(*args) A checkpointing action. .. !! processed by numpydoc !! .. py:property:: args Action parameters. .. !! processed by numpydoc !! .. py:class:: Clear(clear_ics, clear_data) A checkpointing action which clears the intermediate storage. :arg clear_ics: Whether to clear stored forward restart data. :arg clear_data: Whether to clear stored non-linear dependency data. .. !! processed by numpydoc !! .. py:property:: clear_ics Whether to clear stored forward restart data. .. !! processed by numpydoc !! .. py:property:: clear_data Whether to clear stored non-linear dependency data. .. !! processed by numpydoc !! .. py:class:: Configure(store_ics, store_data) A checkpointing action which configures the intermediate storage. :arg store_ics: Whether to store forward restart data. :arg store_data: Whether to store non-linear dependency data. .. !! processed by numpydoc !! .. py:property:: store_ics Whether to store forward restart data. .. !! processed by numpydoc !! .. py:property:: store_data Whether to store non-linear dependency data. .. !! processed by numpydoc !! .. py:class:: Forward(n0, n1) A checkpointing action which indicates forward advancement. :arg n0: The forward should advance from the start of this step. :arg n1: The forward should advance to the start of this step. .. !! processed by numpydoc !! .. py:property:: n0 The forward should advance from the start of this step. .. !! processed by numpydoc !! .. py:property:: n1 The forward should advance to the start of this step. .. !! processed by numpydoc !! .. py:class:: Reverse(n1, n0) A checkpointing action which indicates adjoint advancement. :arg n1: The adjoint should advance from the start of this step. :arg n0: The adjoint should advance to the start of this step. .. !! processed by numpydoc !! .. py:property:: n0 The adjoint should advance to the start of this step. .. !! processed by numpydoc !! .. py:property:: n1 The adjoint should advance from the start of this step. .. !! processed by numpydoc !! .. py:class:: Read(n, storage, delete) A checkpointing action which indicates loading of data from a checkpointing unit. :arg n: The step with which the loaded data is associated. :arg storage: The storage from which the data should be loaded. Either `'RAM'` or `'disk'`. :arg delete: Whether the data should be deleted from the indicated storage after it has been loaded. .. !! processed by numpydoc !! .. py:property:: n The step with which the loaded data is associated. .. !! processed by numpydoc !! .. py:property:: storage The storage from which the data should be loaded. Either `'RAM'` or `'disk'`. .. !! processed by numpydoc !! .. py:property:: delete Whether the data should be deleted from the indicated storage after it has been loaded. .. !! processed by numpydoc !! .. py:class:: Write(n, storage) A checkpointing action which indicates saving of data to a checkpointing unit. :arg n: The step with which the saved data is associated. :arg storage: The storage to which the data should be saved. Either `'RAM'` or `'disk'`. .. !! processed by numpydoc !! .. py:property:: n The step with which the saved data is associated. .. !! processed by numpydoc !! .. py:property:: storage The storage to which the data should be saved. Either `'RAM'` or `'disk'`. .. !! processed by numpydoc !! .. py:class:: EndForward A checkpointing action which indicates the end of the initial forward calculation. .. !! processed by numpydoc !! .. py:class:: EndReverse(exhausted) A checkpointing action which indicates the end of an adjoint calculation. :arg exhausted: Indicates whether the schedule has concluded. If `True` then this action should be the last action in the schedule. .. !! processed by numpydoc !! .. py:property:: exhausted Indicates whether the schedule has concluded. If `True` then this action should be the last action in the schedule. .. !! processed by numpydoc !! .. py:class:: 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. .. code-block:: python @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 :meth:`.CheckpointSchedule.finalize` method. :arg max_n: The number of steps in the initial forward calculation. If not supplied then this should later be provided by calling the :meth:`.CheckpointSchedule.finalize` method. .. !! processed by numpydoc !! .. py:property:: is_exhausted :abstractmethod: Whether the schedule has concluded. Note that some schedules permit multiple adjoint calculation, and may never conclude. .. !! processed by numpydoc !! .. py:property:: uses_disk_storage :abstractmethod: Whether the schedule may use disk storage. If `False` then no disk storage is required. .. !! processed by numpydoc !! .. py:property:: n The forward location. After executing all actions defined so far in the schedule the forward is at the start of this step. .. !! processed by numpydoc !! .. py:property:: r The number of adjoint steps advanced in the current adjoint calculation after executing all actions defined so far in the schedule. .. !! processed by numpydoc !! .. py: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. .. !! processed by numpydoc !! .. py:property:: is_running Whether the schedule is 'running' -- i.e. at least one action has been defined so far in the schedule. .. !! processed by numpydoc !! .. py:method:: iter() :abstractmethod: A generator which should be overridden in derived classes in order to define a checkpointing schedule. .. !! processed by numpydoc !! .. py:method:: finalize(n) Indicate the number of forward steps in the initial forward calculation. :arg n: The number of steps in the initial forward calculation. .. !! processed by numpydoc !!