Merge pull request #419 from JuliaDiff/ox/muladd

Ox/muladd

Author: oxinabox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.7.68"
+version = "0.7.69"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -14,8 +14,8 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 ChainRulesCore = "0.9.44"
 ChainRulesTestUtils = "0.6.8"
-Compat = "3"
-FiniteDifferences = "0.11, 0.12"
+Compat = "3.30"
+FiniteDifferences = "0.12.8"
 Reexport = "0.2, 1"
 Requires = "0.5.2, 1"
 julia = "1"

src/rulesets/Base/arraymath.jl

@@ -95,6 +95,90 @@ function rrule(
 end
 
 
+#####
+##### `muladd`
+#####
+
+function rrule(
+        ::typeof(muladd),
+        A::AbstractMatrix{<:CommutativeMulNumber},
+        B::AbstractVecOrMat{<:CommutativeMulNumber},
+        z::Union{CommutativeMulNumber, AbstractVecOrMat{<:CommutativeMulNumber}},
+    )
+    # The useful case, mul! fused with +
+    function muladd_pullback_1(Ȳ)
+        matmul = (
+            InplaceableThunk(
+                @thunk(Ȳ * B'),
+                dA -> mul!(dA, Ȳ, B', true, true)
+            ),
+            InplaceableThunk(
+                @thunk(A' * Ȳ),
+                dB -> mul!(dB, A', Ȳ, true, true)
+            )
+        )
+        addon = if z isa Bool
+            DoesNotExist()
+        elseif z isa Number
+            @thunk(sum(Ȳ))
+        else
+            InplaceableThunk(
+                @thunk(sum!(similar(z, eltype(Ȳ)), Ȳ)),
+                dz -> sum!(dz, Ȳ; init=false)
+            )
+        end
+        (NO_FIELDS, matmul..., addon)
+    end
+    return muladd(A, B, z), muladd_pullback_1
+end
+
+function rrule(
+        ::typeof(muladd),
+        ut::LinearAlgebra.AdjOrTransAbsVec{<:CommutativeMulNumber},
+        v::AbstractVector{<:CommutativeMulNumber},
+        z::CommutativeMulNumber,
+    )
+    # This case is dot(u,v)+z, but would also match signature above.
+    function muladd_pullback_2(dy)
+        ut_thunk = InplaceableThunk(
+            @thunk(v' .* dy),
+            dut -> dut .+= v' .* dy
+        )
+        v_thunk = InplaceableThunk(
+            @thunk(ut' .* dy),
+            dv -> dv .+= ut' .* dy
+        )
+        (NO_FIELDS, ut_thunk, v_thunk, z isa Bool ? DoesNotExist() : dy)
+    end
+    return muladd(ut, v, z), muladd_pullback_2
+end
+
+function rrule(
+        ::typeof(muladd),
+        u::AbstractVector{<:CommutativeMulNumber},
+        vt::LinearAlgebra.AdjOrTransAbsVec{<:CommutativeMulNumber},
+        z::Union{CommutativeMulNumber, AbstractVecOrMat{<:CommutativeMulNumber}},
+    )
+    # Outer product, just broadcasting
+    function muladd_pullback_3(Ȳ)
+        proj = (
+            @thunk(vec(sum(Ȳ .* conj.(vt), dims=2))),
+            @thunk(vec(sum(u .* conj.(Ȳ), dims=1))'),
+        )
+        addon = if z isa Bool
+            DoesNotExist()
+        elseif z isa Number
+            @thunk(sum(Ȳ))
+        else
+            InplaceableThunk(
+                @thunk(sum!(similar(z, eltype(Ȳ)), Ȳ)),
+                dz -> sum!(dz, Ȳ; init=false)
+            )
+        end
+        (NO_FIELDS, proj..., addon)
+    end
+    return muladd(u, vt, z), muladd_pullback_3
+end
 
 #####
 ##### `/`

test/rulesets/Base/arraymath.jl

@@ -62,6 +62,47 @@
         end
     end
 
+    @testset "muladd: $T" for T in (Float64, ComplexF64)
+        @testset "add $(typeof(z))" for z in [rand(T), rand(T, 3), rand(T, 3, 3), false]
+            @testset "matrix * matrix" begin
+                A = rand(T, 3, 3)
+                B = rand(T, 3, 3)
+                test_rrule(muladd, A, B, z)
+                test_rrule(muladd, A', B, z)
+                test_rrule(muladd, A , B', z)
+
+                C = rand(T, 3, 5)
+                D = rand(T, 5, 3)
+                test_rrule(muladd, C, D, z)
+            end
+            if ndims(z) <= 1
+                @testset "matrix * vector" begin
+                    A, B = rand(T, 3, 3), rand(T, 3)
+                    test_rrule(muladd, A, B, z)
+                    test_rrule(muladd, A, B ⊢ rand(T, 3,1), z)
+                end
+                @testset "adjoint * matrix" begin
+                    At, B = rand(T, 3)', rand(T, 3, 3)
+                    test_rrule(muladd, At, B, z')
+                    test_rrule(muladd, At ⊢ rand(T,1,3), B, z')
+                end
+            end
+            if ndims(z) == 0
+                @testset "adjoint * vector" begin # like dot
+                    A, B = rand(T, 3)', rand(T, 3)
+                    test_rrule(muladd, A, B, z)
+                    test_rrule(muladd, A ⊢ rand(T,1,3), B, z')
+                end
+            end
+            if ndims(z) == 2 # other dims lead to e.g. muladd(ones(4), ones(1,4), 1)
+                @testset "vector * adjoint" begin # outer product
+                    A, B = rand(T, 3), rand(T, 3)'
+                    test_rrule(muladd, A, B, z)
+                    test_rrule(muladd, A, B ⊢ rand(T,1,3), z)
+                end
+            end
+        end
+    end
 
     @testset "$f" for f in (/, \)
         @testset "Matrix" begin
@@ -89,13 +130,13 @@
             end
         end
     end
+
     @testset "/ and \\ Scalar-AbstractArray" begin
         A = randn(3, 4, 5)
         test_rrule(/, A, 7.2)
         test_rrule(\, 7.2, A)
     end
 
-
     @testset "negation" begin
         A = randn(4, 4)
         Ā = randn(4, 4)