Merge 923ad31 into bbe68cc

src/rulesets/Base/arraymath.jl

@@ -104,3 +104,14 @@ function rrule(::typeof(\), b::Real, A::AbstractArray{<:Real})
     end
     return Y, backslash_pullback
 end
+
+#####
+##### Negation (Unary -)
+#####
+
+function rrule(::typeof(-), x::AbstractArray)
+    function negation_pullback(ȳ)
+        return NO_FIELDS, InplaceableThunk(@thunk(-ȳ), ā -> _subtract!!(ā, ȳ))
+    end
+    return -x, negation_pullback
+end

src/rulesets/LinearAlgebra/utils.jl

@@ -17,6 +17,9 @@ _mulsubtrans!!(X::AbstractZero, F::AbstractZero) = X
 _mulsubtrans!!(X::AbstractZero, F::AbstractMatrix{<:Real}) = X
 _mulsubtrans!!(X::AbstractMatrix{<:Real}, F::AbstractZero) = F
 
+_subtract!!(x, y) = x - y
+_subtract!!(x::Array, y::AbstractArray) = x .-= y
+
 # I - X, overwrites X
 function _eyesubx!(X::AbstractMatrix)
     n, m = size(X)

test/rulesets/Base/arraymath.jl

@@ -63,4 +63,13 @@
         rrule_test(/, Ȳ, (A, Ā), (7.2, 2.3))
         rrule_test(\, Ȳ, (7.2, 2.3), (A, Ā))
     end
+
+
+    @testset "negation" begin
+        A = randn(4, 4)
+        Ā = randn(4, 4)
+        Ȳ = randn(4, 4)
+        rrule_test(-, Ȳ, (A, Ā))
+        rrule_test(-, Diagonal(Ȳ), (Diagonal(A), Diagonal(Ā)))
+    end
 end

test/rulesets/LinearAlgebra/utils.jl

@@ -0,0 +1,27 @@
+@testset "utils.jl" begin
+    @testset "_subtract!!" begin
+        _subtract!! = ChainRules._subtract!!
+
+        @testset "Inplace" begin
+            x = [1, 2, 3]
+            ret = _subtract!!(x, ones(3))
+            @test ret === x
+            @test ret == [0, 1, 2]
+        end
+
+        @testset "Out of place" begin
+            x = Diagonal([2, 2])
+            ret = _subtract!!(x, ones(2, 2))
+            @test ret !== x
+            @test ret == [1 -1; -1 1]
+        end
+
+        @testset "Currently out of place, but this could change" begin
+            x = Diagonal([3, 3])
+            ret = _subtract!!(x, Diagonal([1,1]))
+            @test ret !== x
+            @test ret isa Diagonal
+            @test ret == Diagonal([2, 2])
+        end
+    end
+end

test/runtests.jl

@@ -39,6 +39,7 @@ println("Testing ChainRules.jl")
             include(joinpath("rulesets", "LinearAlgebra", "structured.jl"))
             include(joinpath("rulesets", "LinearAlgebra", "factorization.jl"))
             include(joinpath("rulesets", "LinearAlgebra", "blas.jl"))
+            include(joinpath("rulesets", "LinearAlgebra", "utils.jl"))
         end
 
         print(" ")