Use Testsuite for AD tests #126
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
kshyatt
commented
Dec 23, 2025
| L = kron(diagm(I_n), A) + kron(adjoint(B), diagm(I_m)) | ||
| x_vec = L \ -vec(C) | ||
| X = CuMatrix(reshape(x_vec, m, n))=# | ||
| hX = sylvester(collect(A), collect(B), collect(C)) |
Member
Author
There was a problem hiding this comment.
This is very awful but I wasn't able to find a correct way to do it in five minutes so there you go
|
Your PR requires formatting changes to meet the project's style guidelines. Click here to view the suggested changes.diff --git a/ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl b/ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl
index 637012c..75701ba 100644
--- a/ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl
+++ b/ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl
@@ -17,9 +17,9 @@ Mooncake.tangent_type(::Type{<:MatrixAlgebraKit.AbstractAlgorithm}) = Mooncake.N
@is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(copy_input), Any, Any}
function Mooncake.rrule!!(::CoDual{typeof(copy_input)}, f_df::CoDual, A_dA::CoDual)
- Ac = copy_input(Mooncake.primal(f_df), Mooncake.primal(A_dA))
+ Ac = copy_input(Mooncake.primal(f_df), Mooncake.primal(A_dA))
Ac_dAc = Mooncake.zero_fcodual(Ac)
- dAc = Mooncake.tangent(Ac_dAc)
+ dAc = Mooncake.tangent(Ac_dAc)
function copy_input_pb(::NoRData)
Mooncake.increment!!(Mooncake.tangent(A_dA), dAc)
return NoRData(), NoRData(), NoRData()
diff --git a/src/pullbacks/polar.jl b/src/pullbacks/polar.jl
index e4400a3..0caed08 100644
--- a/src/pullbacks/polar.jl
+++ b/src/pullbacks/polar.jl
@@ -48,12 +48,12 @@ function right_polar_pullback!(ΔA::AbstractMatrix, A, PWᴴ, ΔPWᴴ; kwargs...
!iszerotangent(ΔP) && mul!(M, P, ΔP, -1, 1)
C = sylvester(P, P, M' - M)
C .+= ΔP
- ΔA .+= C*Wᴴ
+ ΔA .+= C * Wᴴ
if !iszerotangent(ΔWᴴ)
- PΔWᴴ = P \ ΔWᴴ
- PΔWᴴW = PΔWᴴ * Wᴴ'
+ PΔWᴴ = P \ ΔWᴴ
+ PΔWᴴW = PΔWᴴ * Wᴴ'
PΔWᴴ .-= PΔWᴴW * Wᴴ
- ΔA .+= PΔWᴴ
+ ΔA .+= PΔWᴴ
end
return ΔA
end
diff --git a/src/pullbacks/svd.jl b/src/pullbacks/svd.jl
index c87407f..bc7b9cb 100644
--- a/src/pullbacks/svd.jl
+++ b/src/pullbacks/svd.jl
@@ -90,7 +90,7 @@ function svd_pullback!(
# Add the remaining contributions
if m > r && !iszerotangent(ΔU) # remaining ΔU is already orthogonal to Ur
- Sp = view(S, indU)
+ Sp = view(S, indU)
Vᴴp = view(Vᴴ, indU, :)
ΔA .+= (ΔU ./ Sp') * Vᴴp
end
diff --git a/test/testsuite/ad_utils.jl b/test/testsuite/ad_utils.jl
index dfd61ef..70ea161 100644
--- a/test/testsuite/ad_utils.jl
+++ b/test/testsuite/ad_utils.jl
@@ -128,14 +128,14 @@ function ad_qr_rd_compact_setup(A::Diagonal)
T = eltype(A)
r = minmn - 5
Ard_ = randn!(similar(A, T, m))
- Ard_[r+1:m] .= zero(T)
- Ard = Diagonal(Ard_)
+ Ard_[(r + 1):m] .= zero(T)
+ Ard = Diagonal(Ard_)
Q, R = qr_compact(Ard)
QR = (Q, R)
ΔQ = Diagonal(randn!(similar(A.diag, T, m)))
ΔR = Diagonal(randn!(similar(A.diag, T, m)))
- diagview(ΔQ)[r+1:m] .= zero(T)
- diagview(ΔR)[r+1:m] .= zero(T)
+ diagview(ΔQ)[(r + 1):m] .= zero(T)
+ diagview(ΔR)[(r + 1):m] .= zero(T)
return (Q, R), (ΔQ, ΔR)
end
@@ -283,7 +283,7 @@ function ad_svd_compact_setup(A)
ΔS2 = Diagonal(randn!(similar(A, real(T), minmn)))
ΔVᴴ = randn!(similar(A, T, minmn, n))
U, S, Vᴴ = svd_compact(A)
- ΔU, ΔVᴴ = remove_svdgauge_dependence!(ΔU, ΔVᴴ, U, S, Vᴴ)
+ ΔU, ΔVᴴ = remove_svdgauge_dependence!(ΔU, ΔVᴴ, U, S, Vᴴ)
return (U, S, Vᴴ), (ΔU, ΔS, ΔVᴴ), (ΔU, ΔS2, ΔVᴴ)
end
@@ -296,7 +296,7 @@ function ad_svd_compact_setup(A::Diagonal)
ΔS2 = Diagonal(randn!(similar(A.diag, real(T), minmn)))
ΔVᴴ = randn!(similar(A.diag, T, m, n))
U, S, Vᴴ = svd_compact(A)
- ΔU, ΔVᴴ = remove_svdgauge_dependence!(ΔU, ΔVᴴ, U, S, Vᴴ)
+ ΔU, ΔVᴴ = remove_svdgauge_dependence!(ΔU, ΔVᴴ, U, S, Vᴴ)
return (U, S, Vᴴ), (ΔU, ΔS, ΔVᴴ), (ΔU, ΔS2, ΔVᴴ)
end
diff --git a/test/testsuite/mooncake.jl b/test/testsuite/mooncake.jl
index c587b2e..4736288 100644
--- a/test/testsuite/mooncake.jl
+++ b/test/testsuite/mooncake.jl
@@ -62,14 +62,14 @@ make_mooncake_tangent(ΔD::Diagonal{T}) where {T <: Complex} = Mooncake.build_ta
make_mooncake_tangent(T::Tuple) = Mooncake.build_tangent(typeof(T), make_mooncake_tangent.(T)...)
-make_mooncake_fdata(x) = make_mooncake_tangent(x)
+make_mooncake_fdata(x) = make_mooncake_tangent(x)
make_mooncake_fdata(x::Diagonal) = Mooncake.FData((diag = make_mooncake_tangent(x.diag),))
-make_mooncake_fdata(x::Tuple) = map(make_mooncake_fdata, x)
+make_mooncake_fdata(x::Tuple) = map(make_mooncake_fdata, x)
# no `alg` argument
function _get_copying_derivative(f_c, rrule, A, ΔA, args, Δargs, ::Nothing, rdata)
- dA_copy = make_mooncake_fdata(copy(ΔA))
- A_copy = copy(A)
+ dA_copy = make_mooncake_fdata(copy(ΔA))
+ A_copy = copy(A)
dargs_copy = make_mooncake_fdata(deepcopy(Δargs))
copy_out, copy_pb!! = rrule(Mooncake.CoDual(f_c, Mooncake.NoFData()), Mooncake.CoDual(A_copy, dA_copy), Mooncake.CoDual(args, dargs_copy))
copy_pb!!(rdata)
@@ -78,8 +78,8 @@ end
# `alg` argument
function _get_copying_derivative(f_c, rrule, A, ΔA, args, Δargs, alg, rdata)
- dA_copy = make_mooncake_fdata(copy(ΔA))
- A_copy = copy(A)
+ dA_copy = make_mooncake_fdata(copy(ΔA))
+ A_copy = copy(A)
dargs_copy = make_mooncake_fdata(deepcopy(Δargs))
copy_out, copy_pb!! = rrule(Mooncake.CoDual(f_c, Mooncake.NoFData()), Mooncake.CoDual(A_copy, dA_copy), Mooncake.CoDual(args, dargs_copy), Mooncake.CoDual(alg, Mooncake.NoFData()))
copy_pb!!(rdata) |
kshyatt
force-pushed
the
testsuite-ad
branch
2 times, most recently
from
0669d08 to
958dd36
Compare
kshyatt
force-pushed
the
testsuite-ad
branch
6 times, most recently
from
fad3f4a to
daeab4d
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This lets us reuse a lot of the setup infrastructure for ChainRules, Mooncake, and (soon) Enzyme. Also starts testing the AD rules on GPU.