Skip to content

Commit 43d3961

Browse files
oxinaboxoxinabox
committed
reorganize tests
1 parent 022865e commit 43d3961

File tree

6 files changed

+390
-332
lines changed

6 files changed

+390
-332
lines changed

‎test/forward.jl

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
module forward
2+
using Diffractor
3+
using Diffractor: var"'", ∂⃖, DiffractorRuleConfig
4+
using ChainRules
5+
using ChainRulesCore
6+
using ChainRulesCore: ZeroTangent, NoTangent, frule_via_ad, rrule_via_ad
7+
using LinearAlgebra
8+
9+
using Test
10+
11+
const fwd = Diffractor.PrimeDerivativeFwd
12+
const bwd = Diffractor.PrimeDerivativeBack
13+
14+
15+
16+
# Minimal 2-nd order forward smoke test
17+
@test Diffractor.∂☆{2}()(Diffractor.ZeroBundle{2}(sin),
18+
Diffractor.ExplicitTangentBundle{2}(1.0, (1.0, 1.0, 0.0)))[Diffractor.CanonicalTangentIndex(1)] == sin'(1.0)
19+
20+
# Simple Forward Mode tests
21+
let var"'" = Diffractor.PrimeDerivativeFwd
22+
recursive_sin(x) = sin(x)
23+
ChainRulesCore.frule(∂, ::typeof(recursive_sin), x) = frule(∂, sin, x)
24+
25+
# Integration tests
26+
@test recursive_sin'(1.0) == cos(1.0)
27+
@test recursive_sin''(1.0) == -sin(1.0)
28+
# Error: ArgumentError: Tangent for the primal Tangent{Tuple{Float64, Float64}, Tuple{Float64, Float64}}
29+
# should be backed by a NamedTuple type, not by Tuple{Tangent{Tuple{Float64, Float64}, Tuple{Float64, Float64}}}.
30+
@test_broken recursive_sin'''(1.0) == -cos(1.0)
31+
@test_broken recursive_sin''''(1.0) == sin(1.0)
32+
@test_broken recursive_sin'''''(1.0) == cos(1.0)
33+
@test_broken recursive_sin''''''(1.0) == -sin(1.0)
34+
35+
# Test the special rules for sin/cos/exp
36+
@test sin''''''(1.0) == -sin(1.0)
37+
@test cos''''''(1.0) == -cos(1.0)
38+
@test exp''''''(1.0) == exp(1.0)
39+
@test (x->prod([x, 4]))'(3) == 4
40+
end
41+
42+
# Some Basic Mixed Mode tests
43+
function sin_twice_fwd(x)
44+
let var"'" = Diffractor.PrimeDerivativeFwd
45+
sin''(x)
46+
end
47+
end
48+
let var"'" = Diffractor.PrimeDerivativeFwd
49+
@test sin_twice_fwd'(1.0) == sin'''(1.0)
50+
end
51+
52+
53+
54+
end

‎test/forward_diff_no_inf.jl

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
2+
module forward_diff_no_inf
3+
using Diffractor, Test
4+
# this is needed as transform! is *always* called on Arguments regardless of what visit_custom says
5+
identity_transform!(ir, ssa::Core.SSAValue, order) = ir[ssa]
6+
function identity_transform!(ir, arg::Core.Argument, order)
7+
return Core.Compiler.insert_node!(ir, Core.SSAValue(1), Core.Compiler.NewInstruction(Expr(:call, Diffractor.ZeroBundle{1}, arg), Any))
8+
end
9+
10+
@testset "Constructors in forward_diff_no_inf!" begin
11+
struct Bar148
12+
v
13+
end
14+
foo_148(x) = Bar148(x)
15+
16+
ir = first(only(Base.code_ircode(foo_148, Tuple{Float64})))
17+
Diffractor.forward_diff_no_inf!(ir, [Core.SSAValue(1) => 1]; transform! = identity_transform!)
18+
ir2 = Core.Compiler.compact!(ir)
19+
f = Core.OpaqueClosure(ir2; do_compile=false)
20+
@test f(1.0) == Bar148(1.0) # This would error if we were not handling constructors (%new) right
21+
end
22+
23+
@testset "nonconst globals in forward_diff_no_inf!" begin
24+
@eval global _coeff::Float64=24.5
25+
plus_a_global(x) = x + _coeff
26+
27+
ir = first(only(Base.code_ircode(plus_a_global, Tuple{Float64})))
28+
Diffractor.forward_diff_no_inf!(ir, [Core.SSAValue(1) => 1]; transform! = identity_transform!)
29+
ir2 = Core.Compiler.compact!(ir)
30+
Core.Compiler.verify_ir(ir2) # This would error if we were not handling nonconst globals correctly
31+
f = Core.OpaqueClosure(ir2; do_compile=false)
32+
@test f(3.5) == 28.0
33+
end
34+
35+
@testset "runs of phi nodes" begin
36+
function phi_run(x::Float64)
37+
a = 2.0
38+
b = 2.0
39+
if (@noinline rand()) < 0 # this branch will never actually be taken
40+
a = -100.0
41+
b = 200.0
42+
end
43+
return x - a + b
44+
end
45+
46+
input_ir = first(only(Base.code_ircode(phi_run, Tuple{Float64})))
47+
ir = copy(input_ir)
48+
#Workout where to diff to trigger error
49+
diff_ssa = Core.SSAValue[]
50+
for idx in 1:length(ir.stmts)
51+
if ir.stmts[idx][:inst] isa Core.PhiNode
52+
push!(diff_ssa, Core.SSAValue(idx))
53+
end
54+
end
55+
56+
Diffractor.forward_diff_no_inf!(ir, diff_ssa .=> 1; transform! = identity_transform!)
57+
ir2 = Core.Compiler.compact!(ir)
58+
Core.Compiler.verify_ir(ir2) # This would error if we were not handling nonconst phi nodes correctly (after https://github.com/JuliaLang/julia/pull/50158)
59+
f = Core.OpaqueClosure(ir2; do_compile=false)
60+
@test f(3.5) == 3.5 # this will segfault if we are not handling phi nodes correctly
61+
end
62+
end

‎test/regression.jl

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
module forward
2+
using Diffractor
3+
using Diffractor: var"'", ∂⃖, DiffractorRuleConfig
4+
using ChainRules
5+
using ChainRulesCore
6+
using ChainRulesCore: ZeroTangent, NoTangent, frule_via_ad, rrule_via_ad
7+
using LinearAlgebra
8+
9+
using Test
10+
11+
const fwd = Diffractor.PrimeDerivativeFwd
12+
const bwd = Diffractor.PrimeDerivativeBack
13+
14+
15+
# Regression tests
16+
@test gradient(x -> sum(abs2, x .+ 1.0), zeros(3))[1] == [2.0, 2.0, 2.0]
17+
18+
function f_broadcast(a)
19+
l = a / 2.0 * [[0. 1. 1.]; [1. 0. 1.]; [1. 1. 0.]]
20+
return sum(l)
21+
end
22+
@test fwd(f_broadcast)(1.0) == bwd(f_broadcast)(1.0)
23+
24+
# Make sure that there's no infinite recursion in kwarg calls
25+
g_kw(;x=1.0) = sin(x)
26+
f_kw(x) = g_kw(;x)
27+
@test bwd(f_kw)(1.0) == bwd(sin)(1.0)
28+
29+
function f_crit_edge(a, b, c, x)
30+
# A function with two critical edges. This used to trigger an issue where
31+
# Diffractor would fail to insert edges for the second split critical edge.
32+
y = 1x
33+
if a && b
34+
y = 2x
35+
end
36+
if b && c
37+
y = 3x
38+
end
39+
40+
if c
41+
y = 4y
42+
end
43+
44+
return y
45+
end
46+
@test bwd(x->f_crit_edge(false, false, false, x))(1.0) == 1.0
47+
@test bwd(x->f_crit_edge(true, true, false, x))(1.0) == 2.0
48+
@test bwd(x->f_crit_edge(false, true, true, x))(1.0) == 12.0
49+
@test bwd(x->f_crit_edge(false, false, true, x))(1.0) == 4.0
50+
51+
# Issue #27 - Mixup in lifting of getfield
52+
let var"'" = bwd
53+
@test (x->x^5)''(1.0) == 20.
54+
@test_broken (x->x^5)'''(1.0) == 60.
55+
end
56+
57+
# Issue #38 - Splatting arrays
58+
@test gradient(x -> max(x...), (1,2,3))[1] == (0.0, 0.0, 1.0)
59+
@test gradient(x -> max(x...), [1,2,3])[1] == [0.0, 0.0, 1.0]
60+
61+
# Issue #40 - Symbol type parameters not properly quoted
62+
@test Diffractor.∂⃖recurse{1}()(Val{:transformations})[1] === Val{:transformations}()
63+
64+
# PR #43
65+
loss(res, z, w) = sum(res.U * Diagonal(res.S) * res.V) + sum(res.S .* w)
66+
x43 = rand(10, 10)
67+
@test Diffractor.gradient(x->loss(svd(x), x[:,1], x[:,2]), x43) isa Tuple{Matrix{Float64}}
68+
69+
# PR # 45 - Calling back into AD from ChainRules
70+
y45, back45 = rrule_via_ad(DiffractorRuleConfig(), x -> log(exp(x)), 2)
71+
@test y45 ≈ 2.0
72+
@test back45(1) == (ZeroTangent(), 1.0)
73+
74+
z45, delta45 = frule_via_ad(DiffractorRuleConfig(), (0,1), x -> log(exp(x)), 2)
75+
@test z45 ≈ 2.0
76+
@test delta45 ≈ 1.0
77+
78+
# PR #82 - getindex on non-numeric arrays
79+
@test gradient(ls -> ls[1](1.), [Base.Fix1(*, 1.)])[1][1] isa Tangent{<:Base.Fix1}
80+
81+
@testset "broadcast" begin
82+
@test gradient(x -> sum(x ./ x), [1,2,3]) == ([0,0,0],) # derivatives_given_output
83+
@test gradient(x -> sum(sqrt.(atan.(x, transpose(x)))), [1,2,3])[1] ≈ [0.2338, -0.0177, -0.0661] atol=1e-3
84+
@test gradient(x -> sum(exp.(log.(x))), [1,2,3]) == ([1,1,1],)
85+
86+
@test gradient(x -> sum((exp∘log).(x)), [1,2,3]) == ([1,1,1],) # frule_via_ad
87+
exp_log(x) = exp(log(x))
88+
@test gradient(x -> sum(exp_log.(x)), [1,2,3]) == ([1,1,1],)
89+
@test gradient((x,y) -> sum(x ./ y), [1 2; 3 4], [1,2]) == ([1 1; 0.5 0.5], [-3, -1.75])
90+
@test gradient((x,y) -> sum(x ./ y), [1 2; 3 4], 5) == ([0.2 0.2; 0.2 0.2], -0.4)
91+
@test gradient(x -> sum((y -> y/x).([1,2,3])), 4) == (-0.375,) # closure
92+
93+
@test gradient(x -> sum(sum, (x,) ./ x), [1,2,3])[1] ≈ [-4.1666, 0.3333, 1.1666] atol=1e-3 # array of arrays
94+
@test gradient(x -> sum(sum, Ref(x) ./ x), [1,2,3])[1] ≈ [-4.1666, 0.3333, 1.1666] atol=1e-3
95+
@test gradient(x -> sum(sum, (x,) ./ x), [1,2,3])[1] ≈ [-4.1666, 0.3333, 1.1666] atol=1e-3
96+
@test gradient(x -> sum(sum, (x,) .* transpose(x)), [1,2,3])[1] ≈ [12, 12, 12] # must not take the * fast path
97+
98+
@test gradient(x -> sum(x ./ 4), [1,2,3]) == ([0.25, 0.25, 0.25],)
99+
@test gradient(x -> sum([1,2,3] ./ x), 4) == (-0.375,) # x/y rule
100+
@test gradient(x -> sum(x.^2), [1,2,3]) == ([2.0, 4.0, 6.0],) # x.^2 rule
101+
@test gradient(x -> sum([1,2,3] ./ x.^2), 4) == (-0.1875,) # scalar^2 rule
102+
103+
@test gradient(x -> sum((1,2,3) .- x), (1,2,3)) == (Tangent{Tuple{Int,Int,Int}}(-1.0, -1.0, -1.0),)
104+
@test gradient(x -> sum(transpose([1,2,3]) .- x), (1,2,3)) == (Tangent{Tuple{Int,Int,Int}}(-3.0, -3.0, -3.0),)
105+
@test gradient(x -> sum([1 2 3] .+ x .^ 2), (1,2,3)) == (Tangent{Tuple{Int,Int,Int}}(6.0, 12.0, 18.0),)
106+
107+
@test gradient(x -> sum(x .> 2), [1,2,3]) |> only |> iszero # Bool output
108+
@test gradient(x -> sum(1 .+ iseven.(x)), [1,2,3]) |> only |> iszero
109+
@test gradient((x,y) -> sum(x .== y), [1,2,3], [1 2 3]) == (NoTangent(), NoTangent())
110+
@test gradient(x -> sum(x .+ [1,2,3]), true) |> only |> iszero # Bool input
111+
@test gradient(x -> sum(x ./ [1,2,3]), [true false]) |> only |> iszero
112+
@test gradient(x -> sum(x .* transpose([1,2,3])), (true, false)) |> only |> iszero
113+
114+
tup_adj = gradient((x,y) -> sum(2 .* x .+ log.(y)), (1,2), transpose([3,4,5]))
115+
@test tup_adj[1] == Tangent{Tuple{Int64, Int64}}(6.0, 6.0)
116+
@test tup_adj[2] ≈ [0.6666666666666666 0.5 0.4]
117+
@test tup_adj[2] isa Transpose
118+
@test gradient(x -> sum(atan.(x, (1,2,3))), Diagonal([4,5,6]))[1] isa Diagonal
119+
120+
@test gradient(x -> sum((y -> (x*y)).([1,2,3])), 4.0) == (6.0,) # closure
121+
end
122+
123+
@testset "broadcast, 2nd order" begin
124+
@test gradient(x -> gradient(y -> sum(y .* y), x)[1] |> sum, [1,2,3.0])[1] == [2,2,2] # calls "split broadcasting generic" with f = unthunk
125+
@test gradient(x -> gradient(y -> sum(y .* x), x)[1].^3 |> sum, [1,2,3.0])[1] == [3,12,27]
126+
@test_broken gradient(x -> gradient(y -> sum(y .* 2 .* y'), x)[1] |> sum, [1,2,3.0])[1] == [12, 12, 12] # Control flow support not fully implemented yet for higher-order
127+
128+
@test_broken gradient(x -> sum(gradient(x -> sum(x .^ 2 .+ x'), x)[1]), [1,2,3.0])[1] == [6,6,6] # BoundsError: attempt to access 18-element Vector{Core.Compiler.BasicBlock} at index [0]
129+
@test_broken gradient(x -> sum(gradient(x -> sum((x .+ 1) .* x .- x), x)[1]), [1,2,3.0])[1] == [2,2,2]
130+
@test_broken gradient(x -> sum(gradient(x -> sum(x .* x ./ 2), x)[1]), [1,2,3.0])[1] == [1,1,1]
131+
132+
@test_broken gradient(x -> sum(gradient(x -> sum(exp.(x)), x)[1]), [1,2,3])[1] ≈ exp.(1:3) # MethodError: no method matching copy(::Nothing)
133+
@test_broken gradient(x -> sum(gradient(x -> sum(atan.(x, x')), x)[1]), [1,2,3.0])[1] ≈ [0,0,0]
134+
@test_broken gradient(x -> sum(gradient(x -> sum(transpose(x) .* x), x)[1]), [1,2,3]) == ([6,6,6],) # accum(a::Transpose{Float64, Vector{Float64}}, b::ChainRulesCore.Tangent{Transpose{Int64, Vector{Int64}}, NamedTuple{(:parent,), Tuple{ChainRulesCore.NoTangent}}})
135+
@test_broken gradient(x -> sum(gradient(x -> sum(transpose(x) ./ x.^2), x)[1]), [1,2,3])[1] ≈ [27.675925925925927, -0.824074074074074, -2.1018518518518516]
136+
137+
@test_broken gradient(z -> gradient(x -> sum((y -> (x^2*y)).([1,2,3])), z)[1], 5.0) == (12.0,)
138+
end
139+
140+
end

‎test/reverse.jl

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
module reverse
2+
using Diffractor
3+
using Diffractor: var"'", ∂⃖, DiffractorRuleConfig
4+
using ChainRules
5+
using ChainRulesCore
6+
using ChainRulesCore: ZeroTangent, NoTangent, frule_via_ad, rrule_via_ad
7+
using Symbolics
8+
using LinearAlgebra
9+
10+
using Test
11+
12+
const fwd = Diffractor.PrimeDerivativeFwd
13+
const bwd = Diffractor.PrimeDerivativeBack
14+
15+
# Unit tests
16+
function tup2(f)
17+
a, b = ∂⃖{2}()(f, 1)
18+
c, d = b((2,))
19+
e, f = d(ZeroTangent(), 3)
20+
f((4,))
21+
end
22+
23+
@test tup2(tuple) == (NoTangent(), 4)
24+
25+
my_tuple(args...) = args
26+
ChainRules.rrule(::typeof(my_tuple), args...) = args, Δ->Core.tuple(NoTangent(), Δ...)
27+
@test tup2(my_tuple) == (ZeroTangent(), 4)
28+
29+
# Check characteristic of exp rule
30+
# broken see: https://github.com/JuliaDiff/Diffractor.jl/pull/142
31+
@variables ω α β γ δ ϵ ζ η
32+
@test_broken ((x1, c1) = ∂⃖{3}()(exp, ω)) isa Any
33+
#==
34+
@test isequal(simplify(x1), simplify(exp(ω)))
35+
((_, x2), c2) = c1(α)
36+
@test isequal(simplify(x2), simplify(α*exp(ω)))
37+
(x3, c3) = c2(ZeroTangent(), β)
38+
@test isequal(simplify(x3), simplify(β*exp(ω)))
39+
((_, x4), c4) = c3(γ)
40+
@test isequal(simplify(x4), simplify(exp(ω)*(γ + (α*β))))
41+
(x5, c5) = c4(ZeroTangent(), δ)
42+
@test isequal(simplify(x5), simplify(δ*exp(ω)))
43+
((_, x6), c6) = c5(ϵ)
44+
@test isequal(simplify(x6), simplify(ϵ*exp(ω) + α*δ*exp(ω)))
45+
(x7, c7) = c6(ZeroTangent(), ζ)
46+
@test isequal(simplify(x7), simplify(ζ*exp(ω) + β*δ*exp(ω)))
47+
(_, x8) = c7(η)
48+
@test isequal(simplify(x8), simplify((η + (α*ζ) + (β*ϵ) + (δ*(γ + (α*β))))*exp(ω)))
49+
==#
50+
51+
52+
function simple_control_flow(b, x)
53+
if b
54+
return sin(x)
55+
else
56+
return cos(x)
57+
end
58+
end
59+
60+
function myprod(xs)
61+
s = 1
62+
for x in xs
63+
s *= x
64+
end
65+
return s
66+
end
67+
68+
function mypow(x, n)
69+
r = one(x)
70+
while n > 0
71+
n -= 1
72+
r *= x
73+
end
74+
return r
75+
end
76+
77+
function times_three_while(x)
78+
z = x
79+
i = 3
80+
while i > 1
81+
z += x
82+
i -= 1
83+
end
84+
z
85+
end
86+
87+
isa_control_flow(::Type{T}, x) where {T} = isa(x, T) ? x : T(x)
88+
89+
# Simple Reverse Mode tests
90+
let var"'" = Diffractor.PrimeDerivativeBack
91+
# Integration tests
92+
@test @inferred(sin'(1.0)) == cos(1.0)
93+
@test sin''(1.0) == -sin(1.0)
94+
@test sin'''(1.0) == -cos(1.0)
95+
# TODO These currently cause segfaults c.f. https://github.com/JuliaLang/julia/pull/48742
96+
# @test sin''''(1.0) == sin(1.0)
97+
# @test sin'''''(1.0) == cos(1.0)
98+
# @test sin''''''(1.0) == -sin(1.0)
99+
100+
f_getfield(x) = getfield((x,), 1)
101+
@test f_getfield'(1) == 1
102+
@test f_getfield''(1) == 0
103+
@test f_getfield'''(1) == 0
104+
105+
# Higher order mixed mode tests
106+
107+
complicated_2sin(x) = (x = map(sin, Diffractor.xfill(x, 2)); x[1] + x[2])
108+
@test @inferred(complicated_2sin'(1.0)) == 2sin'(1.0)
109+
@test @inferred(complicated_2sin''(1.0)) == 2sin''(1.0) broken=true
110+
@test @inferred(complicated_2sin'''(1.0)) == 2sin'''(1.0) broken=true
111+
# TODO This currently causes a segfault, c.f. https://github.com/JuliaLang/julia/pull/48742
112+
# @test @inferred(complicated_2sin''''(1.0)) == 2sin''''(1.0) broken=true
113+
114+
# Control flow cases
115+
@test @inferred((x->simple_control_flow(true, x))'(1.0)) == sin'(1.0)
116+
@test @inferred((x->simple_control_flow(false, x))'(1.0)) == cos'(1.0)
117+
@test (x->sum(isa_control_flow(Matrix{Float64}, x)))'(Float32[1 2;]) == [1.0 1.0;]
118+
@test times_three_while'(1.0) == 3.0
119+
120+
pow5p(x) = (x->mypow(x, 5))'(x)
121+
@test pow5p(1.0) == 5.0
122+
end
123+
124+
125+
end

0 commit comments

Comments
 (0)