add subtract rule

Author: nmheim

src/rulesets/Base/arraymath.jl

@@ -434,6 +434,16 @@ function rrule(::typeof(-), x::AbstractArray)
     return -x, negation_pullback
 end
 
+#####
+##### Subtraction
+#####
+
+frule((_, Δx, Δy), ::typeof(-), x::AbstractArray, y::AbstractArray) = x-y, Δx-Δy
+
+function rrule(::typeof(-), x::AbstractArray, y::AbstractArray)
+    subtract_pullback(dy) = (NoTangent(), dy, -dy)
+    return x-y, subtract_pullback
+end
 
 #####
 ##### Addition (Multiarg `+`)

test/rulesets/Base/arraymath.jl

@@ -217,4 +217,12 @@
         @gpu test_rrule(+, randn(4, 4), randn(4, 4), randn(4, 4))
         @gpu test_rrule(+, randn(3), randn(3,1), randn(3,1,1))
     end
+
+    @testset "subtraction" begin
+        # fwd
+        @gpu test_frule(-, randn(2), randn(2))
+        # rev
+        @gpu test_rrule(-, randn(4, 4), randn(4, 4))
+        @gpu test_rrule(-, randn(3), randn(3,1))
+    end
 end