add subtract rule

JuliaDiff · Feb 9, 2024 · 2c3e48d · 2c3e48d
1 parent b3f9f9b
commit 2c3e48d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/src/rulesets/Base/arraymath.jl b/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 `+`)

diff --git a/test/rulesets/Base/arraymath.jl b/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