Add a pedagogical example (#516)

Author: mzgubic

docs/make.jl

@@ -50,6 +50,7 @@ makedocs(;
         "Introduction" => "index.md",
         "How to use ChainRules as a rule author" => [
             "Introduction" => "rule_author/intro.md",
+            "Pedagogical example" => "rule_author/example.md",
             "Tangent types" => "rule_author/tangents.md",
             #"`frule` and `rrule`" => "rule_author/rules.md", # TODO: a complete example
             "Writing good rules" => "rule_author/writing_good_rules.md",

docs/src/index.md

@@ -10,7 +10,7 @@ They differ in the way they break down complicated functions into simple ones, b
 [ChainRules](https://github.com/JuliaDiff/ChainRules.jl) is an AD-independent set of rules, and a system for defining and testing rules.
 
 !!! note "What is a rule?"
-    A rule encodes knowledge about propagating derivatives, e.g. that the derivative (with respect to `x`) of `a*x` is `a`, and the derivative of `sin(x)` is `cos(x)`, etc.
+    A rule encodes knowledge about propagating derivatives, e.g. that the derivative with respect to `x` of `a*x` is `a`, and the derivative of `sin(x)` is `cos(x)`, etc.
 
 ## ChainRules ecosystem organisation
 

docs/src/rule_author/example.md

@@ -0,0 +1,88 @@
+# Pedagogical Example
+
+This pedagogical example will show you how to write an `rrule`.
+See [On writing good `rrule` / `frule` methods](@ref) section for more tips and gotchas.
+If you want to learn about `frule`s, you should still read and understand this example as many concepts are shared, and then look for real world `frule` examples in ChainRules.jl.
+
+## The primal
+
+We define a struct `Foo`
+```julia
+struct Foo
+    A::Matrix
+    c::Float64
+end
+```
+and a function that multiplies `Foo` with an `AbstractArray`:
+```julia
+function foo_mul(foo::Foo, b::AbstractArray)
+    return foo.A * b
+end
+```
+Note that field `c` is ignored in the calculation.
+
+## The `rrule`
+
+The `rrule` method for our primal computation should extend the `ChainRulesCore.rrule` function.
+```julia
+function ChainRulesCore.rrule(::typeof(foo_mul), foo::Foo, b::AbstractArray)
+    y = foo_mul(foo, b)
+    function foo_mul_pullback(ȳ)
+        f̄ = NoTangent()
+        f̄oo = Tangent{Foo}(; A=ȳ * b', c=ZeroTangent())
+        b̄ = @thunk(foo.A' * ȳ)
+        return f̄, f̄oo, b̄
+    end
+    return y, foo_mul_pullback
+end
+```
+Now let's examine the rule in more detail:
+```julia
+function ChainRulesCore.rrule(::typeof(foo_mul), foo::Foo, b::AbstractArray)
+    ...
+    return y, foo_mul_pullback
+end
+```
+The `rrule` dispatches on the `typeof` of the function we are writing the `rrule` for, as well as the types of its arguments.
+Read more about writing rules for constructors and callable objects [here](@ref structs).
+The `rrule` returns the primal result `y`, and the pullback function.
+It is a _very_ good idea to name your pullback function, so that they are helpful when appearing in the stacktrace.
+```julia
+    y = foo_mul(foo, b)
+```
+Computes the primal result.
+It is possible to change the primal computation so that work can be shared between the primal and the pullback.
+See e.g. [the rule for `sort`](https://github.com/JuliaDiff/ChainRules.jl/blob/a75193768775975fac5578c89d1e5f50d7f358c2/src/rulesets/Base/sort.jl#L19-L35), where the sorting is done only once.
+```julia
+    function foo_mul_pullback(ȳ)
+        ...
+        return f̄, f̄oo, b̄
+    end
+```
+The pullback function takes in the tangent of the primal output (`ȳ`) and returns the tangents of the primal inputs.
+Note that it returns a tangent for the primal function in addition to the tangents of primal arguments.
+
+Finally, computing the tangents of primal inputs:
+```julia
+        f̄ = NoTangent()
+```
+The function `foo_mul` has no fields (i.e. it is not a closure) and can not be perturbed.
+Therefore its tangent (`f̄`) is a `NoTangent`.
+```julia
+        f̄oo = Tangent{Foo}(; A=ȳ * b', c=ZeroTangent())
+```
+The struct `foo::Foo` gets a `Tangent{Foo}` structural tangent, which stores the tangents of fields of `foo`.
+
+The tangent of the field `A` is `ȳ * b'`,
+
+The tangent of the field `c` is `ZeroTangent()`, because `c` can be perturbed but has no effect on the primal output.
+```julia
+        b̄ = @thunk(foo.A' * ȳ)
+```
+The tangent of `b` is `foo.A' * ȳ`, but we have wrapped it into a `Thunk`, a tangent type that represents delayed computation.
+The idea is that in case the tangent is not used anywhere, the computation never happens.
+Use [`InplaceableThunk`](@ref) if you are interested in [accumulating gradients inplace](@ref grad_acc).
+Note that in practice one would also `@thunk` the `f̄oo.A` tangent, but it was omitted in this example for clarity.
+
+As a final note, Since `b` is an `AbstractArray`, its tangent `b̄` should be projected to the right subspace.
+See the [`ProjectTo` the primal subspace](@ref projectto) section for more information and an example that motivates the projection operation.

docs/src/rule_author/intro.md

@@ -10,6 +10,6 @@ This section also outlines some ChainRules superpowers that can be considered ad
 Most users can ignore these.
 However:
 - If you are writing rules with abstractly typed arguments, read about [`ProjectTo`](@ref projectto).
-- If you want to opt out of using the abstractly typed rule for certain argument types, read [`@opt_out`](@ref opt_out).
+- If you want to opt out of using the abstractly typed rule for certain argument types, read about [`@opt_out`](@ref opt_out).
 - If you are writing rules for higher order functions, read about [calling back into AD](@ref config).
 - If you want to accumulate gradients inplace to avoid extra allocations, read about [gradient accumulation](@ref grad_acc).

docs/src/rule_author/rules.md

@@ -71,7 +71,7 @@ Examples being:
 - There is only one derivative being returned, so from the fact that the user called
   `frule`/`rrule` they clearly will want to use that one.
 
-## Structs: constructors and functors
+## [Structs: constructors and functors](@id structs)
 
 To define an `frule` or `rrule` for a _function_ `foo` we dispatch on the type of `foo`, which is `typeof(foo)`.
 For example, the `rrule` signature would be like: