@Michel_Schanen and I have created a package that provides checkpointing schemes for adjoint computation using automatic differentiation/autodiff (AD) of time stepping loops. The schemes are agnostic to the ADTool being used and can be easily interfaced with any Julia AD tool.
We welcome contributions and suggestions for improvement. The repository is at:
Currently the package provides the following checkpointing schemes:
- Revolve/Binomial checkpointing 
- Periodic checkpointing
Checkpointing.jl interfaces with an AD tool through the computation of a Jacobian by implementing a
jacobian method. The following describes the interface for
using Checkpointing struct ReverseDiffADTool <: AbstractADTool end function Checkpointing.jacobian(tobedifferentiated, F_H, ::ReverseDiffADTool) return ReverseDiff.jacobian(tobedifferentiated, F_H) end
The interfaces for
Zygote.jl are implemented in
@checkpointing covers the transformation of
for loops with
1:steps ranges where
steps is the number of timesteps:
@checkpoint scheme adtool for i in 1:steps F_H = [F, F] F = advance(F_H,t,h) t += h end
adtool is one of the interface AD tools (e.g.
ReverseDiffADTool()) and scheme is a adjoint checkpointing scheme like for example revolve.
function store(x::Vector, c::Vector, t::Int64, s::Int64) c[1,s] = x c[2,s] = x c[3,s] = t return end function restore(c, i) x = [c[1,i], c[2,i]] t = c[3,i] return x, t end Revolve(steps, checkpoints, store, store; verbose=verbose)
Revolve needs the functions
restore for storing and restoring the
c[i] with variables
steps is the total number of timesteps while
checkpoints is the number of available checkpoints.
The following features are planned for development:
- Integration with
ChainRules.jlto generate rules for timestepping loops
- Online checkpointing schemes for adaptive timestepping
- Composition of checkpointing schemes
- Multi-level checkpointing schemes
 Andreas Griewank and Andrea Walther. 2000. Algorithm 799: revolve: an implementation of checkpointing for the reverse or adjoint mode of computational differentiation. ACM Trans. Math. Softw. 26, 1 (March 2000), 19–45. DOI:https://doi.org/10.1145/347837.347846