code review comments

Author: Miha Zgubic

docs/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 

docs/src/index.md

@@ -12,99 +12,126 @@ For information about ChainRules, including how to write rules, refer to the gen
 ## Canonical example
 
 Let's suppose a custom transformation has been defined
-```
-function two2three(a::Float64, b::Float64)
-    return 1.0, 2.0*a, 3.0*b
+```jldoctest ex; output = false
+function two2three(x1::Float64, x2::Float64)
+    return 1.0, 2.0*x1, 3.0*x2
 end
+
+# output
+two2three (generic function with 1 method)
 ```
 along with the `frule`
-```
-function ChainRulesCore.frule((Δf, Δa, Δb), ::typeof(two2three), a, b)
-    y = two2three(a, b)
-    ∂y = Composite{Tuple{Float64, Float64, Float64}}(Zero(), 2.0*Δa, 3.0*Δb)
+```jldoctest ex; output = false
+using ChainRulesCore
+
+function ChainRulesCore.frule((Δf, Δx1, Δx2), ::typeof(two2three), x1, x2)
+    y = two2three(x1, x2)
+    ∂y = Composite{Tuple{Float64, Float64, Float64}}(Zero(), 2.0*Δx1, 3.0*Δx2)
     return y, ∂y
 end
+# output
+
 ```
 and `rrule`
-```
-function ChainRulesCore.rrule(::typeof(two2three), a, b)
-    y = two2three(a, b)
+```jldoctest ex; output = false
+function ChainRulesCore.rrule(::typeof(two2three), x1, x2)
+    y = two2three(x1, x2)
     function two2three_pullback(Ȳ)
         return (NO_FIELDS, 2.0*Ȳ[2], 3.0*Ȳ[3])
     end
     return y, two2three_pullback
 end
+# output
+
 ```
 
-The `test_frule`/`test_rrule` helper function compares the `frule`/`rrule` outputs
+The [`frule_test`](@ref)/[`rrule_test`](@ref) helper function compares the `frule`/`rrule` outputs
 to the gradients obtained by finite differencing.
 They can be used for any type and number of inputs and outputs.
 
 ### Testing the `frule`
 
-`frule_test` takes in the function `f` and tuples `(x, ẋ)` for each function argument `x`.
+[`frule_test`](@ref) takes in the function `f` and tuples `(x, ẋ)` for each function argument `x`.
 The call will test the `frule` for function `f` at the point `x` in the domain. Keep
 this in mind when testing discontinuous rules for functions like
 [ReLU](https://en.wikipedia.org/wiki/Rectifier_(neural_networks)), which should ideally
 be tested at both `x` being above and below zero.
 Additionally, choosing `ẋ` in an unfortunate way (e.g. as zeros) could hide
 underlying problems with the defined `frule`.
 
-```
-xs = (3.33, -7.77)
-ẋs = (rand(), rand())
-frule_test(two2three, (xs[1], ẋs[1]), (xs[2], ẋs[2]))
+```jldoctest ex; output = false
+using ChainRulesTestUtils
+
+x1, x2 = (3.33, -7.77)
+ẋ1, ẋ2 = (rand(), rand())
+
+frule_test(two2three, (x1, ẋ1), (x2, ẋ2))
+# output
+Test Summary:                    | Pass  Total
+Tuple{Float64,Float64,Float64}.1 |    1      1
+Test Summary:                    | Pass  Total
+Tuple{Float64,Float64,Float64}.2 |    1      1
+Test Summary:                    | Pass  Total
+Tuple{Float64,Float64,Float64}.3 |    1      1
+Test Passed
 ```
 
 ### Testing the `rrule`
 
-`rrule_test` takes in the function `f`, sensitivities of the function outputs `ȳ`,
+[`rrule_test`](@ref) takes in the function `f`, sensitivities of the function outputs `ȳ`,
 and tuples `(x, x̄)` for each function argument `x`.
-`x̄` is the accumulated adjoint which should be set randomly.
+`x̄` is the accumulated adjoint which can be set arbitrarily.
 The call will test the `rrule` for function `f` at the point `x`, and similarly to
 `frule` some rules should be tested at multiple points in the domain.
 Choosing `ȳ` in an unfortunate way (e.g. as zeros) could hide underlying problems with
 the `rrule`. 
-```
-xs = (3.33, -7.77)
+```jldoctest ex; output = false
+x1, x2 = (3.33, -7.77)
+x̄1, x̄2 = (rand(), rand())
 ȳs = (rand(), rand(), rand())
-x̄s = (rand(), rand())
-rrule_test(two2three, ȳs, (xs[1], x̄s[1]), (xs[2], x̄s[2]))
+
+rrule_test(two2three, ȳs, (x1, x̄1), (x2, x̄2))
+
+# output
+Test Summary:                      |
+Don't thunk only non_zero argument | No tests
+Test.DefaultTestSet("Don't thunk only non_zero argument", Any[], 0, false)
 ```
 
 ## Scalar example
 
-For functions with a single argument and a single output, such as e.g. `ReLU`,
-```
+For functions with a single argument and a single output, such as e.g. ReLU,
+```jldoctest ex; output = false
 function relu(x::Real)
     return max(0, x)
 end
+
+# output
+relu (generic function with 1 method)
 ```
-with the `frule`
-```
-function ChainRulesCore.frule((Δf, Δx), ::typeof(relu), x::Real)
-    y = relu(x)
-    dydx = x <= 0 ? zero(x) : one(x)
-    return y, dydx .* Δx
-end
-```
-and `rrule` defined,
-```
-function ChainRulesCore.rrule(::typeof(relu), x::Real)
-    y = relu(x)
-    dydx = x <= 0 ? zero(x) : one(x)
-    function relu_pullback(Ȳ)
-        return (NO_FIELDS, Ȳ .* dydx)
-    end
-    return y, relu_pullback
-end
+with the `frule` and `rrule` defined with the help of `@scalar_rule` macro
+```jldoctest ex; output = false
+@scalar_rule relu(x::Real) x <= 0 ? zero(x) : one(x)
+
+# output
+
 ```
 
 `test_scalar` function is provided to test both the `frule` and the `rrule` with a single
 call.
-```
+```jldoctest ex; output = false
 test_scalar(relu, 0.5)
 test_scalar(relu, -0.5)
+
+# output
+Test Summary:                 | Pass  Total
+relu at 0.5, with tangent 1.0 |    3      3
+Test Summary:                   | Pass  Total
+relu at 0.5, with cotangent 1.0 |    4      4
+Test Summary:                  | Pass  Total
+relu at -0.5, with tangent 1.0 |    3      3
+Test Summary:                    | Pass  Total
+relu at -0.5, with cotangent 1.0 |    4      4
 ```