RFC: rules related to 3-arg * (#412)

* gradients for 3-arg *, untested

* fixup, add tests

* version bound defn

* fix depwarns & indents

* update InplaceableThunk

* fix complaints

* simplify + update

* tweak

* version

* tweak

Co-authored-by: Michael Abbott <me@escbook>

Author: mcabbott

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.18.1"
+version = "1.19"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/arraymath.jl

@@ -128,6 +128,77 @@ function rrule(
     return A * B, times_pullback
 end
 
+#####
+##### fused 3-argument *
+#####
+
+if VERSION > v"1.7.0-"
+
+    @eval using LinearAlgebra: mat_mat_scalar, mat_vec_scalar, StridedMaybeAdjOrTransMat
+
+    function rrule(
+        ::typeof(mat_mat_scalar),
+        A::StridedMaybeAdjOrTransMat{<:CommutativeMulNumber},
+        B::StridedMaybeAdjOrTransMat{<:CommutativeMulNumber},
+        γ::CommutativeMulNumber
+    )
+        project_A = ProjectTo(A)
+        project_B = ProjectTo(B)
+        project_γ = ProjectTo(γ)
+        C = mat_mat_scalar(A, B, γ)
+        function mat_mat_scalar_back(Ȳraw)
+            Ȳ = unthunk(Ȳraw)
+            Athunk = InplaceableThunk(
+                dA -> mul!(dA, Ȳ, B', conj(γ), true),
+                @thunk(project_A(mat_mat_scalar(Ȳ, B', conj(γ)))),
+            )
+            Bthunk = InplaceableThunk(
+                dB -> mul!(dB, A', Ȳ, conj(γ), true),
+                @thunk(project_B(mat_mat_scalar(A', Ȳ, conj(γ)))),
+            )
+            γthunk = @thunk if iszero(γ)
+                # Could save A*B on the forward pass, but it's messy.
+                # This ought to be rare, should guarantee the same type:
+                project_γ(dot(mat_mat_scalar(A, B, oneunit(γ)), Ȳ) / one(γ))
+            else
+                project_γ(dot(C, Ȳ) / conj(γ))
+            end
+            return (NoTangent(), Athunk, Bthunk, γthunk)
+        end
+        return C, mat_mat_scalar_back
+    end
+
+    function rrule(
+        ::typeof(mat_vec_scalar),
+        A::StridedMaybeAdjOrTransMat{<:CommutativeMulNumber},
+        b::StridedVector{<:CommutativeMulNumber},
+        γ::CommutativeMulNumber
+    )
+        project_A = ProjectTo(A)
+        project_b = ProjectTo(b)
+        project_γ = ProjectTo(γ)
+        y = mat_vec_scalar(A, b, γ)
+        function mat_vec_scalar_back(dy_raw)
+            dy = unthunk(dy_raw)
+            Athunk = InplaceableThunk(
+                dA -> mul!(dA, dy, b', conj(γ), true),
+                @thunk(project_A(*(dy, b', conj(γ)))),
+            )
+            Bthunk = InplaceableThunk(
+                db -> mul!(db, A', dy, conj(γ), true),
+                @thunk(project_b(*(A', dy, conj(γ)))),
+            )
+            γthunk = @thunk if iszero(γ)
+                project_γ(dot(mat_vec_scalar(A, b, oneunit(γ)), dy))
+            else
+                project_γ(dot(y, dy) / conj(γ))
+            end
+            return (NoTangent(), Athunk, Bthunk, γthunk)
+        end
+        return y, mat_vec_scalar_back
+    end
+
+end # VERSION
 
 #####
 ##### `muladd`

test/rulesets/Base/arraymath.jl

@@ -111,6 +111,28 @@
         end
     end
 
+    if VERSION > v"1.7.0-"
+        @eval using LinearAlgebra: mat_mat_scalar, mat_vec_scalar, StridedMaybeAdjOrTransMat
+
+        @testset "3-arg *, $T" for T in [Float64, ComplexF64]
+
+            test_rrule(mat_mat_scalar, rand(T,4,4), rand(T,4,4), rand(T))
+            test_rrule(mat_mat_scalar, rand(T,4,4), rand(T,4,4), 0.0)
+            test_rrule(mat_mat_scalar, rand(T,4,4)' ⊢ rand(T,4,4), rand(T,4,4), rand(T))
+
+            test_rrule(mat_vec_scalar, rand(T,4,4), rand(T,4), rand(T))
+            test_rrule(mat_vec_scalar, rand(T,4,4), rand(T,4), 0.0)
+
+            T == ComplexF64 && continue
+            # Test with γ of a wider type
+            A, B, b, γ = rand(3,3), rand(3,3), rand(3), rand()
+            dZ, dz = rand(3,3), rand(3)
+            unthunk(rrule(mat_mat_scalar, A, B, γ + 0im)[2](dZ)[4]) ≈ unthunk(rrule(mat_mat_scalar, A, B, γ)[2](dZ)[4])
+            unthunk(rrule(mat_mat_scalar, A, B, 0 + 0im)[2](dZ)[4]) ≈ unthunk(rrule(mat_mat_scalar, A, B, 0)[2](dZ)[4])
+            unthunk(rrule(mat_vec_scalar, A, b, γ + 0im)[2](dz)[4]) ≈ unthunk(rrule(mat_vec_scalar, A, b, γ)[2](dz)[4])
+        end
+    end # VERSION
+
     @testset "$f" for f in (/, \)
         @testset "Matrix" begin
             for n in 3:5, m in 3:5