From bc37fa0a4d6c5b5ff6253cf6f0a25190ee0bc5aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 24 Sep 2026 11:26:32 +0200 Subject: [PATCH] Move examples into doctests and fix the in-place Jacobian example (fixes #114) #128 ported `A_mul_Bc!(out, A, B)` (= `A * B'`) to `mul!(out, a + b, a * b)`, dropping the adjoint, so `g!` no longer matched `f`. The examples are now doctests in the Gradients and Jacobians sections of the API docs, so CI checks them, and `examples/` is removed. Co-Authored-By: Claude Opus 5.5 (1M context) --- README.md | 13 +-- docs/Project.toml | 6 +- docs/src/api.md | 224 +++++++++++++++++++++++++++++++++++++++++++ docs/src/index.md | 4 +- examples/gradient.jl | 58 ----------- examples/jacobian.jl | 66 ------------- 6 files changed, 238 insertions(+), 133 deletions(-) delete mode 100644 examples/gradient.jl delete mode 100644 examples/jacobian.jl diff --git a/README.md b/README.md index 36a105c0..29c0f08b 100644 --- a/README.md +++ b/README.md @@ -5,8 +5,6 @@ [![Build status](https://github.com/JuliaDiff/ReverseDiff.jl/workflows/CI/badge.svg)](https://github.com/JuliaDiff/ReverseDiff.jl/actions) [![codecov.io](https://codecov.io/github/JuliaDiff/ReverseDiff.jl/coverage.svg?branch=master)](https://codecov.io/github/JuliaDiff/ReverseDiff.jl?branch=master) -[**See ReverseDiff Usage Examples**](https://github.com/JuliaDiff/ReverseDiff.jl/tree/master/examples) - ReverseDiff is a fast and compile-able tape-based **reverse mode automatic differentiation (AD)** that implements methods to take **gradients**, **Jacobians**, **Hessians**, and higher-order derivatives of native Julia functions (or any callable object, really). @@ -45,13 +43,16 @@ of them (as far as I know at the time of this writing): - higher-order `map` and `broadcast` optimizations - it's well tested -...and, simply put, it's fast (for gradients, at least). Using the code from `examples/gradient.jl`: +...and, simply put, it's fast (for gradients, at least). For example: ```julia -julia> using BenchmarkTools, Pkg +julia> using BenchmarkTools, ReverseDiff + +julia> f(a, b) = sum(a' * b + a * b'); + +julia> compiled_f_tape = ReverseDiff.compile(ReverseDiff.GradientTape(f, (rand(100, 100), rand(100, 100)))); -# this script defines f and ∇f! -julia> include(joinpath(Pkg.dir("ReverseDiff"), "examples/gradient.jl")); +julia> ∇f!(results, inputs) = ReverseDiff.gradient!(results, compiled_f_tape, inputs); julia> a, b = rand(100, 100), rand(100, 100); diff --git a/docs/Project.toml b/docs/Project.toml index 4b8c876b..d771fc8f 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,9 +1,13 @@ [deps] +DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" [compat] +DiffResults = "1" Documenter = "1" +LinearAlgebra = "1" [sources] -ReverseDiff = { path = ".." } \ No newline at end of file +ReverseDiff = { path = ".." } diff --git a/docs/src/api.md b/docs/src/api.md index 044f8360..905df051 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -11,6 +11,98 @@ ReverseDiff.gradient ReverseDiff.gradient! ``` +### [Example](@id gradient-example) + +```jldoctest gradient +julia> f(a, b) = sum(a' * b + a * b'); + +julia> a, b = [1.0 2.0 3.0; 4.0 5.0 6.0; 7.0 8.0 9.0], [2.0 0.0 1.0; 1.0 3.0 0.0; 0.0 1.0 2.0]; + +julia> inputs = (a, b); +``` + +For a one-off gradient, simply call `gradient`: + +```jldoctest gradient +julia> using ReverseDiff: gradient + +julia> ∇f = gradient(f, inputs); +``` + +The in-place `gradient!` writes the gradients into pre-allocated `results`: + +```jldoctest gradient +julia> using ReverseDiff: gradient! + +julia> results = (similar(a), similar(b)); + +julia> gradient!(results, f, inputs) == ∇f +true +``` + +In addition to the gradients, `DiffResult` instances from +[DiffResults.jl](https://github.com/JuliaDiff/DiffResults.jl) also store the value +`f(a, b)`: + +```jldoctest gradient +julia> using DiffResults + +julia> all_results = map(DiffResults.GradientResult, (similar(a), similar(b))); + +julia> gradient!(all_results, f, inputs); + +julia> map(DiffResults.value, all_results) == (f(a, b), f(a, b)) +true + +julia> map(DiffResults.gradient, all_results) == ∇f +true +``` + +Every such call allocates a new [`GradientConfig`](@ref ReverseDiff.GradientConfig) and +records a new tape. For repeated calls with inputs of the same shape and element type, a +pre-allocated `GradientConfig` avoids the former: + +```jldoctest gradient +julia> using ReverseDiff: GradientConfig + +julia> cfg = GradientConfig(inputs); + +julia> gradient!(results, f, inputs, cfg) == ∇f +true +``` + +A pre-recorded [`GradientTape`](@ref ReverseDiff.GradientTape) also avoids re-recording the +tape. It is recorded once with inputs of the same shape and element type: + +```jldoctest gradient +julia> using ReverseDiff: GradientTape + +julia> f_tape = GradientTape(f, (rand(3, 3), rand(3, 3))); + +julia> gradient!(results, f_tape, inputs) == ∇f +true +``` + +!!! warning + A recorded tape, compiled or not, only replays the operations that were executed + while recording it. It is only valid for inputs that take exactly the same path + through `f`. For example, a branch that depends on the input values won't be + re-evaluated. `f` above has no such control flow, so reusing its tape is safe. See + [The `AbstractTape` API](@ref). + +[`compile`](@ref ReverseDiff.compile) turns the tape into a more optimized representation. +Compiling has an upfront cost, but executing the compiled tape is typically the fastest +method when many gradients are needed: + +```jldoctest gradient +julia> using ReverseDiff: compile + +julia> compiled_f_tape = compile(f_tape); + +julia> gradient!(results, compiled_f_tape, inputs) == ∇f +true +``` + ## Jacobians of `f(x::AbstractArray{<:Real}...)::AbstractArray{<:Real}` ```@docs @@ -18,6 +110,138 @@ ReverseDiff.jacobian ReverseDiff.jacobian! ``` +### [Example](@id jacobian-example) + +Consider an out-of-place function `f` and an in-place function `g!` that computes the same +output: + +```jldoctest jacobian +julia> f(a, b) = (a + b) * (a * b)'; + +julia> using LinearAlgebra: mul! + +julia> g!(out, a, b) = mul!(out, a + b, (a * b)'); + +julia> a, b = [1.0 2.0 3.0; 4.0 5.0 6.0; 7.0 8.0 9.0], [2.0 0.0 1.0; 1.0 3.0 0.0; 0.0 1.0 2.0]; + +julia> inputs = (a, b); + +julia> out = similar(a); + +julia> g!(out, a, b) == f(a, b) +true +``` + +For a one-off Jacobian, simply call `jacobian`: + +```jldoctest jacobian +julia> using ReverseDiff: jacobian + +julia> J = jacobian(f, inputs); + +julia> jacobian(g!, out, inputs) == J +true +``` + +The in-place `jacobian!` writes the Jacobians into pre-allocated `results`: + +```jldoctest jacobian +julia> using ReverseDiff: jacobian! + +julia> results = (similar(a, 9, 9), similar(b, 9, 9)); + +julia> jacobian!(results, f, inputs) == J +true + +julia> jacobian!(results, g!, out, inputs) == J +true +``` + +In addition to the Jacobians, `DiffResult` instances from +[DiffResults.jl](https://github.com/JuliaDiff/DiffResults.jl) also store the value +`f(a, b)`: + +```jldoctest jacobian +julia> using DiffResults + +julia> all_results = map(x -> DiffResults.JacobianResult(out, x), inputs); + +julia> jacobian!(all_results, f, inputs); + +julia> map(DiffResults.value, all_results) == (f(a, b), f(a, b)) +true + +julia> map(DiffResults.jacobian, all_results) == J +true + +julia> jacobian!(all_results, g!, out, inputs); + +julia> map(DiffResults.value, all_results) == (f(a, b), f(a, b)) +true + +julia> map(DiffResults.jacobian, all_results) == J +true +``` + +Every such call allocates a new [`JacobianConfig`](@ref ReverseDiff.JacobianConfig) and +records a new tape. For repeated calls with inputs (and, for `g!`, an output) of the same +shape and element type, a pre-allocated `JacobianConfig` avoids the former: + +```jldoctest jacobian +julia> using ReverseDiff: JacobianConfig + +julia> f_cfg = JacobianConfig(inputs); + +julia> jacobian!(results, f, inputs, f_cfg) == J +true + +julia> g!_cfg = JacobianConfig(out, inputs); + +julia> jacobian!(results, g!, out, inputs, g!_cfg) == J +true +``` + +A pre-recorded [`JacobianTape`](@ref ReverseDiff.JacobianTape) also avoids re-recording the +tape. It is recorded once with inputs (and, for `g!`, an output) of the same shape and +element type: + +```jldoctest jacobian +julia> using ReverseDiff: JacobianTape + +julia> f_tape = JacobianTape(f, (rand(3, 3), rand(3, 3))); + +julia> jacobian!(results, f_tape, inputs) == J +true + +julia> g!_tape = JacobianTape(g!, rand(3, 3), (rand(3, 3), rand(3, 3))); + +julia> jacobian!(results, g!_tape, inputs) == J +true +``` + +!!! warning + As for gradients, a recorded tape, compiled or not, is only valid for inputs that take + exactly the same path through the function as the ones it was recorded with. See + [The `AbstractTape` API](@ref). + +[`compile`](@ref ReverseDiff.compile) turns the tapes into more optimized representations. +Compiling has an upfront cost, but executing the compiled tapes is typically the fastest +method when many Jacobians are needed: + +```jldoctest jacobian +julia> using ReverseDiff: compile + +julia> compiled_f_tape = compile(f_tape); + +julia> jacobian!(results, compiled_f_tape, inputs) == J +true + +julia> compiled_g!_tape = compile(g!_tape); + +julia> jacobian!(results, compiled_g!_tape, inputs) == J +true +``` + ## Hessians of `f(x::AbstractArray{<:Real})::Real` ```@docs diff --git a/docs/src/index.md b/docs/src/index.md index 4830c653..259f970f 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -6,5 +6,5 @@ and details about the package's implementation. For a basic introduction to ReverseDiff, see [the package's README](https://github.com/JuliaDiff/ReverseDiff.jl). -For usage examples, see -[the examples directory](https://github.com/JuliaDiff/ReverseDiff.jl/tree/master/examples). +For usage examples, see the examples for [gradients](@ref gradient-example) and +[Jacobians](@ref jacobian-example). diff --git a/examples/gradient.jl b/examples/gradient.jl deleted file mode 100644 index 9d6dab53..00000000 --- a/examples/gradient.jl +++ /dev/null @@ -1,58 +0,0 @@ -using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile, DiffResults - -######### -# setup # -######### - -# some objective function to work with -f(a, b) = sum(a' * b + a * b') - -# pre-record a GradientTape for `f` using inputs of shape 100x100 with Float64 elements -const f_tape = GradientTape(f, (rand(100, 100), rand(100, 100))) - -# compile `f_tape` into a more optimized representation -const compiled_f_tape = compile(f_tape) - -# some inputs and work buffers to play around with -a, b = rand(100, 100), rand(100, 100) -inputs = (a, b) -results = (similar(a), similar(b)) -all_results = map(DiffResults.GradientResult, results) -cfg = GradientConfig(inputs) - -#################### -# taking gradients # -#################### - -# with pre-recorded/compiled tapes (generated in the setup above) # -#-----------------------------------------------------------------# - -# this should be the fastest method, and non-allocating -gradient!(results, compiled_f_tape, inputs) - -# the same as the above, but in addition to calculating the gradients, the value `f(a, b)` -# is loaded into the the provided `DiffResult` instances (see DiffResults.jl documentation). -gradient!(all_results, compiled_f_tape, inputs) - -# this should be the second fastest method, and also non-allocating -gradient!(results, f_tape, inputs) - -# you can also make your own function if you want to abstract away the tape -∇f!(results, inputs) = gradient!(results, compiled_f_tape, inputs) - -# with a pre-allocated GradientConfig # -#-------------------------------------# -# these methods are more flexible than a pre-recorded tape, but can be -# wasteful since the tape will be re-recorded for every call. - -gradient!(results, f, inputs, cfg) - -gradient(f, inputs, cfg) - -# without a pre-allocated GradientConfig # -#----------------------------------------# -# convenient, but pretty wasteful since it has to allocate the GradientConfig itself - -gradient!(results, f, inputs) - -gradient(f, inputs) diff --git a/examples/jacobian.jl b/examples/jacobian.jl deleted file mode 100644 index afe7e280..00000000 --- a/examples/jacobian.jl +++ /dev/null @@ -1,66 +0,0 @@ -using ReverseDiff: JacobianTape, JacobianConfig, jacobian, jacobian!, compile -using LinearAlgebra: mul! - -######### -# setup # -######### - -# some objective functions to work with -f(a, b) = (a + b) * (a * b)' -g!(out, a, b) = mul!(out, a + b, a * b) - -# pre-record JacobianTapes for `f` and `g` using inputs of shape 10x10 with Float64 elements -const f_tape = JacobianTape(f, (rand(10, 10), rand(10, 10))) -const g_tape = JacobianTape(g!, rand(10, 10), (rand(10, 10), rand(10, 10))) - -# compile `f_tape` and `g_tape` into more optimized representations -const compiled_f_tape = compile(f_tape) -const compiled_g_tape = compile(g_tape) - -# some inputs and work buffers to play around with -a, b = rand(10, 10), rand(10, 10) -inputs = (a, b) -output = rand(10, 10) -results = (similar(a, 100, 100), similar(b, 100, 100)) -fcfg = JacobianConfig(inputs) -gcfg = JacobianConfig(output, inputs) - -#################### -# taking Jacobians # -#################### - -# with pre-recorded/compiled tapes (generated in the setup above) # -#-----------------------------------------------------------------# - -# these should be the fastest methods, and non-allocating -jacobian!(results, compiled_f_tape, inputs) -jacobian!(results, compiled_g_tape, inputs) - -# these should be the second fastest methods, and also non-allocating -jacobian!(results, f_tape, inputs) -jacobian!(results, g_tape, inputs) - -# with a pre-allocated JacobianConfig # -#-------------------------------------# -# this is more flexible than a pre-recorded tape, but can be wasteful since the tape -# will be re-recorded for every call. - -jacobian!(results, f, inputs, fcfg) - -jacobian(f, inputs, fcfg) - -jacobian!(results, g!, output, inputs, gcfg) - -jacobian(g!, output, inputs, gcfg) - -# without a pre-allocated JacobianConfig # -#----------------------------------------# -# convenient, but pretty wasteful since it has to allocate the JacobianConfig itself - -jacobian!(results, f, inputs) - -jacobian(f, inputs) - -jacobian!(results, g!, output, inputs) - -jacobian(g!, output, inputs)