Broadcast the propagation_expr for vector mode AD (#93)

* Broadcast the `propagation_expr` for vector mode AD

* Use the `muladd` macro to optimize `propagation_expr`

* New release

* Fix propagation_expr

* Explicit broadcast

* Revert "Explicit broadcast"

This reverts commit b7d1da2.

Project.toml

@@ -1,9 +1,13 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.5.1"
+version = "0.5.2"
+
+[deps]
+MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 
 [compat]
 julia = "^1.0"
+MuladdMacro = "0.2.1"
 
 [extras]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/rule_definition_tools.jl

@@ -1,4 +1,5 @@
 # These are some macros (and supporting functions) to make it easier to define rules.
+using MuladdMacro: @muladd
 
 """
     @scalar_rule(f(x₁, x₂, ...),
@@ -208,9 +209,25 @@ end
 function propagation_expr(Δs, ∂s)
     # This is basically Δs ⋅ ∂s
     ∂s = map(esc, ∂s)
+    n∂s = length(∂s)
 
-    ∂_mul_Δs = ntuple(i->:($(∂s[i]) * $(Δs[i])), length(∂s))
-    return :(+($(∂_mul_Δs...)))
+    # Due to bugs in Julia 1.0, we can't use `.+`  or `.*` inside expression
+    # literals.
+    ∂_mul_Δs = ntuple(i->:($(∂s[i]) * $(Δs[i])), n∂s)
+
+    # Avoiding the extra `+` operation, it is potentially expensive for vector
+    # mode AD.
+    sumed_∂_mul_Δs = if n∂s > 1
+        # we use `@.` to broadcast `*` and `+`
+        :(@. +($(∂_mul_Δs...)))
+    else
+        # Note: we don't want to do broadcasting with only 1 multiply (no `+`),
+        # because some arrays overload multiply with scalar. Avoiding
+        # broadcasting saves compilation time.
+        ∂_mul_Δs[1]
+    end
+
+    return :(@muladd $sumed_∂_mul_Δs)
 end
 
 """