juliadiff / chainrulestestutils.jl Goto Github PK

Currently, if someone mistakenly uses Zero() instead of DoesNotExist() inside the rrule

function partly_dne(a, i)
    return a[i]
end

function ChainRulesCore.rrule(::typeof(partly_dne), a, i)
    y = partly_dne(a, i)
    function partly_dne_pullback(ȳ)
        grad = zeros(size(a))
        grad[i] = ȳ
        return (NO_FIELDS, grad, Zero())
    end
    return y, partly_dne_pullback
end

The error message isn't very clear

julia> rrule_test(partly_dne, rand(), (rand(4), rand(4)), (1, nothing))
Test Failed at /Users/mzgubic/Projects/ChainRules.jl/dev/ChainRulesTestUtils/src/testers.jl:295
  Expression: x̄_ad isa DoesNotExist
   Evaluated: Zero() isa DoesNotExist
ERROR: There was an error during testing

Currently, if you run rrule_test on something like the example below, FiniteDifferences makes its estimate of the cotangent of x an Adjoint, rather than a Composite. This isn't desirable behaviour in general, so we might want to consider providing an escape hatch in FiniteDifferences that lets us provide a function to get the cotangent of x from its vector-form, and utilising that escape hatch in this package to enable the user to provide the correct form of the cotangent, or to insist that FiniteDifferences uses a to_vec that is derived from an appropriate cotangent.

Example:

rrule(parent, randn(5, 3), (randn(5, 3)', Composite(parent=randn(5, 3)))

We should have reverse dependency intregration tests on thigns that have a dependency (including test time)
on ChainRulesTestUtils.

Adding more of these should not be hard because it should basically be copying what we already have.
Perhaps smartly using a matrix

Us not chatching JuliaArrays/BlockDiagonals.jl#49 until nightly ran, when we could have avoided it in #57 is not great.
(Though can't deny that Nightly build is given value there)

Clicking on the badge in the repo gives a 404

This may be the same issue as "do we want extern at all?" (JuliaDiff/ChainRulesCore.jl#56)

But I think if we have it, we may need it to be "smarter" if we're going to have isapprox(a::Differential, b) = isapprox(extern(a), b) as we currently do.

Because we will need to extern(a) (or maybe just unthunk?) the differential type before calling isapprox anyway if a it is not a scalar (e.g. the Zero Matix case in SVD tests: https://github.com/JuliaDiff/ChainRules.jl/blob/8c4985a9a941df165740badfbee83cdf93661742/test/rulesets/LinearAlgebra/factorization.jl#L20-L22)

Proposal

A very helpful function for testing adjoints is one that accepts an AbstractRNG and a primal, and returns a valid differential for the primal. I don't think we necessarily care too much about the statistics of the object produced, we just care that it's a valid differential (meaning it satisfies any implicit constraints, such as symmetry for Symmetric matrix differentials) that is sufficiently random to avoid things like e.g. only testing differentials full of 1s or something.

Implementation sketch:

differential_randn(rng::AbstractRNG, a::Real) = randn(rng)

function differential_randn(rng::AbstractRNG, x::Array)
    Δx = Array{eltype(x)}(undef, size(x))
    for n in eachindex(Δx)
        Δx[n] = differential_randn(rng, x[n])
    end
    return Δx
end

function differential_randn(rng::AbstractRNG, x::Fill)
    return (value=differential_randn(rng, x.value), axes=nothing)
end

function differential_randn(rng::AbstractRNG, x::Zeros)
    return (axes=nothing,)
end

(please mentally replace the NamedTuples with appropriate Composites.

Why is this useful

Say you're writing a lot of gradient tests because, e.g. you want to ensure that your favourite reverse-mode AD works well with a package that you're implementing involving Gaussian processes. Quite often you'll want to implement batteries of automated unit tests that run various consistency checks, gradient-related things included.

Say you're testing a collection of functions, most of which output an Array{<:Real}, but some of which output a FillArray{<:Real}, and you're either not sure before running any particular function what you're going to get, or you cba to figure it out for each test case (because life it too short).

The differentials for Array{<:Real}s and FillArray{<:Real}s are really very different, so if you don't want to have to figure out what one should look like, and you do want to test reverse-mode AD, you need something to construct an appropriate one for you. The proposed differential_randn does precisely this.

nb

I'm open to better naming suggestions, because naming things is hard and I don't think that this is a good name.

We should probably have some kinda of structural_zeros test
which for rrule checks that the structural zeros of the sensitivity are at least a superset of the input,
and for frule similar but for output

The fundamental theorem of calculus says that differential problems can also be expressed as Integral problems.
Certain functions like clamp, max, min lack a definition of derivatives at every point.
When you perform a finite difference approximation the result will certainly be in doubt around these points.
Integrals would be mostly immune against those defects.

Do you have any thoughts on idea?

We have frule_test for frules, rrule_test for rrule... but test_scalar for @scalar_rule created rules.

I think it'd be easier to remember these names if we were consistent on whether we used the suffix _test or the prefix test_.

Either (a) rename test_scalar -> scalar_test (or scalar_rule_test)
Or (b) rename frule_test -> test_frule and rrule_test -> test_rrule

Renaming is breaking, and that's annoying, but we may want to consider doing it for a v1.0 of this package?

During the fixing of JuliaDiff/ChainRules.jl#100
I came in to a problem with defining a test for ldexp which takes a float and an integer.
Wrapping it in another functor or temporary function did not work.
Neiter did the approach to verity binary/ternary functions.

right now all the tests check they are not funning on functors.

With the new API in #116
we can have f ⟂ ḟ and `f ⟂ f̄``, which would be fine API-wise.

See discussion at https://discourse.julialang.org/t/chainrulescore-rrule-for-custom-struct-does-the-pullback-need-to-support-composite-explicitly/56734

We can maybe do something with zero and similar like Zygote does for getindex
https://github.com/FluxML/Zygote.jl/blob/956cbcf3c572c0eb09c146189bb38b1b434634ff/src/lib/array.jl#L47-L49
I find thought when i was working on the getindex for ChainRules (JuliaDiff/ChainRules.jl#240) trying to work this out for general case is actually really hard, i gave up and restricteds to x::Array{<:Number}.

But maybe for rand_tangent we don't need to do so well, we just need to make something that is more likely to be what the user expects than Composite is.
It can still screw up, but as long as it screws up less, or with a clearer error.

JuliaDiff/ChainRules.jl#223 manually implements a nograd. Once someone has gotten around to sorting JuliaDiff/ChainRulesCore.jl#150 we'll be able to nograd loads of things really easily, at which point it starts to make sense to want to have a simple way to test that things have been successfully nograd-ed and regressions don't happen.

The tests I had to implement in JuliaDiff/ChainRules.jl#223 are quite verbose and boilerplate-y, so it would be good to have a one-liner that covers both frules and rrules.

Using the example from the docs:

using ChainRulesCore
using ChainRulesTestUtils

function two2three(x1::Float64, x2::Float64)
    return 1.0, 2.0*x1, 3.0*x2
end

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

test_rrule(two2three, 3.33, -7.77)

Test Summary:                      |
Don't thunk only non_zero argument | No tests
Test.DefaultTestSet("Don't thunk only non_zero argument", Any[], 0, false)

version info:

  [d360d2e6] ChainRulesCore v0.9.29
  [cdddcdb0] ChainRulesTestUtils v0.6.3
Julia Version 1.5.2
Commit 539f3ce943 (2020-09-23 23:17 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Core(TM) i7-9800X CPU @ 3.80GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-9.0.1 (ORCJIT, skylake-avx512)

If this is not implemented yet, then probably the ChainRules docs shouldn't point to this package at this stage.

We need to setup a TravisCI environment, here

The frule and rrule testing functionality is now pretty decent. However, there is not good utility functionality to assist in writing tests for how each of the various AbstractDifferentials play with each of the frules and rrules.

Ideally frule_test and rrule_test will be set up such that each possible combination of arguments to each pushforward and pullback are tested separately, and for the most part automatically.

edit: on second thoughts, we should consider whether or not we actually want to test all of the possible combinations, or just make sure that each argument sees each AbstractDifferential, since the pushforwards and pullbacks are linear in their arguments.

Right now we hit the fallback definition of rrule/frulewhich is to just returnnothing`.
which then gives an error complaining about iteration failing.

We should instead detect that an error ourselves with somehting more clear.
Possibly throwing a MethodError which would show what is closest.
(though it would also show the Any method that does match but maybe that is ok)

Since we can mostly generate random tangents correctly based on the primals.
we should let the using opt to not provide them.

I am not sure the API we want for it.

Option 1

Maybe we keep the same API as is current, but let the user replace the tangent with, Auto().
Then we go and generate tangents automatically for anything marked Auto().
(or :auto or 🎩 or something.)

Option 2

Or maybe we just let them skip it, and instead do something else:

Option 2A

via passing in a 1-tuple instead of a 2-tuple for xẋs,
and by skipping it for output.

Optiom 2B

Or maybe we let them skip it by just passing the primal.
Then it it is a 2-Tuple it needs to be primal, and differential.
But then it is ambigious if the primal itself is a Tuple.

Option 2C

We could change tuple to be a Pair insead, so one would write x => ẋ
instead of (x, ẋ).
Then Pair is ambigious, but that is less of a problem.

Option 2D

We could introduce out own type, and overload some other nice operator to construct it.
One option if <| which is ANSI and in 1.0 and kind of looks like a tangent.
Julia 1.5 adds --> and <--.
Or we could go some unicode character.

The default tolerance of atol=1e-9, rtol=1e-9 for the isapprox tests seem like they might be too stringent for 32-bit checks. Perhaps we should have a utility method for settings these tolerances based on the float eltype of the inputs?

rrule_test takes in sample sensitivities (tangents) for all inputs,even though rrule. doesn't use them (since it is reverse mode) which gives us something to test accumulation against.

ffule_test should do the same.

needful to be more consistent in how we handle testign that: see #59

Sorry if this is a duplicate of an existing issue, but I'd like to recommend 2 improvements to rrule_test.

1. relax the equality check of the outputs to use isapprox. Sometimes a slightly different algorithm is used to compute the output when doing AD and will cause a strict equality test to fail. e.g. the reverse-mode derivatives of eigvals(::Symmetric) uses eigen(::Symmetric), which produces approximately the same eigvals but not identical: https://github.com/FluxML/Zygote.jl/blob/0b3e32d5c0f8e5ef3b061a6f84ed5505f5a202a0/src/lib/array.jl#L590-L593
2. If the function returns a tuple of outputs or any of the arguments are tuples, iterate over them.

I had adapted rrule_test for some Zygote functions and implemented these changes: https://gist.github.com/sethaxen/fa67e541c4a2a5e773b475349ed87fb9/677f58af483e270fc5bdc1a318825e9fef1b0256#file-zygote_power_series_tests-jl-L4-L48

We should move to FiniteDifferences v0.11

I know mutating rules aren't fully supported yet, but I've started implementing some, following the guidelines in this tweet. To support testing mutating functions, I think only a few changes are needed:

1. Because the mutating function might use the values of an input before overwriting it, it is only safe to call the function multiple times as the testers do on copies of the inputs. Specifically deep copies. So the idea is to only pass deep copies of inputs to the functions in the testers.
2. Because the pullback must undo the mutation, after calling the pullback, check that the copies of the inputs that were passed to the rrule are approximately equal to the original inputs.

For an example of how this might work, see this working implementation.

Am I missing anything, or are there any objections to me preemptively making these changes?

A follow-up to #60
we could make a new package in JuliaTesting called MetaTesting.jl to hold things like the function that captures all failures.

See JuliaDiff/ChainRules.jl#233 and reference therein for possible issues when unexpected casts happen

When using test_scalar to test functions with multiple outputs the error messages are uninformative:

julia> simo(x) = (x, 2x)

julia> @scalar_rule simo(x) 1.0 2.0

julia> test_scalar(simo, 3.0)
Test Summary:                 | Pass  Total
simo at 3.0, with tangent 1.0 |    3      3
ERROR: MethodError: no method matching one(::Tuple{Float64,Float64})
Closest candidates are:
  one(::Type{Missing}) at missing.jl:103
  one(::BitArray{2}) at bitarray.jl:400
  one(::Missing) at missing.jl:100
  ...
Stacktrace:
 [1] #test_scalar#26(::Float64, ::Float64, ::FiniteDifferences.FiniteDifferenceMethod{Array{Int64,1},Array{Float64,1},FiniteDifferences.var"#27#29"{FiniteDifferences.FiniteDifferenceMethod{Array{Int64,1},Array{Float64,1},FiniteDifferences.var"#default_bound_estimator#19"{Int64}}}}, ::NamedTuple{(),Tuple{}}, ::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, ::typeof(test_scalar), ::typeof(simo), ::Float64) at /Users/mzgubic/Projects/ChainRules.jl/dev/ChainRulesTestUtils/src/testers.jl:154
 [2] test_scalar(::Function, ::Float64) at /Users/mzgubic/Projects/ChainRules.jl/dev/ChainRulesTestUtils/src/testers.jl:125
 [3] top-level scope at REPL[23]:1

We could also support multiple outputs though I don't think it is a priority given rrule_test and frule_test are available.

In order to make frule_test and rrule_test more generally useful for testing AD systems,
we should make them take a keyword argument that specifies what function that what to use.
This function would be used in the place of the default frule/rrule.

It would need to match the API of frule and rrule.
Just like the function that needs to be provided for JuliaDiff/ChainRulesCore.jl#68

Then we would be able to use these tools to test AD systems

Consider this (From #36)

simo(x) = (x, 2x)
function ChainRulesCore.frule((_, ẋ), simo, x)
    y = simo(x)
    return y, Composite{typeof(y)}(ẋ, 2ẋ)
end

I believe that that the following should work:

frule_test(simo, (randn(), randn(4)))  # chunked mode, scalar/vector
frule_test(simo, (randn(3), randn(2, 3)))  # chunked mode, vector/matrix

as I believe the following is the correct chunked mode behavour.
@YingboMa am i right?

at least for the scalar primal and vector differential

julia> frule((Zero(), [1, 2, 3]), simo, π)
((π, 6.283185307179586), ([1, 2, 3], [2, 4, 6]))

julia> frule((Zero(), [1 1; 0 1; 1 0]), simo, [1, 1, 1])
(([1, 1, 1], [2, 2, 2]), Composite{Tuple{Array{Int64,1},Array{Int64,1}}}([1 1; 0 1; 1 0], [2 2; 0 2; 2 0]))

If I am getting the rrule for something that takes e.g. a Diagonal matrix as input,
the senstivity should also be Diagonal etc.

There may be some special cases where it might need to be even more sparse, e.g one hot,
but that is a special case.

Likely we need to disable this, so it should be made configuable.

This functionality now lives in FiniteDifferences.jl, so should be removed from here.

We have a fairly solid idea of what it means to be a valid differnetials type.
Needs to support multipleation with a scalar
needs to support addition to self and to primal

It is rare to need to implement new once, but not inconceave-able.

Using ChainRulesTestUtils as a test time dependency of ChainRulesCore is not technically a circular dependency
though it is close, i think it is fine?

We have check_equal for Zero() but not for DoesNotExist(), which currently breaks

julia> check_equal(2.0, DoesNotExist())
ERROR: MethodError: no method matching length(::DoesNotExist)
...

julia> check_equal(Zero(), DoesNotExist())
ERROR: StackOverflowError:
...

We should probably have DoesNotExist() be equal to itself, and nothing else

check_equal(::DoesNotExist, x; kwargs...) = @test false
check_equal(x, ::DoesNotExist; kwargs...) = @test false
check_equal(x::DoesNotExist, y::DoesNotExist; kwargs...) = @test true

We've currently got this package and, by extension, ChainRules.jl fixed to FiniteDifferences v0.7. This is bad because there are newer and better versions available. There was a breaking API change between 0.7 and 0.9, so would be good to upgrade and change bounds.

Related to JuliaDiff/FiniteDifferences.jl#52, it seems there is currently no way to specify constraints on function arguments to FiniteDifferences. For instance, when testing the ChainRule adjoint of (an extension of) binomlogpdf(n, p, x) input p has to be constrained to the interval [0, 1] since otherwise the forward pass errors.

This issue is used to trigger TagBot; feel free to unsubscribe.

If you haven't already, you should update your TagBot.yml to include issue comment triggers.
Please see this post on Discourse for instructions and more details.

If you'd like for me to do this for you, comment TagBot fix on this issue.
I'll open a PR within a few hours, please be patient!

Presently testing a function

function relu(x::Real)
    return max(0, x)
end

gives

julia> frule_test(relu, (1.2, .23))
Test Passed

julia> rrule_test(relu, 2.1, (-1.2, .23))
Test Summary:                      |
Don't thunk only non_zero argument | No tests
Test.DefaultTestSet("Don't thunk only non_zero argument", Any[], 0, false)

julia> test_scalar(relu, 0.5)
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

Should we make the outputs all consistent?

Creating an in-depth README explaining what this package is, how to use it, etc.

This came up in JuliaDiff/ChainRules.jl#182, where rrule_test fails because one of the args equals :U. _make_fdm_call seems to evaluate that argument as a variable U, which does not exist.

We used to have something like this, but when we removed accumulate etc we removed it.

To do this well it would be good to bring back accumulate!
See JuliaDiff/ChainRulesCore.jl#113

Would be good to update the version bound for FiniteDifferences to v0.10.0, which fixed some issues with handling complex numbers: JuliaDiff/FiniteDifferences.jl#76.

To really test rrules with complex numbers with FD, we'll need to make sure the inputs/output/sensitivity types are correct. I think the rough logic we discussed is:

If the output is complex, then the "sensitivity" of the output must be complex, and the sensitivities of all inputs will be complex, so all inputs must be complex

This is just for FD; the rrule should be tested with the user-provided inputs.
I used this heuristic in this gist testing Zygote. However, cases like these when testing Zygote rely on other rrules, so this won't work with ChainRules. The only way I can think to handle this is for the rrule_test or _make_fdm_call to pass a modified version of the function to FiniteDifferences with modified inputs.

I haven't thought at all about what to do when the output is something besides a number or array. I assume these aren't supported by rrule_test. This proposal also doesn't handle cases where one wants to test rrules whose inputs are e.g. structured arrays, where FD will force the pulled back sensitivity to have the same structure as the inputs. See FluxML/Zygote.jl#608 (comment) for an example.

Would be good also to test holomorphicness when complex numbers are involved, but I don't know much about this.

Reusing similar example as in #112:

function partly_dne(a, i)
    return a[i]
end

function ChainRulesCore.rrule(::typeof(partly_dne), a, i)
    y = partly_dne(a, i)
    function partly_dne_pullback(ȳ)
        grad = zeros(size(a))
        grad[i] = ȳ
        return (NO_FIELDS, grad, DoesNotExist())
    end
    return y, partly_dne_pullback
end

It seems un ChainRules like that nothing should be used instead of DoesNotExist(). Should we change this for v1?

julia> rrule_test(partly_dne, rand(), (rand(4), rand(4)), (1, nothing))

We will need to write more unit tests for this package.

From

https://github.com/JuliaDiff/ChainRules.jl/blob/24318b0321ccd48f16cbbd59dba6ae8bb9e90860/test/rulesets/LinearAlgebra/structured.jl#L120

using ChainRules
using ChainRulesCore
using ChainRulesTestUtils
using LinearAlgebra

f = adjoint
T = Float64
n = 5
m = 3

A = randn(T, n, m)
Y = f(A)
Ȳ_mat = randn(T, m, n)
Ȳ_composite = Composite{typeof(Y)}(parent=collect(f(Ȳ_mat)))

test_rrule(f, A; output_tangent=Ȳ_mat) # works
test_rrule(f, A; output_tangent=Ȳ_composite) # breaks
test_rrule(f, A) # breaks

Pushforwards and pullbacks are requred to be linear operators.
Which formally is defined as φ(𝛼x) = 𝛼φ(x) and φ(a+b) = φ(a) + φ(b) for a,b,x differentials/vector space elements, and 𝛼 being a scalar.

A important collory of that is φ(0) = 0 where two 0s are the additive identity for the relivent vector spaces.
Which we can test with iszero(φ(zero(d)), since we always have access to a differential d, and all differentials must define zero and iszero.

Similarly we can test φ(𝛼x) = 𝛼φ(x) via φ(2d) == 2φ(d) since all differnetial define multiplication with a scalar. and we have one provided d.

We can't easily test φ(a+b) = φ(a) + φ(b) though, as we only have one differential.
We can test φ(d+d) = φ(d) + φ(d) but that is less interesting.
If we are generating them using rand_tangent we can easily generate another.
Though perhaps we should be insisting that things define the addition of there natural differentials with the structural differentials that rand_tangent will most often provide (though sometimes rand_tangent's structural differential will have extra nondifferentiable fields so we can't trust it, unless we insist the user does actually overload it).
Or we could avoid direct additon and do (primal + rand_tangent(d)) + d) though this is getting off-topic..
Possibly just skip that in initial implementation of this

The linear pullback can be extracted from the rrule via (f,primal...) -> (ds...) -> last(rrule(f, primal...))(ds...) (or equiv: (f,primal) -> last(rrule(f, primal))) The linear pushforward can be extracted from frulevia(f,primal) -> (ds...) -> frule(ds..., f, primal...)`

We need to release an initial commit with a Project.toml, etc. and then we can register the package with the Registrator, here

Port over test_utils from ChainRules into this package.

needed for JuliaDiff/ChainRules.jl#175

This PR is the latest to propose a rule for AbstractArrays that, in my opinion, is too broadly-typed. While I obviously sympathise with rule authors wanting to do this (it feels like the natural thing to do), we know that it causes real problems.

I wonder whether we should implement a test util comprising rrule_tests and frule_tests on a suite of different types of arrays that we insist any new rule that purports to support AbstractArrays, or some subtype thereof, must pass. The idea being would would include

1. examples of every concrete type of array that we can find in Base / the standard libraries
2. examples from popular external packages, such as StaticArrays, FillArrays, and BlockArrays.

Crucially, I propose that the tests should pass if, for all arrays in the suite, either

1. the standard rule tests pass, or
2. no rule is found for that array.

This means that the rule author can pick and choose the kinds of arrays that they support, and obvious "gaming" of the tests would be addressed at review time (i.e. the idea is not to implement a rule of AbstractArrays and then implement rules that return nothing for each non-Array array in the test suite).

This approach is obviously not a panacea since you clearly can't test against all concrete subtypes of AbstractArray that will ever be written. Rather, the goal is to have enough coverage of types of arrays that break assumptions that hold for Array-like AbstractArrays to ensure that rule authors and reviewers don't accidentally implement something too loosely-typed. My hope is that this will encourage rule authors (and reviewers) to be cognisant of the issues surrounding implementing rules for abstract types.

While Numbers have generally been less troublesome in practice, they obviously suffer from the same issues and have caused problems from time to time. It might make sense to consider a similar thing for Numbers.

Maybe there are some other common types where this stuff goes wrong on a regular basis?

I'm not entirely sure how this should look in practice (do we want a formal testing function that actually runs a load of tests, or do we just want to provide a global const Vector of subtypes of AbstractArray that we insist gets used when writing tests?), but if people like the general idea then we can think a little more carefully about how to make it work.

rrule_test use collect to get the field of composite objects, but it is not always possible, as for e.g. qr

F = qr(randn(5, 5))
Q, R = F  # Valid since F is iterable
collect(F)  # Errors since length is not defined on LinearAlgebra.QRCompactWY

Using [y...] instead of collect(y) should do the trick.