Some basic svd forward rules and tests by kshyatt · Pull Request #247 · QuantumKitHub/MatrixAlgebraKit.jl

kshyatt · 2026-06-08T09:45:16Z

Definitely not optimized...

kshyatt · 2026-06-11T09:59:31Z

Took another look at this. The Enzyme tests seem to be failing because of finite differences, for svd_full for example when size(A) == (17, 19), the finite differences result for dU is all over the place compared to that from the rule for the section of U which is only present in the full case. I'll try to think of a nice way to handle this, I think it is not occurring for Mooncake because Mooncake uses a different technique for checking against FD.

github-actions · 2026-06-11T13:33:40Z

Your PR no longer requires formatting changes. Thank you for your contribution!

Co-authored-by: Jutho <Jutho@users.noreply.github.com>

Jutho · 2026-06-15T09:20:45Z

In the gauge fixing parts, we could use more views for the slicing ΔU₁[I] etc, but I did not want to push my luck on the GPU with views into views using CartesianIndices.

Jutho · 2026-06-15T09:22:40Z

Anyway, I will leave this aside now for a while and look at some other PRs. In principle, tests can now be expanded to cover the complex case and svd_compact of rectangular matrices as well (but as always, I would only consider double precision for finite difference comparisons).

lkdvos

Left some final small comments, otherwise good to go?

lkdvos · 2026-06-15T13:59:19Z

+        hUᴴΔAV₁ = inv_safe.(transpose(S₁) .- S₁) .* project_hermitian(UᴴΔAV₁)
+        aUᴴΔAV₁ = inv_safe.(transpose(S₁) .+ S₁) .* project_antihermitian(UᴴΔAV₁)


I think below only the sum and difference are actually used, we could use a kernel like

MatrixAlgebraKit.jl/src/implementations/polar.jl

Lines 209 to 220 in e271b4a

function _avgdiff!(A::AbstractArray, B::AbstractArray)

axes(A) == axes(B) || throw(DimensionMismatch())

@simd for I in eachindex(A, B)

@inbounds begin

a = A[I]

b = B[I]

A[I] = (a + b) / 2

B[I] = b - a

end

end

return A, B

end

to avoid the two extra allocations, but I'm also happy to just leave them as-is, it's hard to imagine this really making that huge of a difference

kshyatt · 2026-06-15T14:03:47Z

The GPU tests will probably fail until a new version is tagged at GPUArrays (JuliaGPU/GPUArrays.jl#738)

kshyatt · 2026-06-17T18:37:51Z

Ran it locally and it worked, let's hope!

Jutho · 2026-06-18T20:34:13Z

+            S, dS = arrayify(USVᴴ[2], dUSVᴴ[2])
+            Vᴴ, dVᴴ = arrayify(USVᴴ[3], dUSVᴴ[3])
+            $f!(A, USVᴴ, Mooncake.primal(alg_dalg))
+            svd_pushforward!(dA, A, (U, S, Vᴴ), (dU, dS, dVᴴ))


This again works because svd_pushforward! doesn't actually need A, since A is destroyed at this point. Not sure if it is worth to add a comment.

Jutho · 2026-06-18T20:43:39Z

+# have to override this as methods are missing in GPUArrays for the various
+# views of Diagonal of ΔA


This is a bit of a confusing comment: what exactly is missing?
Might be useful to keep track of this, in case it gets fixed.

It's again a situation of mul!(::Diagonal{T, CuVector{T}}, [horrific view of adjoint of view], CuArray) which GPUArrays cannot dispatch onto at all.

kshyatt · 2026-06-19T09:29:54Z

Of course when I ran this locally (5 times!!!) it passed every single time.

kshyatt · 2026-06-19T09:30:15Z

Some of these look like tolerance issues to me at least

Jutho · 2026-06-19T09:36:50Z

        if eltype(U) <: Complex && !iszerotangent(ΔU) && !iszerotangent(ΔVᴴ) # fix gauge for `gaugefix!` compatibility
            _, I = findmax(abs, U₁; dims = 1)
-            infinitesimal_phases = imag.(ΔU₁[I] ./ U₁[I])
+            infinitesimal_phases = imag.(ΔU₁[I] .* inv_safe.(U₁[I]))


I would be suprised if that is needed or makes a difference. U₁[I] is the maximum element of every column of U₁. As U₁ is an isometric matrix, all of its columns have norm 1, and therefore, the largest element needs to be at least 1/sqrt(m), in magnitude, with m the number of rows. Typically, it will be larger.

I would be quite surprised too but I'm otherwise very confused where the NaN are emerging from (the ./transpose(S1) line has the same objection). I'll try stepping through the pushfoward to see if I can find the culprit.

Jutho · 2026-06-19T09:37:26Z

I am also happy to take a look at the test failures.

kshyatt · 2026-06-19T09:39:36Z

I am also happy to take a look at the test failures.

If you like! I'm about to get on a plane and unsure if it will have WiFi with which to push if I can find the problem

Jutho · 2026-06-21T00:04:56Z

Ok, I located the problem, but also started changing various other things so this turned out a bit bigger than anticipated. The problem was that the initialization of the tangent in the frule is random introducing NaNs, subnormals and other random stuff. But svd_pushforward only touches the diagonal of S, and did not set the off-diagonal entries (in the case of svd_full) to zero.

Jutho · 2026-06-21T00:06:24Z

I also tried to make some steps in the various implementations more consistent, and to simplify them a bit here and there. I introduced a has_equal_storage for simplifying and generalizing the detection of the input matrix A overlapping with some of the output arguments.

lkdvos · 2026-06-21T16:51:59Z

        ∂K .*= inv_safe.(transpose(diagview(D)) .- diagview(D), degeneracy_atol)
-        mul!(ΔV, V, ∂K, 1, 0)
+        mul!(ΔV, V, ∂K)
+        if eltype(V) <: Complex # fix gauge for `gaugefix!` compatibility


Are we enforcing that the primal computation is always gauge fixing, and is that relevant for this?

No we are not enforcing this, but it is the default. This is relevant to the extent that if gaugefix = false, there is no way to compute the pushforward at all. So we could throw a warning. There is some asymmetry between the pullbacks and the pushforwards in that the latter don't throw warnings for cases where there is gauge freedom.

But the way to deal with this is also very different between the two.

lkdvos

Minor detail, but otherwise this looks ready to me?

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

Jutho · 2026-06-24T10:35:31Z

 end
 has_equal_storage(A::AbstractMatrix, B::AbstractVector) = false
-has_equal_storage(A::AbstractVector, B::AbstractMatrix) = false
+has_equal_storage(A::AbstractVector, B::AbstractMatrix) = false


So this was my formatting issue? I don't actually see any change? Is it some invisible character?

I think it's that we were missing a \n after the false

kshyatt force-pushed the ksh/svd_fwd branch from b919563 to c88e5c9 Compare June 8, 2026 09:46

Jutho reviewed Jun 9, 2026

View reviewed changes

Comment thread src/pushforwards/svd.jl Outdated

lkdvos reviewed Jun 11, 2026

View reviewed changes

Comment thread src/pushforwards/svd.jl Outdated

kshyatt force-pushed the ksh/svd_fwd branch from 35571fa to fc4f760 Compare June 11, 2026 13:26

Katharine Hyatt and others added 9 commits June 12, 2026 13:35

Some basic svd forward rules and tests

33ab47b

Try to fix Enzyme?

838ecab

Use GPU safe rank calculation

be9afda

Another fix for GPU

883dadb

Some more fixes

1e03d65

Update src/pushforwards/svd.jl

d8c9427

Co-authored-by: Jutho <Jutho@users.noreply.github.com>

A little cleanup

ec2a674

Some small fixes again

9d0f666

Use sylvester fallback

8940756

kshyatt force-pushed the ksh/svd_fwd branch from f446751 to 8940756 Compare June 12, 2026 11:35

kshyatt and others added 2 commits June 12, 2026 17:18

Fix diagms

aa3a2ee

Add a one! method for vectors

267886d

Jutho reviewed Jun 15, 2026

View reviewed changes

Comment thread src/common/initialization.jl Outdated

Use UniformScaling for diagm

64c849f

Jutho reviewed Jun 15, 2026

View reviewed changes

Comment thread src/pushforwards/svd.jl Outdated

some changes

dc7bbeb

kshyatt commented Jun 15, 2026

View reviewed changes

Comment thread src/pushforwards/svd.jl Outdated

more gauge fixing in eig/eigh

c5cf3de

lkdvos reviewed Jun 15, 2026

View reviewed changes

kshyatt added 2 commits June 17, 2026 02:38

Get stuff working on CUDA

17c2158

Fix missing import

c89643d

Jutho reviewed Jun 17, 2026

View reviewed changes

Comment thread test/testsuite/mooncake/svd.jl Outdated

Jutho reviewed Jun 17, 2026

View reviewed changes

Comment thread test/testsuite/mooncake/svd.jl Outdated

Run more complex test cases

fbcd972

Jutho reviewed Jun 18, 2026

View reviewed changes

Comment thread ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl Outdated

The case of the missing excalamtion point

dd2c5a3

Jutho reviewed Jun 18, 2026

View reviewed changes

Comment thread ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl Outdated

Jutho reviewed Jun 18, 2026

View reviewed changes

Comment thread ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl Outdated

More touchups for comments

36a22ca

Try invsafe for gaugefixing in pushforward

10ca00a

Jutho reviewed Jun 19, 2026

View reviewed changes

various changes

0af8ec4

lkdvos reviewed Jun 21, 2026

View reviewed changes

Format and zero_dual touchups

d5cb3c5

lkdvos approved these changes Jun 23, 2026

View reviewed changes

Comment thread ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl Outdated

Update ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

7dfc1c4

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

lkdvos approved these changes Jun 23, 2026

View reviewed changes

kshyatt enabled auto-merge (squash) June 23, 2026 23:55

kshyatt disabled auto-merge June 23, 2026 23:55

kshyatt merged commit 0dfcb52 into main Jun 23, 2026
36 checks passed

kshyatt deleted the ksh/svd_fwd branch June 23, 2026 23:55

Jutho reviewed Jun 24, 2026

View reviewed changes

		hUᴴΔAV₁ = inv_safe.(transpose(S₁) .- S₁) .* project_hermitian(UᴴΔAV₁)
		aUᴴΔAV₁ = inv_safe.(transpose(S₁) .+ S₁) .* project_antihermitian(UᴴΔAV₁)

	function _avgdiff!(A::AbstractArray, B::AbstractArray)
	axes(A) == axes(B) \|\| throw(DimensionMismatch())
	@simd for I in eachindex(A, B)
	@inbounds begin
	a = A[I]
	b = B[I]
	A[I] = (a + b) / 2
	B[I] = b - a
	end
	end
	return A, B
	end

		# have to override this as methods are missing in GPUArrays for the various
		# views of Diagonal of ΔA

Uh oh!

Conversation

kshyatt commented Jun 8, 2026

Uh oh!

Uh oh!

kshyatt commented Jun 11, 2026

Uh oh!

Uh oh!

github-actions Bot commented Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Jutho commented Jun 15, 2026

Uh oh!

Jutho commented Jun 15, 2026

Uh oh!

lkdvos left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kshyatt commented Jun 15, 2026

Uh oh!

Uh oh!

Uh oh!

kshyatt commented Jun 17, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

kshyatt commented Jun 19, 2026

Uh oh!

kshyatt commented Jun 19, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jutho commented Jun 19, 2026

Uh oh!

kshyatt commented Jun 19, 2026

Uh oh!

Jutho commented Jun 21, 2026

Uh oh!

Jutho commented Jun 21, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lkdvos left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

github-actions Bot commented Jun 11, 2026 •

edited

Loading