Skip to content

Commit 5cdd986

Browse files
kshyattlkdvos
kshyatt
authored and
lkdvos
committed
Add Mooncake tests for Diagonal
Fix stupid typos Diagonal + eig(h) working SVD working for diag LQ/QR working Polar and orthnull Some Mooncake progress Formatter Working Mooncake eig Remove comment No diag for orthnull Diag tests working for Mooncake Formatter Update ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl Co-authored-by: Lukas Devos <ldevos98@gmail.com> A bit of Enzyme cleanup A little Mooncake cleanup
1 parent dd05ae9 commit 5cdd986

20 files changed

Lines changed: 180 additions & 83 deletions

File tree

‎ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl‎

Lines changed: 39 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -43,37 +43,6 @@ using LinearAlgebra
4343
# other provided input variable.
4444
#--------------------------------------------
4545

46-
# this rule is necessary for now as without it,
47-
# a segfault occurs both on 1.10 and 1.12 -- likely
48-
# a deeper internal bug
49-
function EnzymeRules.augmented_primal(
50-
config::EnzymeRules.RevConfigWidth{1},
51-
func::Const{typeof(copy_input)},
52-
::Type{RT},
53-
f::Annotation,
54-
A::Annotation
55-
) where {RT}
56-
ret = func.val(f.val, A.val)
57-
primal = EnzymeRules.needs_primal(config) ? ret : nothing
58-
shadow = EnzymeRules.needs_shadow(config) ? zero(A.dval) : nothing
59-
return EnzymeRules.AugmentedReturn(primal, shadow, shadow)
60-
end
61-
62-
function EnzymeRules.reverse(
63-
config::EnzymeRules.RevConfigWidth{1},
64-
func::Const{typeof(copy_input)},
65-
::Type{RT},
66-
cache,
67-
f::Annotation,
68-
A::Annotation
69-
) where {RT}
70-
copy_shadow = cache
71-
if !isa(A, Const) && !isnothing(copy_shadow)
72-
A.dval .+= copy_shadow
73-
end
74-
return (nothing, nothing)
75-
end
76-
7746
# two-argument factorizations like LQ, QR, EIG
7847
for (f, pb) in (
7948
(qr_full!, qr_pullback!),
@@ -101,9 +70,14 @@ for (f, pb) in (
10170
# if arg.val == ret, the annotation must be Duplicated or DuplicatedNoNeed
10271
# if arg isa Const, ret may still be modified further down the call graph so we should
10372
# copy it to protect ourselves
104-
cache_arg = (arg.val !== ret) || EnzymeRules.overwritten(config)[3] ? copy.(ret) : nothing
105-
dret = if EnzymeRules.needs_shadow(config)
106-
(TA == Nothing && TB == Nothing) || isa(arg, Const) ? zero.(ret) : arg.dval
73+
A_is_arg1 = !isa(A, Const) && A.val === arg.val[1]
74+
A_is_arg2 = !isa(A, Const) && A.val === arg.val[2]
75+
A_is_arg = A_is_arg1 || A_is_arg2
76+
cache_arg = (arg.val !== ret && !A_is_arg) || EnzymeRules.overwritten(config)[3] ? copy.(ret) : nothing
77+
dret = if EnzymeRules.needs_shadow(config) && ((TA == Nothing && TB == Nothing) || isa(arg, Const))
78+
make_zero.(ret)
79+
elseif EnzymeRules.needs_shadow(config)
80+
arg.dval
10781
else
10882
nothing
10983
end
@@ -125,11 +99,19 @@ for (f, pb) in (
12599
# use A (so that whoever does this is forced to handle caching A
126100
# appropriately here)
127101
Aval = nothing
102+
A_is_arg1 = !isa(A, Const) && A.dval === arg.dval[1]
103+
A_is_arg2 = !isa(A, Const) && A.dval === arg.dval[2]
104+
A_is_arg = A_is_arg1 || A_is_arg2
128105
argval = something(cache_arg, arg.val)
129106
if !isa(A, Const)
130-
$pb(A.dval, Aval, argval, darg)
107+
ΔA = A_is_arg ? make_zero(A.dval) : A.dval
108+
$pb(ΔA, Aval, argval, darg)
109+
A_is_arg && (A.dval .= ΔA)
110+
end
111+
if !isa(arg, Const)
112+
A_is_arg1 || make_zero!(arg.dval[1])
113+
A_is_arg2 || make_zero!(arg.dval[2])
131114
end
132-
!isa(arg, Const) && make_zero!(arg.dval)
133115
return (nothing, nothing, nothing)
134116
end
135117
end
@@ -343,7 +325,13 @@ for (f, trunc_f, full_f, pb) in (
343325
if !isa(A, Const)
344326
$pb(A.dval, Aval, DVval, dDVtrunc, ind)
345327
end
346-
!isa(DV, Const) && make_zero!(DV.dval)
328+
if !isa(DV, Const)
329+
if A.dval !== DV.dval[1]
330+
make_zero!(DV.dval)
331+
else
332+
make_zero!(DV.dval[2])
333+
end
334+
end
347335
return (nothing, nothing, nothing)
348336
end
349337
end
@@ -379,21 +367,24 @@ for (f!, f_full!, pb!) in (
379367
func::Const{typeof($f!)},
380368
::Type{RT},
381369
cache,
382-
A::Annotation,
370+
A::Annotation{TA},
383371
D::Annotation,
384372
alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
385-
) where {RT}
373+
) where {RT, TA}
386374
cache_D, dD, V = cache
387375
Dval = something(cache_D, D.val)
388376
# A is NOT used in the pullback, so we assign Aval = nothing
389377
# to trigger an error in case the pullback is modified to directly
390378
# use A (so that whoever does this is forced to handle caching A
391379
# appropriately here)
392380
Aval = nothing
381+
A_is_arg = !isa(A, Const) && TA <: Diagonal && diagview(A.dval) === D.dval
393382
if !isa(A, Const)
394-
$pb!(A.dval, Aval, (Diagonal(Dval), V), dD)
383+
ΔA = A_is_arg ? make_zero(A.dval) : A.dval
384+
$pb!(ΔA, Aval, (Diagonal(Dval), V), dD)
385+
A_is_arg && (A.dval .= ΔA)
395386
end
396-
!isa(D, Const) && make_zero!(D.dval)
387+
!isa(D, Const) && !A_is_arg && make_zero!(D.dval)
397388
return (nothing, nothing, nothing)
398389
end
399390
end
@@ -425,21 +416,24 @@ function EnzymeRules.reverse(
425416
func::Const{typeof(svd_vals!)},
426417
::Type{RT},
427418
cache,
428-
A::Annotation,
419+
A::Annotation{TA},
429420
S::Annotation,
430421
alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
431-
) where {RT}
422+
) where {RT, TA}
432423
cache_S, dS, U, Vᴴ = cache
433424
# A is NOT used in the pullback, so we assign Aval = nothing
434425
# to trigger an error in case the pullback is modified to directly
435426
# use A (so that whoever does this is forced to handle caching A
436427
# appropriately here)
437428
Aval = nothing
438429
Sval = something(cache_S, S.val)
430+
A_is_arg = !isa(A, Const) && TA <: Diagonal && diagview(A.dval) === S.dval
439431
if !isa(A, Const)
440-
svd_vals_pullback!(A.dval, Aval, (U, Diagonal(Sval), Vᴴ), dS)
432+
ΔA = A_is_arg ? make_zero(A.dval) : A.dval
433+
svd_vals_pullback!(ΔA, Aval, (U, Diagonal(Sval), Vᴴ), dS)
434+
A_is_arg && (A.dval .= ΔA)
441435
end
442-
!isa(S, Const) && make_zero!(S.dval)
436+
!isa(S, Const) && !A_is_arg && make_zero!(S.dval)
443437
return (nothing, nothing, nothing)
444438
end
445439

‎ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl‎

Lines changed: 79 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -15,19 +15,6 @@ using LinearAlgebra
1515

1616
Mooncake.tangent_type(::Type{<:MatrixAlgebraKit.AbstractAlgorithm}) = Mooncake.NoTangent
1717

18-
@is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(copy_input), Any, Any}
19-
function Mooncake.rrule!!(::CoDual{typeof(copy_input)}, f_df::CoDual, A_dA::CoDual)
20-
Ac = copy_input(Mooncake.primal(f_df), Mooncake.primal(A_dA))
21-
Ac_dAc = Mooncake.zero_fcodual(Ac)
22-
dAc = Mooncake.tangent(Ac_dAc)
23-
function copy_input_pb(::NoRData)
24-
Mooncake.increment!!(Mooncake.tangent(A_dA), dAc)
25-
return NoRData(), NoRData(), NoRData()
26-
end
27-
return Ac_dAc, copy_input_pb
28-
end
29-
30-
Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(initialize_output), Any, Any, Any}
3118
# two-argument in-place factorizations like LQ, QR, EIG
3219
for (f!, f, pb, adj) in (
3320
(:qr_full!, :qr_full, :qr_pullback!, :qr_adjoint),
@@ -53,13 +40,29 @@ for (f!, f, pb, adj) in (
5340
arg2c = copy(arg2)
5441
$f!(A, args, Mooncake.primal(alg_dalg))
5542
function $adj(::NoRData)
56-
copy!(A, Ac)
57-
$pb(dA, A, (arg1, arg2), (darg1, darg2))
58-
copy!(arg1, arg1c)
43+
# DON'T copy Ac to A if A === one
44+
# of the output args -- this can
45+
# mess up the pullback because
46+
# generally the args are used there
47+
if !(A === arg1 || A === arg2)
48+
copy!(A, Ac)
49+
$pb(dA, A, (arg1, arg2), (darg1, darg2))
50+
else
51+
ΔA = zero(A)
52+
$pb(ΔA, A, (arg1, arg2), (darg1, darg2))
53+
dA .= ΔA
54+
end
55+
if A === arg1
56+
zero!(darg2)
57+
elseif A === arg2
58+
zero!(darg1)
59+
else
60+
zero!(darg1)
61+
zero!(darg2)
62+
end
5963
copy!(arg2, arg2c)
60-
zero!(darg1)
61-
zero!(darg2)
62-
return NoRData(), NoRData(), NoRData(), NoRData()
64+
copy!(arg1, arg1c)
65+
return ntuple(Returns(NoRData()), 4)
6366
end
6467
return args_dargs, $adj
6568
end
@@ -140,9 +143,19 @@ for (f!, f, f_full, pb, adj) in (
140143
copy!(D, diagview(DV[1]))
141144
V = DV[2]
142145
function $adj(::NoRData)
143-
$pb(dA, A, DV, dD)
144-
copy!(D, Dc)
145-
zero!(dD)
146+
if A !== D
147+
$pb(dA, A, DV, dD)
148+
else
149+
ΔA = zero(A)
150+
$pb(ΔA, A, DV, dD)
151+
dA .= A
152+
end
153+
if A !== D
154+
zero!(dD)
155+
copy!(D, Dc)
156+
else
157+
copy!(A, Ac)
158+
end
146159
return NoRData(), NoRData(), NoRData(), NoRData()
147160
end
148161
return D_dD, $adj
@@ -199,15 +212,27 @@ for f in (:eig, :eigh)
199212
# not for nested structs with various fields (like Diagonal{Complex})
200213
output_codual = Mooncake.zero_fcodual(output)
201214
function $f_adjoint!(dy::Tuple{NoRData, NoRData, <:Real})
202-
copy!(A, Ac)
203215
Dtrunc, Vtrunc, ϵ = Mooncake.primal(output_codual)
204216
dDtrunc_, dVtrunc_, dϵ = Mooncake.tangent(output_codual)
205217
_warn_pullback_truncerror(dy[3])
206218
D′, dD′ = arrayify(Dtrunc, dDtrunc_)
207219
V′, dV′ = arrayify(Vtrunc, dVtrunc_)
208-
$f_trunc_pullback!(dA, A, (D′, V′), (dD′, dV′))
209-
copy!(DV[1], DVc[1])
210-
copy!(DV[2], DVc[2])
220+
D, dD = arrayify(DV[1], dDV[1])
221+
V, dV = arrayify(DV[2], dDV[2])
222+
copy!(A, Ac)
223+
if !(A === D || A === V)
224+
$f_trunc_pullback!(dA, A, (D′, V′), (dD′, dV′))
225+
else
226+
ΔA = zero(A)
227+
$f_trunc_pullback!(ΔA, A, (D′, V′), (dD′, dV′))
228+
dA .= ΔA
229+
end
230+
if A === D
231+
copy!(DV[2], DVc[2])
232+
else
233+
copy!(DV[1], DVc[1])
234+
copy!(DV[2], DVc[2])
235+
end
211236
zero!(dD′)
212237
zero!(dV′)
213238
return NoRData(), NoRData(), NoRData(), NoRData()
@@ -239,12 +264,22 @@ for f in (:eig, :eigh)
239264
_warn_pullback_truncerror(dϵ)
240265

241266
# compute pullbacks
242-
$f_pullback!(dA, Ac, DV, dDVtrunc, ind)
243-
zero!.(dDVtrunc) # since this is allocated in this function this is probably not required
244-
267+
if !(A === DV[1] || A === DV[2])
268+
$f_pullback!(dA, Ac, DV, dDVtrunc, ind)
269+
else
270+
ΔA = zero(A)
271+
$f_pullback!(ΔA, Ac, DV, dDVtrunc, ind)
272+
dA .= ΔA
273+
end
245274
# restore state
246275
copy!(A, Ac)
247-
copy!.(DV, DVc)
276+
if A === DV[1]
277+
copy!(DV[2], DVc[2])
278+
zero!(dDV[2])
279+
else
280+
copy!.(DV, DVc)
281+
zero!.(dDV)
282+
end
248283

249284
return ntuple(Returns(NoRData()), 4)
250285
end
@@ -351,12 +386,23 @@ for f in (:eig, :eigh)
351386
dDVtrunc = last.(arrayify.(DVtrunc, Mooncake.tangent(DVtrunc_dDVtrunc)))
352387
function $f_adjoint!(::NoRData)
353388
# compute pullbacks
354-
$f_pullback!(dA, Ac, DV, dDVtrunc, ind)
355-
zero!.(dDV)
389+
if !(A === DV[1] || A === DV[2])
390+
$f_pullback!(dA, Ac, DV, dDVtrunc, ind)
391+
else
392+
ΔA = zero(A)
393+
$f_pullback!(ΔA, Ac, DV, dDVtrunc, ind)
394+
dA .= ΔA
395+
end
356396

357397
# restore state
358398
copy!(A, Ac)
359-
copy!.(DV, DVc)
399+
if A === DV[1]
400+
copy!(DV[2], DVc[2])
401+
zero!(dDV[2])
402+
else
403+
copy!.(DV, DVc)
404+
zero!.(dDV)
405+
end
360406

361407
return ntuple(Returns(NoRData()), 4)
362408
end

‎test/enzyme/eig.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...)
1515
TestSuite.seed_rng!(1234)
1616
if !is_buildkite
1717
TestSuite.test_enzyme_eig(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
18+
AT = Diagonal{T, Vector{T}}
19+
TestSuite.test_enzyme_eig(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
1820
end
1921
end

‎test/enzyme/eigh.jl‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ m = 19
1414
for T in (BLASFloats..., GenericFloats...)
1515
TestSuite.seed_rng!(1234)
1616
if !is_buildkite
17-
TestSuite.test_enzyme_eigh(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
17+
#TestSuite.test_enzyme_eigh(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
18+
AT = Diagonal{T, Vector{T}}
19+
TestSuite.test_enzyme_eigh(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
1820
end
1921
end

‎test/enzyme/lq.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
1515
TestSuite.seed_rng!(1234)
1616
if !is_buildkite
1717
TestSuite.test_enzyme_lq(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
18+
AT = Diagonal{T, Vector{T}}
19+
m == n && TestSuite.test_enzyme_lq(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
1820
end
1921
end

‎test/enzyme/orthnull.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
1515
TestSuite.seed_rng!(1234)
1616
if !is_buildkite
1717
TestSuite.test_enzyme_orthnull(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
18+
AT = Diagonal{T, Vector{T}}
19+
m == n && TestSuite.test_enzyme_orthnull(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
1820
end
1921
end

‎test/enzyme/polar.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
1515
TestSuite.seed_rng!(1234)
1616
if !is_buildkite
1717
TestSuite.test_enzyme_polar(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
18+
AT = Diagonal{T, Vector{T}}
19+
#m == n && TestSuite.test_enzyme_polar(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
1820
end
1921
end

‎test/enzyme/qr.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
1515
TestSuite.seed_rng!(1234)
1616
if !is_buildkite
1717
TestSuite.test_enzyme_qr(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
18+
AT = Diagonal{T, Vector{T}}
19+
m == n && TestSuite.test_enzyme_qr(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
1820
end
1921
end

‎test/enzyme/svd.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
1515
TestSuite.seed_rng!(1234)
1616
if !is_buildkite
1717
TestSuite.test_enzyme_svd(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
18+
AT = Diagonal{T, Vector{T}}
19+
m == n && TestSuite.test_enzyme_svd(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
1820
end
1921
end

‎test/mooncake/eig.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...)
1515
TestSuite.seed_rng!(1234)
1616
if !is_buildkite
1717
TestSuite.test_mooncake_eig(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
18+
AT = Diagonal{T, Vector{T}}
19+
TestSuite.test_mooncake_eig(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
1820
end
1921
end

0 commit comments

Comments
 (0)