Fix ReverseDiff bug (#278)

* Fix ReverseDiff bug

* Add tests for MWE

* Add tests with +

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "FillArrays"
 uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
-version = "1.4.0"
+version = "1.4.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -15,8 +15,9 @@ julia = "1.6"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test", "Base64", "StaticArrays"]
+test = ["Aqua", "Test", "Base64", "ReverseDiff", "StaticArrays"]

src/fillbroadcast.jl

@@ -195,20 +195,18 @@ function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::AbstractRange, b::O
     return _range_convert(AbstractVector{TT}, a)
 end
 
-# To fix AD issues with `broadcast(T, x)`
-# Avoids type inference issues with x -> T(x)
-struct Constructor{T} end
-
-function (::Constructor{T})(x) where {T}
-    return T(x)
-end
-
 for op in (:+, :-)
     @eval begin
         function broadcasted(::DefaultArrayStyle{1}, ::typeof($op), a::AbstractVector, b::ZerosVector)
             broadcast_shape(axes(a), axes(b)) == axes(a) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $b to a Vector first."))
             TT = typeof($op(zero(eltype(a)), zero(eltype(b))))
-            eltype(a) === TT ? a : broadcasted(Constructor{TT}(), a)
+            # Use `TT ∘ (+)` to fix AD issues with `broadcasted(TT, x)`
+            eltype(a) === TT ? a : broadcasted(TT ∘ (+), a)
+        end
+        function broadcasted(::DefaultArrayStyle{1}, ::typeof($op), a::ZerosVector, b::AbstractVector)
+            broadcast_shape(axes(a), axes(b)) == axes(b) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $a to a Vector first."))
+            TT = typeof($op(zero(eltype(a)), zero(eltype(b))))
+            $op === (+) && eltype(b) === TT ? b : broadcasted(TT ∘ ($op), b)
         end
 
         broadcasted(::DefaultArrayStyle{1}, ::typeof($op), a::AbstractFillVector, b::ZerosVector) =
@@ -219,18 +217,6 @@ for op in (:+, :-)
     end
 end
 
-function broadcasted(::DefaultArrayStyle{1}, ::typeof(+), a::ZerosVector, b::AbstractVector)
-    broadcast_shape(axes(a), axes(b)) == axes(b) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $a to a Vector first."))
-    TT = typeof(zero(eltype(a)) + zero(eltype(b)))
-    eltype(b) === TT ? b : broadcasted(Constructor{TT}(), b)
-end
-
-function broadcasted(::DefaultArrayStyle{1}, ::typeof(-), a::ZerosVector, b::AbstractVector)
-    broadcast_shape(axes(a), axes(b)) == axes(b) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $a to a Vector first."))
-    TT = typeof(zero(eltype(a)) - zero(eltype(b)))
-    broadcasted(TT ∘ (-), b)
-end
-
 # Need to prevent array-valued fills from broadcasting over entry
 _broadcast_getindex_value(a::AbstractFill{<:Number}) = getindex_value(a)
 _broadcast_getindex_value(a::AbstractFill) = Ref(getindex_value(a))

test/runtests.jl

@@ -1,4 +1,4 @@
-using FillArrays, LinearAlgebra, SparseArrays, StaticArrays, Random, Base64, Test, Statistics
+using FillArrays, LinearAlgebra, SparseArrays, StaticArrays, ReverseDiff, Random, Base64, Test, Statistics
 import FillArrays: AbstractFill, RectDiagonal, SquareEye
 
 using Aqua
@@ -2071,3 +2071,18 @@ end
         end
     end
 end
+
+@testset "ReverseDiff with Zeros" begin
+    # MWE in https://github.com/JuliaArrays/FillArrays.jl/issues/252
+    @test ReverseDiff.gradient(x -> sum(abs2.((Zeros(5) .- zeros(5)) ./ x)), rand(5)) == zeros(5)
+    @test ReverseDiff.gradient(x -> sum(abs2.((zeros(5) .- Zeros(5)) ./ x)), rand(5)) == zeros(5)
+    # MWE in https://github.com/JuliaArrays/FillArrays.jl/pull/278
+    @test ReverseDiff.gradient(x -> sum(abs2.((Zeros{eltype(x)}(5) .- zeros(5)) ./ x)), rand(5)) == zeros(5)
+    @test ReverseDiff.gradient(x -> sum(abs2.((zeros(5) .- Zeros{eltype(x)}(5)) ./ x)), rand(5)) == zeros(5)
+
+    # Corresponding tests with +
+    @test ReverseDiff.gradient(x -> sum(abs2.((Zeros(5) .+ zeros(5)) ./ x)), rand(5)) == zeros(5)
+    @test ReverseDiff.gradient(x -> sum(abs2.((zeros(5) .+ Zeros(5)) ./ x)), rand(5)) == zeros(5)
+    @test ReverseDiff.gradient(x -> sum(abs2.((Zeros{eltype(x)}(5) .+ zeros(5)) ./ x)), rand(5)) == zeros(5)
+    @test ReverseDiff.gradient(x -> sum(abs2.((zeros(5) .+ Zeros{eltype(x)}(5)) ./ x)), rand(5)) == zeros(5)
+end