Merge pull request #116 from JuliaDiff/ox/autotangent

Automatically provide tangents

Author: oxinabox

.github/workflows/Documenter.yml

@@ -13,7 +13,7 @@ jobs:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1
         with:
-          version: '1'
+          version: '1.5'
       - uses: julia-actions/julia-docdeploy@latest
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesTestUtils"
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.6.2"
+version = "0.6.3"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/index.md

@@ -45,24 +45,20 @@ end
 
 ```
 
-The [`frule_test`](@ref)/[`rrule_test`](@ref) helper function compares the `frule`/`rrule` outputs
+The [`test_frule`](@ref)/[`test_rrule`](@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`](@ref) takes in the function `f` and tuples `(x, ẋ)` for each function argument `x`.
+[`test_frule`](@ref) takes in the function `f` and the primal input `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`.
 
 ```jldoctest ex; output = false
 using ChainRulesTestUtils
 
-x1, x2 = (3.33, -7.77)
-ẋ1, ẋ2 = (rand(), rand())
-
-frule_test(two2three, (x1, ẋ1), (x2, ẋ2))
+test_frule(two2three, 3.33, -7.77)
 # output
 Test Summary:                    | Pass  Total
 Tuple{Float64,Float64,Float64}.1 |    1      1
@@ -75,17 +71,11 @@ Test Passed
 
 ### Testing the `rrule`
 
-[`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 can be set arbitrarily.
+[`test_rrule`](@ref) takes in the function `f`, and primal inputsr `x`.
 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`. 
-```jldoctest ex; output = false
-x1, x2 = (3.33, -7.77)
-x̄1, x̄2 = (rand(), rand())
-ȳs = (rand(), rand(), rand())
-
-rrule_test(two2three, ȳs, (x1, x̄1), (x2, x̄2))
 
+```jldoctest ex; output = false
+test_rrule(two2three, 3.33, -7.77)
 # output
 Test Summary:                      |
 Don't thunk only non_zero argument | No tests
@@ -128,13 +118,30 @@ Test Summary:                    | Pass  Total
 relu at -0.5, with cotangent 1.0 |    4      4
 ```
 
+## Specifying Tangents
+[`test_frule`](@ref) and [`test_rrule`](@ref) allow you to specify the tangents used for testing.
+This is done by passing in `x ⊢ Δx`, where `x` is the primal and `Δx` is the tangent, in the place of the primal inputs.
+If this is not done the tangent will be automatically generated via [`ChainRulesTestUtils.rand_tangent`](@ref).
+A special case of this is that if you specify it as `x ⊢ nothing` then finite differencing will not be used on that input.
+Similarly, by setting the `output_tangent` keyword argument, you can specify the tangent for the primal output.
+
+This can be useful when the default provided [`ChainRulesTestUtils.rand_tangent`](@ref) doesn't produce the desired tangent for your type.
+For example the default tangent for an `Int` is `DoesNotExist()`.
+Which is correct e.g. when the `Int` represents a discrete integer like in indexing.
+But if you are testing something where the `Int` is actually a special case of a real number, then you would want to specify the tangent as a `Float64`.
+
+Care must be taken when manually specifying tangents.
+In particular, when specifying the input tangents to [`test_frule`](@ref) and the output tangent to [`test_rrule`](@ref).
+As these tangents are used to seed the derivative computation.
+Inserting inappropriate zeros can thus hide errors.
+
 ## Custom finite differencing
 
 If a package is using a custom finite differencing method of testing the `frule`s and `rrule`s, `check_equal` function provides a convenient way of comparing [various types](https://www.juliadiff.org/ChainRulesCore.jl/dev/design/many_differentials.html#Design-Notes:-The-many-to-many-relationship-between-differential-types-and-primal-types.) of differentials.
 
 It is effectively `(a, b) -> @test isapprox(a, b)`, but it preprocesses `thunk`s and `ChainRules` differential types `Zero()`, `DoesNotExist()`, and `Composite`, such that the error messages are helpful.
 
-For example, 
+For example,
 ```julia
 check_equal((@thunk 2*2.0), 4.1)
 ```
@@ -159,3 +166,7 @@ which should have passed the test.
 Modules = [ChainRulesTestUtils]
 Private = false
 ```
+
+```@docs
+ChainRulesTestUtils.rand_tangent
+```

src/ChainRulesTestUtils.jl

@@ -12,7 +12,8 @@ using Test
 const _fdm = central_fdm(5, 1; max_range=1e-2)
 
 export TestIterator
-export check_equal, test_scalar, frule_test, rrule_test, generate_well_conditioned_matrix
+export check_equal, test_scalar, test_frule, test_rrule, generate_well_conditioned_matrix
+export ⊢
 
 
 include("generate_tangent.jl")

src/deprecated.jl

@@ -27,3 +27,15 @@ end
 
 # Must be for same primal
 Base.isapprox(d_ad::Composite{P}, d_fd::Composite{Q}; kwargs...) where {P, Q} = false
+
+
+# From when primal and tangent was passed as a tuple
+@deprecate(
+    rrule_test(f, ȳ, inputs::Tuple{Any,Any}...; kwargs...),
+    test_rrule(f, ((x ⊢ dx) for (x, dx) in inputs)...; output_tangent=ȳ, kwargs...)
+)
+
+@deprecate(
+    frule_test(f, inputs::Tuple{Any,Any}...; kwargs...),
+    test_frule(f, ((x ⊢ dx) for (x, dx) in inputs)...; kwargs...)
+)

src/generate_tangent.jl

@@ -1,3 +1,44 @@
+"""
+    Auto()
+
+Use this in the place of a tangent/cotangent in [`test_frule`](@ref) or
+[`test_rrule`](@ref) to have that tangent/cotangent generated automatically based on the
+primal. Uses [`rand_tangent`](@ref)
+"""
+struct Auto end
+
+"""
+    PrimalAndTangent
+
+A struct that represents a primal value paired with its tangent or cotangent.
+For conciseness we refer to both tangent and cotangent as "tangent".
+"""
+struct PrimalAndTangent{P,D}
+    primal::P
+    tangent::D
+end
+primal(p::PrimalAndTangent) = p.primal
+tangent(p::PrimalAndTangent) = p.tangent
+
+"""
+    primal ⊢ tangent
+
+Infix shorthand method to construct a `PrimalAndTangent`.
+Enter via `\\vdash` + tab on supporting editors.
+"""
+const ⊢ = PrimalAndTangent
+
+"""
+    auto_primal_and_tangent(primal; rng=Random.GLOBAL_RNG)
+    auto_primal_and_tangent(::PrimalAndTangent; rng=Random.GLOBAL_RNG)
+
+Convience constructor for `PrimalAndTangent` where the primal is provided
+
+This function is idempotent. If you pass it a `PrimalAndTangent` it doesn't change it.
+"""
+auto_primal_and_tangent(primal; rng=Random.GLOBAL_RNG) = primal ⊢ rand_tangent(rng, primal)
+auto_primal_and_tangent(both::PrimalAndTangent; kwargs...) = both
+
 """
     rand_tangent([rng::AbstractRNG,] x)
 

src/testers.jl

@@ -26,7 +26,7 @@ function test_scalar(f, z; rtol=1e-9, atol=1e-9, fdm=_fdm, fkwargs=NamedTuple(),
     Δx = one(z)
     @testset "$f at $z, with tangent $Δx" begin
         # check ∂u_∂x and (if Ω is complex) ∂v_∂x via forward mode
-        frule_test(f, (z, Δx); rule_test_kwargs...)
+        test_frule(f, z ⊢ Δx; rule_test_kwargs...)
         if z isa Complex
             # check that same tangent is produced for tangent 1.0 and 1.0 + 0.0im
             _, real_tangent = frule((Zero(), real(Δx)), f, z; fkwargs...)
@@ -38,15 +38,15 @@ function test_scalar(f, z; rtol=1e-9, atol=1e-9, fdm=_fdm, fkwargs=NamedTuple(),
         Δy = one(z) * im
         @testset "$f at $z, with tangent $Δy" begin
             # check ∂u_∂y and (if Ω is complex) ∂v_∂y via forward mode
-            frule_test(f, (z, Δy); rule_test_kwargs...)
+            test_frule(f, z ⊢ Δy; rule_test_kwargs...)
         end
     end
 
     # test jacobian transpose using reverse mode
     Δu = one(Ω)
     @testset "$f at $z, with cotangent $Δu" begin
         # check ∂u_∂x and (if z is complex) ∂u_∂y via reverse mode
-        rrule_test(f, Δu, (z, Δx); rule_test_kwargs...)
+        test_rrule(f, z ⊢ Δx; output_tangent=Δu, rule_test_kwargs...)
         if Ω isa Complex
             # check that same cotangent is produced for cotangent 1.0 and 1.0 + 0.0im
             _, back = rrule(f, z)
@@ -59,34 +59,48 @@ function test_scalar(f, z; rtol=1e-9, atol=1e-9, fdm=_fdm, fkwargs=NamedTuple(),
         Δv = one(Ω) * im
         @testset "$f at $z, with cotangent $Δv" begin
             # check ∂v_∂x and (if z is complex) ∂v_∂y via reverse mode
-            rrule_test(f, Δv, (z, Δx); rule_test_kwargs...)
+            test_rrule(f, z ⊢ Δx; output_tangent=Δv,  rule_test_kwargs...)
         end
     end
 end
 
 
 """
-    frule_test(f, (x, ẋ)...; rtol=1e-9, atol=1e-9, fdm=central_fdm(5, 1), fkwargs=NamedTuple(), check_inferred=true, kwargs...)
+    test_frule(f, inputs...; kwargs...)
 
 # Arguments
 - `f`: Function for which the `frule` should be tested.
-- `x`: input at which to evaluate `f` (should generally be set to an arbitary point in the domain).
-- `ẋ`: differential w.r.t. `x` (should generally be set randomly).
-
-Non-differentiable arguments, such as indices, should have `ẋ` set as `nothing`.
-`fkwargs` are passed to `f` as keyword arguments.
-If `check_inferred=true`, then the inferrability of the `frule` is checked, as long as `f`
-is itself inferrable.
-All remaining keyword arguments are passed to `isapprox`.
+- `inputs` either the primal inputs `x`, or primals and their tangents: `x ⊢ ẋ`
+   - `x`: input at which to evaluate `f` (should generally be set to an arbitary point in the domain).
+   - `ẋ`: differential w.r.t. `x`, will be generated automatically if not provided
+     Non-differentiable arguments, such as indices, should have `ẋ` set as `nothing`.
+
+# Keyword Arguments
+   - `output_tangent` tangent to test accumulation of derivatives against
+     should be a differential for the output of `f`. Is set automatically if not provided.
+   - `fdm::FiniteDifferenceMethod`: the finite differencing method to use.
+   - If `check_inferred=true`, then the inferrability of the `rrule` is checked
+   - If `check_inferred=true`, then the inferrability of the `frule` is checked,
+     as long as `f` is itself inferrable.
+   - `fkwargs` are passed to `f` as keyword arguments.
+   - All remaining keyword arguments are passed to `isapprox`.
 """
-function frule_test(f, xẋs::Tuple{Any, Any}...; rtol::Real=1e-9, atol::Real=1e-9, fdm=_fdm, fkwargs::NamedTuple=NamedTuple(), check_inferred::Bool=true, kwargs...)
+function test_frule(
+    f, inputs...;
+    output_tangent=Auto(),
+    fdm=_fdm,
+    check_inferred::Bool=true,
+    fkwargs::NamedTuple=NamedTuple(),
+    rtol::Real=1e-9, atol::Real=1e-9, kwargs...
+)
     # To simplify some of the calls we make later lets group the kwargs for reuse
     isapprox_kwargs = (; rtol=rtol, atol=atol, kwargs...)
 
-    _ensure_not_running_on_functor(f, "frule_test")
+    _ensure_not_running_on_functor(f, "test_frule")
 
-    xs = first.(xẋs)
-    ẋs = last.(xẋs)
+    xẋs = auto_primal_and_tangent.(inputs)
+    xs = primal.(xẋs)
+    ẋs = tangent.(xẋs)
     if check_inferred && _is_inferrable(f, deepcopy(xs)...; deepcopy(fkwargs)...)
         _test_inferred(frule, (NO_FIELDS, deepcopy(ẋs)...), f, deepcopy(xs)...; deepcopy(fkwargs)...)
     end
@@ -102,38 +116,48 @@ function frule_test(f, xẋs::Tuple{Any, Any}...; rtol::Real=1e-9, atol::Real=1e
     dΩ_fd = _make_jvp_call(fdm, (xs...) -> f(deepcopy(xs)...; deepcopy(fkwargs)...), Ω, xs, ẋs, ẋs_is_ignored)
     check_equal(dΩ_ad, dΩ_fd; isapprox_kwargs...)
 
-    # No tangent is passed in to test accumlation, so generate one
-    # See: https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/66
-    acc = rand_tangent(Ω)
+    acc = output_tangent isa Auto ? rand_tangent(Ω) : output_tangent
     _check_add!!_behaviour(acc, dΩ_ad; rtol=rtol, atol=atol, kwargs...)
 end
 
 
+
 """
-    rrule_test(f, ȳ, (x, x̄)...; rtol=1e-9, atol=1e-9, fdm=central_fdm(5, 1), fkwargs=NamedTuple(), check_inferred=true, kwargs...)
+    test_rrule(f, inputs...; kwargs...)
 
 # Arguments
 - `f`: Function to which rule should be applied.
-- `ȳ`: adjoint w.r.t. output of `f` (should generally be set randomly).
-  Should be same structure as `f(x)` (so if multiple returns should be a tuple)
-- `x`: input at which to evaluate `f` (should generally be set to an arbitary point in the domain).
-- `x̄`: currently accumulated adjoint (should generally be set randomly).
-
-Non-differentiable arguments, such as indices, should have `x̄` set as `nothing`.
-`fkwargs` are passed to `f` as keyword arguments.
-If `check_inferred=true`, then the inferrability of the `rrule` is checked — if `f` is
-itself inferrable — along with the inferrability of the pullback it returns.
-All remaining keyword arguments are passed to `isapprox`.
+- `inputs` either the primal inputs `x`, or primals and their tangents: `x ⊢ ẋ`
+    - `x`: input at which to evaluate `f` (should generally be set to an arbitary point in the domain).
+    - `x̄`: currently accumulated cotangent, will be generated automatically if not provided
+      Non-differentiable arguments, such as indices, should have `x̄` set as `nothing`.
+
+# Keyword Arguments
+ - `output_tangent` the seed to propagate backward for testing (techncally a cotangent).
+   should be a differential for the output of `f`. Is set automatically if not provided.
+ - `fdm::FiniteDifferenceMethod`: the finite differencing method to use.
+ - If `check_inferred=true`, then the inferrability of the `rrule` is checked
+   — if `f` is itself inferrable — along with the inferrability of the pullback it returns.
+ - `fkwargs` are passed to `f` as keyword arguments.
+ - All remaining keyword arguments are passed to `isapprox`.
 """
-function rrule_test(f, ȳ, xx̄s::Tuple{Any, Any}...; rtol::Real=1e-9, atol::Real=1e-9, fdm=_fdm, check_inferred::Bool=true, fkwargs::NamedTuple=NamedTuple(), kwargs...)
+function test_rrule(
+    f, inputs...;
+    output_tangent=Auto(),
+    fdm=_fdm,
+    check_inferred::Bool=true,
+    fkwargs::NamedTuple=NamedTuple(),
+    rtol::Real=1e-9, atol::Real=1e-9, kwargs...
+)
     # To simplify some of the calls we make later lets group the kwargs for reuse
     isapprox_kwargs = (; rtol=rtol, atol=atol, kwargs...)
 
-    _ensure_not_running_on_functor(f, "rrule_test")
+    _ensure_not_running_on_functor(f, "test_rrule")
 
     # Check correctness of evaluation.
-    xs = first.(xx̄s)
-    accumulated_x̄ = last.(xx̄s)
+    xx̄s = auto_primal_and_tangent.(inputs)
+    xs = primal.(xx̄s)
+    accumulated_x̄ = tangent.(xx̄s)
     if check_inferred && _is_inferrable(f, xs...; fkwargs...)
         _test_inferred(rrule, f, xs...; fkwargs...)
     end
@@ -143,6 +167,8 @@ function rrule_test(f, ȳ, xx̄s::Tuple{Any, Any}...; rtol::Real=1e-9, atol::Re
     y = f(xs...; fkwargs...)
     check_equal(y_ad, y; isapprox_kwargs...)  # make sure primal is correct
 
+    ȳ = output_tangent isa Auto ? rand_tangent(y) : output_tangent
+
     check_inferred && _test_inferred(pullback, ȳ)
     ∂s = pullback(ȳ)
     ∂s isa Tuple || error("The pullback must return (∂self, ∂args...), not $∂s.")