diff --git a/main/.documenter-siteinfo.json b/main/.documenter-siteinfo.json index 6a4d69e..98d6970 100644 --- a/main/.documenter-siteinfo.json +++ b/main/.documenter-siteinfo.json @@ -1 +1 @@ -{"documenter":{"julia_version":"1.10.4","generation_timestamp":"2024-08-04T01:10:22","documenter_version":"1.5.0"}} \ No newline at end of file +{"documenter":{"julia_version":"1.10.4","generation_timestamp":"2024-08-23T21:10:00","documenter_version":"1.6.0"}} \ No newline at end of file diff --git a/main/assets/documenter.js b/main/assets/documenter.js index b2bdd43..82252a1 100644 --- a/main/assets/documenter.js +++ b/main/assets/documenter.js @@ -77,30 +77,35 @@ require(['jquery'], function($) { let timer = 0; var isExpanded = true; -$(document).on("click", ".docstring header", function () { - let articleToggleTitle = "Expand docstring"; - - debounce(() => { - if ($(this).siblings("section").is(":visible")) { - $(this) - .find(".docstring-article-toggle-button") - .removeClass("fa-chevron-down") - .addClass("fa-chevron-right"); - } else { - $(this) - .find(".docstring-article-toggle-button") - .removeClass("fa-chevron-right") - .addClass("fa-chevron-down"); +$(document).on( + "click", + ".docstring .docstring-article-toggle-button", + function () { + let articleToggleTitle = "Expand docstring"; + const parent = $(this).parent(); + + debounce(() => { + if (parent.siblings("section").is(":visible")) { + parent + .find("a.docstring-article-toggle-button") + .removeClass("fa-chevron-down") + .addClass("fa-chevron-right"); + } else { + parent + .find("a.docstring-article-toggle-button") + .removeClass("fa-chevron-right") + .addClass("fa-chevron-down"); - articleToggleTitle = "Collapse docstring"; - } + articleToggleTitle = "Collapse docstring"; + } - $(this) - .find(".docstring-article-toggle-button") - .prop("title", articleToggleTitle); - $(this).siblings("section").slideToggle(); - }); -}); + parent + .children(".docstring-article-toggle-button") + .prop("title", articleToggleTitle); + parent.siblings("section").slideToggle(); + }); + } +); $(document).on("click", ".docs-article-toggle-button", function (event) { let articleToggleTitle = "Expand docstring"; @@ -110,7 +115,7 @@ $(document).on("click", ".docs-article-toggle-button", function (event) { debounce(() => { if (isExpanded) { $(this).removeClass("fa-chevron-up").addClass("fa-chevron-down"); - $(".docstring-article-toggle-button") + $("a.docstring-article-toggle-button") .removeClass("fa-chevron-down") .addClass("fa-chevron-right"); @@ -119,7 +124,7 @@ $(document).on("click", ".docs-article-toggle-button", function (event) { $(".docstring section").slideUp(animationSpeed); } else { $(this).removeClass("fa-chevron-down").addClass("fa-chevron-up"); - $(".docstring-article-toggle-button") + $("a.docstring-article-toggle-button") .removeClass("fa-chevron-right") .addClass("fa-chevron-down"); diff --git a/main/index.html b/main/index.html index 803227b..c1d23a0 100644 --- a/main/index.html +++ b/main/index.html @@ -1,2 +1,2 @@ -Home · Checkpointing.jl

Checkpointing

Checkpointing.jl provides checkpointing schemes for adjoint computations using automatic differentiation (AD) of time-stepping loops. Currently, we support the macro @checkpoint_struct, which differentiates and checkpoints a struct used in a while or for the loop with a UnitRange.

Each loop iteration is differentiated using Enzyme.jl. We rely on ChainRulesCore.jl to integrate with AD tools applied to the code outside of the loop.

The schemes are agnostic to the AD tool being used and can be easily interfaced with any Julia AD tool. Currently, the package supports:

Scheme

  • Revolve/Binomial checkpointing [1]
  • Periodic checkpointing
  • Online r=2 checkpointing for a while loops with a priori unknown number of iterations [2]

Rules

Storage

  • ArrayStorage: Stores all checkpoints values in an array of type Array
  • HDF5Storage: Stores all checkpoints values in an HDF5 file

Limitations

  • Currently, the package only supports UnitRange ranges in for loops. We will add range types on a per-need basis. Please, open an issue if you need support for a specific range type.
  • We only support Enzyme as the differentiation tool of the loop body. This is due to our strict requirement for a mutation-enabled AD tool in our projects. However, there is no fundamental reason why we could not support other AD tools. Please, open an issue if you need support for a specific AD tool.
  • We don't support any activity analysis. This implies that loop iterators have to be part of the checkpointed struct if they are used in the loop body. Currently, we store the entire struct at each checkpoint. This is not necessary, and we will add support for storing only the required fields in the future.

Future

The following features are planned for development:

  • Support checkpoints on GPUs

Quick Start

API

References

[1] 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

+Home · Checkpointing.jl

Checkpointing

Checkpointing.jl provides checkpointing schemes for adjoint computations using automatic differentiation (AD) of time-stepping loops. Currently, we support the macro @checkpoint_struct, which differentiates and checkpoints a struct used in a while or for the loop with a UnitRange.

Each loop iteration is differentiated using Enzyme.jl. We rely on ChainRulesCore.jl to integrate with AD tools applied to the code outside of the loop.

The schemes are agnostic to the AD tool being used and can be easily interfaced with any Julia AD tool. Currently, the package supports:

Scheme

  • Revolve/Binomial checkpointing [1]
  • Periodic checkpointing
  • Online r=2 checkpointing for a while loops with a priori unknown number of iterations [2]

Rules

Storage

  • ArrayStorage: Stores all checkpoints values in an array of type Array
  • HDF5Storage: Stores all checkpoints values in an HDF5 file

Limitations

  • Currently, the package only supports UnitRange ranges in for loops. We will add range types on a per-need basis. Please, open an issue if you need support for a specific range type.
  • We only support Enzyme as the differentiation tool of the loop body. This is due to our strict requirement for a mutation-enabled AD tool in our projects. However, there is no fundamental reason why we could not support other AD tools. Please, open an issue if you need support for a specific AD tool.
  • We don't support any activity analysis. This implies that loop iterators have to be part of the checkpointed struct if they are used in the loop body. Currently, we store the entire struct at each checkpoint. This is not necessary, and we will add support for storing only the required fields in the future.

Future

The following features are planned for development:

  • Support checkpoints on GPUs

Quick Start

API

References

[1] 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

diff --git a/main/lib/checkpointing/index.html b/main/lib/checkpointing/index.html index a036668..08b580a 100644 --- a/main/lib/checkpointing/index.html +++ b/main/lib/checkpointing/index.html @@ -1,6 +1,6 @@ -API · Checkpointing.jl

API

Decorator macros for loops

Checkpointing.@checkpoint_struct — Macro
@checkpoint_struct(
+API · Checkpointing.jl
API
Decorator macros for loops
Checkpointing.@checkpoint_struct — Macro
@checkpoint_struct(
     alg,
     model,
     loop,
-)
This macro is supposed to be only used in conjunction with ChainRules. It does not initialize the shadowcopy. Apply the checkpointing scheme alg on the loop loop expression. model is the primal struct. shadowmodel contains the adjoints and is created here.  It is supposed to be initialized by ChainRules.
source
Supported Schemes
Checkpointing.Revolve — Type
Revolve
This is a Julia adaptation of the functionality of Revolve; see Alg. 799 published by Griewank et al. A minor extension is the  optional bundle parameter that allows to treat as many loop iterations in one tape/adjoint sweep. If bundle is 1, the default, then the behavior is that of Alg. 799.
source
Supported Storages
Developer variables for implementing new schemes
Checkpointing.ActionFlag — Type
ActionFlag
Each checkpointing algorithm currently uses the same ActionFlag type for setting the next action in the checkpointing scheme none: no action store: store a checkpoint now equivalent to TAKESHOT in Alg. 79 restore: restore a checkpoint now equivalent to RESTORE in Alg. 79 forward: execute iteration(s) forward equivalent to ADVANCE in Alg. 79 firstuturn: tape iteration(s); optionally leave to return later;  and (upon return) do the adjoint(s) equivalent to FIRSTTURN in Alg. 799 uturn: tape iteration(s) and do the adjoint(s) equivalent to YOUTURN in Alg. 79 done: we are done with adjoining the loop equivalent to the terminate enum value in Alg. 79
source
Checkpointing.Action — Type
Action
Stores the state of the checkpointing scheme after an action is taken.     * actionflag is the next action     * iteration is number of iterations for move forward     * startiteration is the loop step to start from     * cpnum is the checkpoint index number
source

+)

This macro is supposed to be only used in conjunction with ChainRules. It does not initialize the shadowcopy. Apply the checkpointing scheme alg on the loop loop expression. model is the primal struct. shadowmodel contains the adjoints and is created here. It is supposed to be initialized by ChainRules.

source

Supported Schemes

Checkpointing.Revolve — Type
Revolve

This is a Julia adaptation of the functionality of Revolve; see Alg. 799 published by Griewank et al. A minor extension is the optional bundle parameter that allows to treat as many loop iterations in one tape/adjoint sweep. If bundle is 1, the default, then the behavior is that of Alg. 799.

source

Supported Storages

Developer variables for implementing new schemes

Checkpointing.ActionFlag — Type
ActionFlag

Each checkpointing algorithm currently uses the same ActionFlag type for setting the next action in the checkpointing scheme none: no action store: store a checkpoint now equivalent to TAKESHOT in Alg. 79 restore: restore a checkpoint now equivalent to RESTORE in Alg. 79 forward: execute iteration(s) forward equivalent to ADVANCE in Alg. 79 firstuturn: tape iteration(s); optionally leave to return later; and (upon return) do the adjoint(s) equivalent to FIRSTTURN in Alg. 799 uturn: tape iteration(s) and do the adjoint(s) equivalent to YOUTURN in Alg. 79 done: we are done with adjoining the loop equivalent to the terminate enum value in Alg. 79

source
Checkpointing.Action — Type
Action

Stores the state of the checkpointing scheme after an action is taken. * actionflag is the next action * iteration is number of iterations for move forward * startiteration is the loop step to start from * cpnum is the checkpoint index number

source
diff --git a/main/quickstart/7cbbfea4.svg b/main/quickstart/34b0827a.svg similarity index 83% rename from main/quickstart/7cbbfea4.svg rename to main/quickstart/34b0827a.svg index f2a09fc..a996df5 100644 --- a/main/quickstart/7cbbfea4.svg +++ b/main/quickstart/34b0827a.svg @@ -1,44 +1,44 @@ - + - + - + - + - + - - - - - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/main/quickstart/index.html b/main/quickstart/index.html index 09eb067..0815fbe 100644 --- a/main/quickstart/index.html +++ b/main/quickstart/index.html @@ -54,4 +54,4 @@ return heat.Tnext, dheat.Tnext[2:end-1] end
heat (generic function with 1 method)

Plot function values:

tsteps = 500
-T, dT = heat(Revolve{Heat}(tsteps,4), tsteps)
([20.0, 19.28670708893302, 18.5762559822926, 17.865804875652184, 17.163811520281293, 16.461818164910405, 15.773698196009292, 15.085578227108181, 14.416422194357889, 13.747266161607596  …  0.0010844643164675824, 0.0008681767992163476, 0.0007135124080400669, 0.0005588480168637862, 0.000445160866081522, 0.00033147371529925775, 0.0002427191223882328, 0.00015396452947720788, 7.698226473860394e-5, 0.0], [0.0356607964401121, 0.07117950265889618, 0.10669762026127153, 0.14179285030017058, 0.17688683371117586, 0.2112871477986924, 0.24568541302418895, 0.27913548144214495, 0.31258246870379675, 0.3448485923782427  …  0.3448485923782427, 0.31258246870379675, 0.27913548144214495, 0.24568541302418895, 0.2112871477986924, 0.17688683371117586, 0.14179285030017058, 0.10669762026127153, 0.07117950265889618, 0.0356607964401121])

Plot gradient with respect to sum(T):

plot(dT)
Example block output +T, dT = heat(Revolve{Heat}(tsteps,4), tsteps)
([20.0, 19.28670708893302, 18.5762559822926, 17.865804875652184, 17.163811520281293, 16.461818164910405, 15.773698196009292, 15.085578227108181, 14.416422194357889, 13.747266161607596  …  0.0010844643164675824, 0.0008681767992163476, 0.0007135124080400669, 0.0005588480168637862, 0.000445160866081522, 0.00033147371529925775, 0.0002427191223882328, 0.00015396452947720788, 7.698226473860394e-5, 0.0], [0.0356607964401121, 0.07117950265889618, 0.10669762026127153, 0.14179285030017058, 0.17688683371117586, 0.2112871477986924, 0.24568541302418895, 0.27913548144214495, 0.31258246870379675, 0.3448485923782427  …  0.3448485923782427, 0.31258246870379675, 0.27913548144214495, 0.24568541302418895, 0.2112871477986924, 0.17688683371117586, 0.14179285030017058, 0.10669762026127153, 0.07117950265889618, 0.0356607964401121])

Plot gradient with respect to sum(T):

plot(dT)
Example block output diff --git a/main/rules/index.html b/main/rules/index.html index 0bef9ef..f792d38 100644 --- a/main/rules/index.html +++ b/main/rules/index.html @@ -13,4 +13,4 @@ model_input::MT, shadowmodel::MT, condition::Function -) where {MT} end +) where {MT} end diff --git a/main/schemes/index.html b/main/schemes/index.html index 2fa37d2..7fd26a3 100644 --- a/main/schemes/index.html +++ b/main/schemes/index.html @@ -13,4 +13,4 @@ model_input::MT, shadowmodel::MT, condition::Function -) where {MT}

The reverse pass for the checkpointing scheme will be called with the loop body body, the scheme alg, the model input model_input, the initialized shadow model shadowmodel, and the loop range range or the loop condition condition.

+) where {MT}

The reverse pass for the checkpointing scheme will be called with the loop body body, the scheme alg, the model input model_input, the initialized shadow model shadowmodel, and the loop range range or the loop condition condition.

diff --git a/main/storage/index.html b/main/storage/index.html index 7ec83e4..ef908c7 100644 --- a/main/storage/index.html +++ b/main/storage/index.html @@ -2,4 +2,4 @@ Storage · Checkpointing.jl

Storage

The checkpoint storage types are derived from AbstractStorage 

using Checkpointing
 struct MyStorage{MT} <: AbstractStorage where {MT} end

and have to implement the following functions.

  • A constructor invoked by the user
function MyStorage{MT}(n::Int) where {MT} end
  • Set and get functions for the storage type
Base.getindex(storage::MyStorage{MT}, i::Int) where {MT}
 Base.setindex!(storage::MyStorage{MT}, value, i::Int) where {MT}
  • Size and dimension functions
Base.size(storage::MyStorage{MT}) where {MT}
-Base.ndims(storage::MyStorage{MT}) where {MT}
+Base.ndims(storage::MyStorage{MT}) where {MT}