rrule for many-arg * (#547)

* rrule for many-arg *

* tweak

* unrelated, remove a dot from muladd

* bump version

* Revert "unrelated, remove a dot from muladd"

This reverts commit eb5b761.

* fixup

* skip inference on 1.0

Author: mcabbott

Project.toml

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

src/rulesets/Base/fastmath_able.jl

@@ -231,19 +231,46 @@ let
             # Optimized version of `Δx .* y .+ x .* Δy`. Also, it is potentially more
             # accurate on machines with FMA instructions, since there are only two
             # rounding operations, one in `muladd/fma` and the other in `*`.
-            ∂xy = muladd.(Δx, y, x .* Δy)
+            ∂xy = muladd(Δx, y, x * Δy)
             return x * y, ∂xy
         end
+        frule((_, Δx), ::typeof(*), x::Number) = x, Δx
 
         function rrule(::typeof(*), x::Number, y::Number)
-            project_x = ProjectTo(x)
-            project_y = ProjectTo(y)
-            function times_pullback(Ω̇)
+            function times_pullback2(Ω̇)
                 ΔΩ = unthunk(Ω̇)
-                return (NoTangent(),  project_x(ΔΩ * y'), project_y(x' * ΔΩ))
+                return (NoTangent(), ProjectTo(x)(ΔΩ * y'), ProjectTo(y)(x' * ΔΩ))
+            end
+            return x * y, times_pullback2
+        end
+        # While 3-arg * calls 2-arg *, this is currently slow in Zygote:
+        # https://github.com/JuliaDiff/ChainRules.jl/issues/544
+        function rrule(::typeof(*), x::Number, y::Number, z::Number)
+            function times_pullback3(Ω̇)
+                ΔΩ = unthunk(Ω̇)
+                return (
+                    NoTangent(), 
+                    ProjectTo(x)(ΔΩ * y' * z'),
+                    ProjectTo(y)(x' * ΔΩ * z'),
+                    ProjectTo(z)(x' * y' * ΔΩ),
+                )
+            end
+            return x * y * z, times_pullback3
+        end
+        # Instead of this recursive rule for N args, you could write the generic case
+        # directly, by nesting ntuples, but this didn't infer well:
+        # https://github.com/JuliaDiff/ChainRules.jl/pull/547/commits/3558951c9f1b3c70e7135fd61d29cc3b96a68dea
+        function rrule(::typeof(*), x::Number, y::Number, z::Number, more::Number...)
+            Ω3, back3 = rrule(*, x, y, z)
+            Ω4, back4 = rrule(*, Ω3, more...)
+            function times_pullback4(Ω̇)
+                Δ4 = back4(unthunk(Ω̇))  # (0, ΔΩ3, Δmore...)
+                Δ3 = back3(Δ4[2])       # (0, Δx, Δy, Δz)
+                return (Δ3..., Δ4[3:end]...)
             end
-            return x * y, times_pullback
+            return Ω4, times_pullback4
         end
+        rrule(::typeof(*), x::Number) = rrule(identity, x)
     end  # fastable_ast
 
     # Rewrite everything to use fast_math functions, including the type-constraints

test/rulesets/Base/base.jl

@@ -104,6 +104,15 @@
             test_frule(*, x, y)
             test_rrule(*, x, y)
         end
+        @testset "*($x, $y, ...)" for x in test_points, y in test_points
+            # This promotion is only for FiniteDifferences, the rules allow mixtures:
+            x, y = Base.promote(x, y)
+
+            # Inference fails on 1.0, passes on 1.6
+            test_rrule(*, x, y, x+y; check_inferred=VERSION>v"1.5")
+            test_rrule(*, x, y, 17x, 23y; check_inferred=VERSION>v"1.5")
+            test_rrule(*, x, y, 7x, 3y, x+y+pi; check_inferred=VERSION>v"1.5")
+        end
     end
 
     @testset "ldexp" begin

test/rulesets/Base/fastmath_able.jl

@@ -120,6 +120,7 @@ const FASTABLE_AST = quote
             test_scalar(+, x)
             test_scalar(-, x)
             test_scalar(atan, x)
+            test_scalar(*, x)
         end
     end
 
@@ -132,7 +133,7 @@ const FASTABLE_AST = quote
             test_rrule(f, (rand(0:10) + .6rand() + .2) * base, base)
         end
 
-        @testset "$f(x::$T, y::$T)" for f in (/, +, -, hypot), T in (Float64, ComplexF64)
+        @testset "$f(x::$T, y::$T)" for f in (/, +, -, *, hypot), T in (Float64, ComplexF64)
             test_frule(f, 10rand(T), rand(T))
             test_rrule(f, 10rand(T), rand(T))
         end