Fix doctests; Organise and improve some docs (#131)

* Remove docs dependency on ChainRules

* Organise API docs page

* Fix doctests

* Improve docs on unthunking

* Build docs on latest stable Julia

* Add comment re rules in doctest setup

Co-Authored-By: Lyndon White <oxinabox@ucc.asn.au>

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

Author: nickrobinson251

.travis.yml

@@ -15,7 +15,7 @@ jobs:
     - julia: 1.4
   include:
     - stage: "Documentation"
-      julia: 1.0
+      julia: 1
       os: linux
       script:
         - julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'

docs/Manifest.toml

@@ -1,24 +1,20 @@
+# This file is machine-generated - editing it directly is not advised
+
 [[Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
-[[ChainRules]]
-deps = ["ChainRulesCore", "LinearAlgebra", "Reexport", "Requires", "Statistics"]
-git-tree-sha1 = "906cb2ae273ddbc559490117faa7abd36c98f51a"
-uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.3.2"
-
 [[ChainRulesCore]]
 deps = ["MuladdMacro"]
 path = ".."
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.6.1"
+version = "0.7.0"
 
 [[Dates]]
 deps = ["Printf"]
 uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
 
 [[Distributed]]
-deps = ["LinearAlgebra", "Random", "Serialization", "Sockets"]
+deps = ["Random", "Serialization", "Sockets"]
 uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 
 [[DocStringExtensions]]
@@ -34,7 +30,7 @@ uuid = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 version = "0.23.4"
 
 [[InteractiveUtils]]
-deps = ["LinearAlgebra", "Markdown"]
+deps = ["Markdown"]
 uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 
 [[JSON]]
@@ -49,10 +45,6 @@ uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
 [[Libdl]]
 uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 
-[[LinearAlgebra]]
-deps = ["Libdl"]
-uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-
 [[Logging]]
 uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
 
@@ -70,12 +62,12 @@ version = "0.2.2"
 
 [[Parsers]]
 deps = ["Dates", "Test"]
-git-tree-sha1 = "0139ba59ce9bc680e2925aec5b7db79065d60556"
+git-tree-sha1 = "d112c19ccca00924d5d3a38b11ae2b4b268dda39"
 uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
-version = "0.3.10"
+version = "0.3.11"
 
 [[Pkg]]
-deps = ["Dates", "LibGit2", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
+deps = ["Dates", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Test", "UUIDs"]
 uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 
 [[Printf]]
@@ -90,18 +82,6 @@ uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 deps = ["Serialization"]
 uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
-[[Reexport]]
-deps = ["Pkg"]
-git-tree-sha1 = "7b1d07f411bc8ddb7977ec7f377b97b158514fe0"
-uuid = "189a3867-3050-52da-a836-e630ba90ab69"
-version = "0.2.0"
-
-[[Requires]]
-deps = ["UUIDs"]
-git-tree-sha1 = "999513b7dea8ac17359ed50ae8ea089e4464e35e"
-uuid = "ae029012-a4dd-5104-9daa-d747884805df"
-version = "1.0.0"
-
 [[SHA]]
 uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 
@@ -111,20 +91,12 @@ uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 [[Sockets]]
 uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
 
-[[SparseArrays]]
-deps = ["LinearAlgebra", "Random"]
-uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-
-[[Statistics]]
-deps = ["LinearAlgebra", "SparseArrays"]
-uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-
 [[Test]]
 deps = ["Distributed", "InteractiveUtils", "Logging", "Random"]
 uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [[UUIDs]]
-deps = ["Random"]
+deps = ["Random", "SHA"]
 uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [[Unicode]]

docs/Project.toml

@@ -1,5 +1,4 @@
 [deps]
-ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 

docs/make.jl

@@ -1,11 +1,26 @@
-using ChainRules
 using ChainRulesCore
 using Documenter
 
 @show ENV
 
+DocMeta.setdocmeta!(
+    ChainRulesCore,
+    :DocTestSetup,
+    quote
+        using Random
+        Random.seed!(0)  # frule doctest shows output
+
+        using ChainRulesCore
+        # These rules are all actually defined in ChainRules.jl, but we redefine them here to 
+        # avoid the dependency.
+        @scalar_rule(sin(x), cos(x))  # frule and rrule doctest
+        @scalar_rule(sincos(x), @setup((sinx, cosx) = Ω), cosx, -sinx)  # frule doctest
+        @scalar_rule(hypot(x::Real, y::Real), (x / Ω, y / Ω))  # rrule doctest
+    end
+)
+
 makedocs(
-    modules=[ChainRules, ChainRulesCore],
+    modules=[ChainRulesCore],
     format=Documenter.HTML(prettyurls=false, assets=["assets/chainrules.css"]),
     sitename="ChainRules",
     authors="Jarrett Revels and other contributors",

docs/src/api.md

@@ -1,5 +1,33 @@
 # API Documentation
 
+## Rules
 ```@autodocs
 Modules = [ChainRulesCore]
+Pages = ["rules.jl"]
+Private = false
+```
+
+## Rule Definition Tools
+```@autodocs
+Modules = [ChainRulesCore]
+Pages = ["rule_definition_tools.jl"]
+Private = false
+```
+
+## Differentials
+```@autodocs
+Modules = [ChainRulesCore]
+Pages = [
+    "differentials/abstract_zero.jl",
+    "differentials/one.jl",
+    "differentials/composite.jl",
+    "differentials/thunks.jl",
+    "differentials/abstract_differential.jl",
+]
+Private = false
+```
+
+## Internal
+```@docs
+ChainRulesCore.AbstractDifferential
 ```

docs/src/index.md

@@ -1,7 +1,3 @@
-```@meta
-DocTestSetup = :(using ChainRulesCore, ChainRules)
-```
-
 # ChainRules
 
 [ChainRules](https://github.com/JuliaDiff/ChainRules.jl) provides a variety of common utilities that can be used by downstream [automatic differentiation (AD)](https://en.wikipedia.org/wiki/Automatic_differentiation) tools to define and execute forward-, reverse-, and mixed-mode primitives.
@@ -288,62 +284,70 @@ This was once how all neural network code worked.
 
 Using ChainRules directly also helps get a feel for it.
 
-```jldoctest
-using ChainRules
+```jldoctest index; output=false
+using ChainRulesCore
 
 function foo(x)
     a = sin(x)
-    b = 2a
+    b = 0.2 + a
     c = asin(b)
     return c
 end
 
+# Define rules (alternatively get them for free via `using ChainRules`)
+@scalar_rule(sin(x), cos(x))
+@scalar_rule(+(x, y), (One(), One()))
+@scalar_rule(asin(x), inv(sqrt(1 - x^2)))
+# output
+
+```
+```jldoctest index
 #### Find dfoo/dx via rrules
 #### First the forward pass, accumulating rules
 x = 3;
 a, a_pullback = rrule(sin, x);
-b, b_pullback = rrule(*, 2, a);
+b, b_pullback = rrule(+, 0.2, a);
 c, c_pullback = rrule(asin, b)
 
 #### Then the backward pass calculating gradients
 c̄ = 1;  # ∂c/∂c
-_, b̄ = c_pullback(extern(c̄));     # ∂c/∂b
-_, _, ā = b_pullback(extern(b̄));  # ∂c/∂a
-_, x̄ = a_pullback(extern(ā));     # ∂c/∂x = ∂f/∂x
-extern(x̄)
+_, b̄ = c_pullback(unthunk(c̄));     # ∂c/∂b
+_, _, ā = b_pullback(unthunk(b̄));  # ∂c/∂a
+_, x̄ = a_pullback(unthunk(ā));     # ∂c/∂x = ∂f/∂x
+unthunk(x̄)
 # output
--2.0638950738662625
+-1.0531613736418153
 ```
-```jldoctest
+```jldoctest index
 #### Find dfoo/dx via frules
 x = 3;
 ẋ = 1;  # ∂x/∂x
 nofields = Zero();  # ∂self/∂self
 
 a, ȧ = frule((nofields, ẋ), sin, x); # ∂a/∂x
-b, ḃ = frule((nofields, Zero(), unthunk(ȧ)), *, 2, a); # ∂b/∂x = ∂b/∂a⋅∂a/∂x
+b, ḃ = frule((nofields, Zero(), unthunk(ȧ)), +, 0.2, a); # ∂b/∂x = ∂b/∂a⋅∂a/∂x
 
 c, ċ = frule((nofields, unthunk(ḃ)), asin, b); # ∂c/∂x = ∂c/∂b⋅∂b/∂x = ∂f/∂x
 unthunk(ċ)
 # output
--2.0638950738662625
+-1.0531613736418153
 ```
 ```julia
 #### Find dfoo/dx via FiniteDifferences.jl
 using FiniteDifferences
 central_fdm(5, 1)(foo, x)
 # output
--2.0638950738670734
+-1.0531613736418257
 
 #### Find dfoo/dx via ForwardDiff.jl
 using ForwardDiff
 ForwardDiff.derivative(foo, x)
 # output
--2.0638950738662625
+-1.0531613736418153
 
 #### Find dfoo/dx via Zygote.jl
 using Zygote
 Zygote.gradient(foo, x)
 # output
-(-2.0638950738662625,)
+(-1.0531613736418153,)
 ```

src/differentials/abstract_differential.jl

@@ -50,13 +50,17 @@ For two reasons:
 Where it is defined the operation of `extern` for a primal type `P` should be
 `extern(x) = zero(P) + x`.
 
-Because of its limitations, `extern` should only really be used for testing.
-It can be useful, if you know what you are getting out, as it recursively removes thunks,
-and otherwise makes outputs more consistent with finite differencing.
-The more useful action in general is to call `+`, or in the case of thunks: [`unthunk`](@ref).
+!!! note
+    Because of its limitations, `extern` should only really be used for testing.
+    It can be useful, if you know what you are getting out, as it recursively removes
+    thunks, and otherwise makes outputs more consistent with finite differencing.
 
-Note that `extern` may return an alias (not necessarily a copy) to data
-wrapped by `x`, such that mutating `extern(x)` might mutate `x` itself.
+    The more useful action in general is to call `+`, or in the case of a [`Thunk`](@ref)
+    to call [`unthunk`](@ref).
+
+!!! warning
+    `extern` may return an alias (not necessarily a copy) to data
+    wrapped by `x`, such that mutating `extern(x)` might mutate `x` itself.
 """
 @inline extern(x) = x
 

src/differentials/thunks.jl

@@ -24,18 +24,20 @@ It wraps a zero argument closure that when invoked returns a differential.
 `@thunk(v)` is a macro that expands into `Thunk(()->v)`.
 
 Calling a thunk, calls the wrapped closure.
-`extern`ing thunks applies recursively, it also externs the differial that the closure returns.
-If you do not want that, then simply call the thunk
+If you are unsure if you have a `Thunk`, call [`unthunk`](@ref) which is a no-op when the
+argument is not a `Thunk`.
+If you need to unthunk recursively, call [`extern`](@ref), which also externs the differial
+that the closure returns.
 
-```
+```jldoctest
 julia> t = @thunk(@thunk(3))
-Thunk(var"##7#9"())
+Thunk(var"#4#6"())
 
 julia> extern(t)
 3
 
 julia> t()
-Thunk(var"##8#10"())
+Thunk(var"#5#7"())
 
 julia> t()()
 3
@@ -83,7 +85,7 @@ end
 On `AbstractThunk`s this removes 1 layer of thunking.
 On any other type, it is the identity operation.
 
-In contrast to `extern` this is nonrecursive.
+In contrast to [`extern`](@ref) this is nonrecursive.
 """
 @inline unthunk(x) = x
 

src/rules.jl

@@ -17,14 +17,14 @@ Examples:
 
 unary input, unary output scalar function:
 
-```jldoctest
-julia> dself = Zero()
-Zero()
+```jldoctest frule
+julia> dself = NO_FIELDS;
 
-julia> x = rand();
+julia> x = rand()
+0.8236475079774124
 
-julia> sinx, Δsinx = frule(sin, x, dself, 1)
-(0.35696518021277485, 0.9341176907197836)
+julia> sinx, Δsinx = frule((dself, 1), sin, x)
+(0.7336293678134624, 0.6795498147167869)
 
 julia> sinx == sin(x)
 true
@@ -35,10 +35,8 @@ true
 
 unary input, binary output scalar function:
 
-```jldoctest
-julia> x = rand();
-
-julia> sincosx, Δsincosx = frule(sincos, x, dself, 1);
+```jldoctest frule
+julia> sincosx, Δsincosx = frule((dself, 1), sincos, x);
 
 julia> sincosx == sincos(x)
 true
@@ -69,7 +67,7 @@ Examples:
 
 unary input, unary output scalar function:
 
-```
+```jldoctest
 julia> x = rand();
 
 julia> sinx, sin_pullback = rrule(sin, x);
@@ -83,7 +81,7 @@ true
 
 binary input, unary output scalar function:
 
-```
+```jldoctest
 julia> x, y = rand(2);
 
 julia> hypotxy, hypot_pullback = rrule(hypot, x, y);