From 9950d6eb3f86c4b8980cdf481bc465952a725eda Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Mon, 17 Aug 2026 09:19:40 -0400 Subject: [PATCH 1/4] Make svd_trunc pullback GPU friendly --- src/pullbacks/svd.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/pullbacks/svd.jl b/src/pullbacks/svd.jl index 4291bf9e6..365ca9feb 100755 --- a/src/pullbacks/svd.jl +++ b/src/pullbacks/svd.jl @@ -241,8 +241,9 @@ function svd_trunc_pullback!( Y₀ᴴ = iszerotangent(ΔV₊ᴴ) ? zero(Vᴴ) : ldiv!(Diagonal(S), ΔV₊ᴴ) US = mul!(ΔAV, U, Smat) # recycle ΔAV AP = mul!(copy(A), US, Vᴴ, -1, 1) - AP ./= S[end] - S⁻¹ = S[end] ./ S + minS = minimum(S) + AP ./= minS + S⁻¹ = minS ./ S X₁ = rmul!(AP * Y₀ᴴ', Diagonal(S⁻¹)) X₁ .+= X₀ Y₁ᴴ = lmul!(Diagonal(S⁻¹), X₀' * AP) From 702bcccb398cb052448e4670aa78a14b3afb7740 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 19 Aug 2026 05:49:09 -0400 Subject: [PATCH 2/4] Try a view --- src/pullbacks/svd.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pullbacks/svd.jl b/src/pullbacks/svd.jl index 365ca9feb..7e67aca7f 100755 --- a/src/pullbacks/svd.jl +++ b/src/pullbacks/svd.jl @@ -241,7 +241,7 @@ function svd_trunc_pullback!( Y₀ᴴ = iszerotangent(ΔV₊ᴴ) ? zero(Vᴴ) : ldiv!(Diagonal(S), ΔV₊ᴴ) US = mul!(ΔAV, U, Smat) # recycle ΔAV AP = mul!(copy(A), US, Vᴴ, -1, 1) - minS = minimum(S) + minS = S[end:end] AP ./= minS S⁻¹ = minS ./ S X₁ = rmul!(AP * Y₀ᴴ', Diagonal(S⁻¹)) From 695c2b200f43a1d45f13a96ada77e2acc72533ec Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 19 Aug 2026 15:10:20 +0200 Subject: [PATCH 3/4] Update src/pullbacks/svd.jl Co-authored-by: Lukas Devos --- src/pullbacks/svd.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pullbacks/svd.jl b/src/pullbacks/svd.jl index 7e67aca7f..3b2258d1c 100755 --- a/src/pullbacks/svd.jl +++ b/src/pullbacks/svd.jl @@ -241,7 +241,7 @@ function svd_trunc_pullback!( Y₀ᴴ = iszerotangent(ΔV₊ᴴ) ? zero(Vᴴ) : ldiv!(Diagonal(S), ΔV₊ᴴ) US = mul!(ΔAV, U, Smat) # recycle ΔAV AP = mul!(copy(A), US, Vᴴ, -1, 1) - minS = S[end:end] + @view minS = S[end:end] AP ./= minS S⁻¹ = minS ./ S X₁ = rmul!(AP * Y₀ᴴ', Diagonal(S⁻¹)) From ba187a89f58ced336b4ad13baf06581ccbda3b1a Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 19 Aug 2026 16:50:27 +0200 Subject: [PATCH 4/4] Move view right --- src/pullbacks/svd.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pullbacks/svd.jl b/src/pullbacks/svd.jl index 3b2258d1c..6c42bc26d 100755 --- a/src/pullbacks/svd.jl +++ b/src/pullbacks/svd.jl @@ -241,7 +241,7 @@ function svd_trunc_pullback!( Y₀ᴴ = iszerotangent(ΔV₊ᴴ) ? zero(Vᴴ) : ldiv!(Diagonal(S), ΔV₊ᴴ) US = mul!(ΔAV, U, Smat) # recycle ΔAV AP = mul!(copy(A), US, Vᴴ, -1, 1) - @view minS = S[end:end] + minS = @view S[end:end] AP ./= minS S⁻¹ = minS ./ S X₁ = rmul!(AP * Y₀ᴴ', Diagonal(S⁻¹))