juliadiff / chainrulestestutils.jl Goto Github PK
View Code? Open in Web Editor NEWUtilities for testing custom AD primitives.
License: MIT License
Utilities for testing custom AD primitives.
License: MIT License
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)
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 1
s 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 NamedTuple
s with appropriate Composite
s.
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.
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 frule
s, 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 frule
s and rrule
s.
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 pushforward
s and pullback
s are linear in their arguments.
Right now we hit the fallback definition of rrule
/frulewhich is to just return
nothing`.
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.
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.)
Or maybe we just let them skip it, and instead do something else:
via passing in a 1-tuple instead of a 2-tuple for xẋs
,
and by skipping it for output.
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
.
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.
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
.
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-L593I 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:
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
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 0
s 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 AbstractArray
s 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_test
s and frule_test
s on a suite of different types of arrays that we insist any new rule that purports to support AbstractArray
s, or some subtype thereof, must pass. The idea being would would include
Base
/ the standard librariesStaticArrays
, FillArrays
, and BlockArrays
.Crucially, I propose that the tests should pass if, for all arrays in the suite, either
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 AbstractArray
s 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 AbstractArray
s 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 Number
s 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 Number
s.
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.
A declarative, efficient, and flexible JavaScript library for building user interfaces.
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
An Open Source Machine Learning Framework for Everyone
The Web framework for perfectionists with deadlines.
A PHP framework for web artisans
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
Some thing interesting about web. New door for the world.
A server is a program made to process requests and deliver data to clients.
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
Some thing interesting about visualization, use data art
Some thing interesting about game, make everyone happy.
We are working to build community through open source technology. NB: members must have two-factor auth.
Open source projects and samples from Microsoft.
Google ❤️ Open Source for everyone.
Alibaba Open Source for everyone
Data-Driven Documents codes.
China tencent open source team.