use add!! and rename _mulsubtrans!!

JuliaDiff · Oct 5, 2020 · 5c96ffe · 5c96ffe
1 parent 2be1af2
commit 5c96ffe
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 33 deletions.
diff --git a/Project.toml b/Project.toml
@@ -11,7 +11,7 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-ChainRulesCore = "0.9.7"
+ChainRulesCore = "0.9.12"
 ChainRulesTestUtils = "0.4.2, 0.5"
 Compat = "3"
 FiniteDifferences = "0.10"

diff --git a/src/rulesets/Base/arraymath.jl b/src/rulesets/Base/arraymath.jl
@@ -74,8 +74,8 @@ function rrule(::typeof(\), A::AbstractVecOrMat{<:Real}, B::AbstractVecOrMat{<:R
         ∂A = @thunk begin
             B̄ = A' \ Ȳ
             Ā = -B̄ * Y'
-            _add!(Ā, (B - A * Y) * B̄' / A')
-            _add!(Ā, A' \ Y * (Ȳ' - B̄'A))
+            Ā = add!!(Ā, (B - A * Y) * B̄' / A')
+            Ā = add!!(Ā, A' \ Y * (Ȳ' - B̄'A))
             Ā
         end
         ∂B = @thunk A' \ Ȳ

diff --git a/src/rulesets/LinearAlgebra/dense.jl b/src/rulesets/LinearAlgebra/dense.jl
@@ -155,12 +155,12 @@ function frule((_, ΔA), ::typeof(pinv), A::AbstractMatrix{T}; kwargs...) where
     # contract over the largest dimension
     if m ≤ n
         ∂Y = -Y * (ΔA * Y)
-        _add!(∂Y, (ΔA' - Y * (A * ΔA')) * (Y' * Y)) # (I - Y A) ΔA' Y' Y
-        _add!(∂Y, Y * (Y' * ΔA') * (I - A * Y)) # Y Y' ΔA' (I - A Y)
+        ∂Y = add!!(∂Y, (ΔA' - Y * (A * ΔA')) * (Y' * Y))  # (I - Y A) ΔA' Y' Y
+        ∂Y = add!!(∂Y, Y * (Y' * ΔA') * (I - A * Y))  # Y Y' ΔA' (I - A Y)
     else
         ∂Y = -(Y * ΔA) * Y
-        _add!(∂Y, (I - Y * A) * (ΔA' * Y') * Y) # (I - Y A) ΔA' Y' Y
-        _add!(∂Y, (Y * Y') * (ΔA' - (ΔA' * A) * Y)) # Y Y' ΔA' (I - A Y)
+        ∂Y = add!!(∂Y, (I - Y * A) * (ΔA' * Y') * Y)  # (I - Y A) ΔA' Y' Y
+        ∂Y = add!!(∂Y, (Y * Y') * (ΔA' - (ΔA' * A) * Y))  # Y Y' ΔA' (I - A Y)
     end
     return Y, ∂Y
 end
@@ -199,12 +199,12 @@ function rrule(::typeof(pinv), A::AbstractMatrix{T}; kwargs...) where {T}
         # contract over the largest dimension
         if m ≤ n
             ∂A = (Y' * -ΔY) * Y'
-            _add!(∂A, (Y' * Y) * (ΔY' - (ΔY' * Y) * A)) # Y' Y ΔY' (I - Y A)
-            _add!(∂A, (I - A * Y) * (ΔY' * Y) * Y') # (I - A Y) ΔY' Y Y'
+            ∂A = add!!(∂A, (Y' * Y) * (ΔY' - (ΔY' * Y) * A)) # Y' Y ΔY' (I - Y A)
+            ∂A = add!!(∂A, (I - A * Y) * (ΔY' * Y) * Y') # (I - A Y) ΔY' Y Y'
         elseif m > n
             ∂A = Y' * (-ΔY * Y')
-            _add!(∂A, Y' * (Y * ΔY') * (I - Y * A)) # Y' Y ΔY' (I - Y A)
-            _add!(∂A, (ΔY' - A * (Y * ΔY')) * (Y * Y')) # (I - A Y) ΔY' Y Y'
+            ∂A = add!!(∂A, Y' * (Y * ΔY') * (I - Y * A)) # Y' Y ΔY' (I - Y A)
+            ∂A = add!!(∂A, (ΔY' - A * (Y * ΔY')) * (Y * Y')) # (I - A Y) ΔY' Y Y'
         end
         return (NO_FIELDS, ∂A)
     end

diff --git a/src/rulesets/LinearAlgebra/factorization.jl b/src/rulesets/LinearAlgebra/factorization.jl
@@ -49,18 +49,18 @@ function svd_rev(USV::SVD, Ū, s̄, V̄)
     # place functions here are significantly faster than their out-of-place, naively
     # implemented counterparts, and allocate no additional memory.
     Ut = U'
-    FUᵀŪ = _mulsubtrans!(Ut*Ū, F)  # F .* (UᵀŪ - ŪᵀU)
-    FVᵀV̄ = _mulsubtrans!(Vt*V̄, F)  # F .* (VᵀV̄ - V̄ᵀV)
-    ImUUᵀ = _eyesubx!(U*Ut)        # I - UUᵀ
-    ImVVᵀ = _eyesubx!(V*Vt)        # I - VVᵀ
+    FUᵀŪ = _mulsubtrans!!(Ut*Ū, F)  # F .* (UᵀŪ - ŪᵀU)
+    FVᵀV̄ = _mulsubtrans!!(Vt*V̄, F)  # F .* (VᵀV̄ - V̄ᵀV)
+    ImUUᵀ = _eyesubx!(U*Ut)  # I - UUᵀ
+    ImVVᵀ = _eyesubx!(V*Vt)  # I - VVᵀ
 
     S = Diagonal(s)
     S̄ = s̄ isa AbstractZero ? s̄ : Diagonal(s̄)
 
     # TODO: consider using MuladdMacro here
-    Ā = _add!(U * FUᵀŪ * S, ImUUᵀ * (Ū / S)) * Vt
-    Ā = _add!(Ā, U * S̄ * Vt)
-    Ā = _add!(Ā, U * _add!(S * FVᵀV̄ * Vt, (S \ V̄') * ImVVᵀ))
+    Ā = add!!(U * FUᵀŪ * S, ImUUᵀ * (Ū / S)) * Vt
+    Ā = add!!(Ā, U * S̄ * Vt)
+    Ā = add!!(Ā, U * add!!(S * FVᵀV̄ * Vt, (S \ V̄') * ImVVᵀ))
 
     return Ā
 end

diff --git a/src/rulesets/LinearAlgebra/utils.jl b/src/rulesets/LinearAlgebra/utils.jl
@@ -2,7 +2,7 @@
 # to any particular rule definition
 
 # F .* (X - X'), overwrites X if possible
-function _mulsubtrans!(X::AbstractMatrix{T}, F::AbstractMatrix{T}) where T<:Real
+function _mulsubtrans!!(X::AbstractMatrix{T}, F::AbstractMatrix{T}) where T<:Real
     k = size(X, 1)
     @inbounds for j = 1:k, i = 1:j  # Iterate the upper triangle
         if i == j
@@ -13,9 +13,9 @@ function _mulsubtrans!(X::AbstractMatrix{T}, F::AbstractMatrix{T}) where T<:Real
     end
     return X
 end
-_mulsubtrans!(X::AbstractZero, F::AbstractZero) = X
-_mulsubtrans!(X::AbstractZero, F::AbstractMatrix{<:Real}) = X
-_mulsubtrans!(X::AbstractMatrix{<:Real}, F::AbstractZero) = F
+_mulsubtrans!!(X::AbstractZero, F::AbstractZero) = X
+_mulsubtrans!!(X::AbstractZero, F::AbstractMatrix{<:Real}) = X
+_mulsubtrans!!(X::AbstractMatrix{<:Real}, F::AbstractZero) = F
 
 # I - X, overwrites X
 function _eyesubx!(X::AbstractMatrix)
@@ -26,13 +26,4 @@ function _eyesubx!(X::AbstractMatrix)
     return X
 end
 
-# X + Y, overwrites X if possible
-function _add!(X::AbstractVecOrMat, Y::AbstractVecOrMat)
-    @inbounds for i = eachindex(X, Y)
-        X[i] += Y[i]
-    end
-    return X
-end
-_add!(X, Y) = X + Y  # handles all `AbstractZero` overloads
-
 _extract_imag(x) = complex(0, imag(x))
diff --git a/test/rulesets/LinearAlgebra/factorization.jl b/test/rulesets/LinearAlgebra/factorization.jl
@@ -47,9 +47,8 @@ using ChainRules: level2partition, level3partition, chol_blocked_rev, chol_unblo
         @testset "Helper functions" begin
             X = randn(10, 10)
             Y = randn(10, 10)
-            @test ChainRules._mulsubtrans!(copy(X), Y) ≈ Y .* (X - X')
+            @test ChainRules._mulsubtrans!!(copy(X), Y) ≈ Y .* (X - X')
             @test ChainRules._eyesubx!(copy(X)) ≈ I - X
-            @test ChainRules._add!(copy(X), Y) ≈ X + Y
         end
     end
     @testset "cholesky" begin