Merge pull request #226 from JuliaDiff/ox/accum

Introduce add!!

Author: oxinabox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.9.11"
+version = "0.9.12"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

docs/Manifest.toml

@@ -10,10 +10,10 @@ uuid = "b99e7846-7c00-51b0-8f62-c81ae34c0232"
 version = "0.5.10"
 
 [[ChainRulesCore]]
-deps = ["MuladdMacro"]
+deps = ["LinearAlgebra", "MuladdMacro"]
 path = ".."
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.9.5"
+version = "0.9.12"
 
 [[Dates]]
 deps = ["Printf"]
@@ -25,9 +25,9 @@ uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 
 [[DocStringExtensions]]
 deps = ["LibGit2", "Markdown", "Pkg", "Test"]
-git-tree-sha1 = "c5714d9bcdba66389612dc4c47ed827c64112997"
+git-tree-sha1 = "50ddf44c53698f5e784bbebb3f4b21c5807401b1"
 uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
-version = "0.8.2"
+version = "0.8.3"
 
 [[Documenter]]
 deps = ["Base64", "Dates", "DocStringExtensions", "InteractiveUtils", "JSON", "LibGit2", "Logging", "Markdown", "REPL", "Test", "Unicode"]
@@ -50,9 +50,9 @@ uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 
 [[JSON]]
 deps = ["Dates", "Mmap", "Parsers", "Unicode"]
-git-tree-sha1 = "b34d7cef7b337321e97d22242c3c2b91f476748e"
+git-tree-sha1 = "81690084b6198a2e1da36fcfda16eeca9f9f24e4"
 uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
-version = "0.21.0"
+version = "0.21.1"
 
 [[LibGit2]]
 deps = ["Printf"]
@@ -61,6 +61,10 @@ uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
 [[Libdl]]
 uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 
+[[LinearAlgebra]]
+deps = ["Libdl"]
+uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+
 [[Logging]]
 uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
 
@@ -91,7 +95,7 @@ deps = ["Unicode"]
 uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
 [[REPL]]
-deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
+deps = ["InteractiveUtils", "Markdown", "Sockets"]
 uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 
 [[Random]]

docs/src/api.md

@@ -27,6 +27,13 @@ Pages = [
 Private = false
 ```
 
+## Accumulation
+```@autodocs
+Modules = [ChainRulesCore]
+Pages = ["accumulation.jl"]
+Private = false
+```
+
 ## Ruleset Loading
 ```@autodocs
 Modules = [ChainRulesCore]

src/ChainRulesCore.jl

@@ -7,6 +7,7 @@ export on_new_rule, refresh_rules  # generation tools
 export frule, rrule  # core function
 export @non_differentiable, @scalar_rule, @thunk  # definition helper macros
 export canonicalize, extern, unthunk  # differential operations
+export add!!  # gradient accumulation operations
 # differentials
 export Composite, DoesNotExist, InplaceableThunk, One, Thunk, Zero, AbstractZero, AbstractThunk
 export NO_FIELDS
@@ -21,6 +22,7 @@ include("differentials/thunks.jl")
 include("differentials/composite.jl")
 
 include("differential_arithmetic.jl")
+include("accumulation.jl")
 
 include("rules.jl")
 include("rule_definition_tools.jl")

src/accumulation.jl

@@ -0,0 +1,15 @@
+"""
+    add!!(x, y)
+
+Returns `x+y`, potentially mutating `x` in-place to hold this value.
+This avoids allocations when `x` can be mutated in this way.
+
+See also: [`InplaceableThunk`](@ref).
+"""
+add!!(x, y) = x + y
+
+add!!(x, t::InplaceableThunk) = t.add!(x)
+
+function add!!(x::Array{<:Any, N}, y::AbstractArray{<:Any, N}) where N
+    return x .+= y
+end

test/accumulation.jl

@@ -0,0 +1,62 @@
+@testset "accumulation.jl" begin
+    @testset "scalar" begin
+        @test 16 == add!!(12, 4)
+    end
+
+    @testset "Differentials" begin
+        @test 16 == add!!(12, @thunk(2*2))
+        @test 16 == add!!(16, Zero())
+
+        @test 16 == add!!(16, DoesNotExist())  # Should this be an error?
+    end
+
+    @testset "Array" begin
+        @testset "Happy Path" begin
+            @testset "RHS Array" begin
+                A = [1.0 2.0; 3.0 4.0]
+                result = -1.0*ones(2,2)
+                ret = add!!(result, A)
+                @test ret === result  # must be same object
+                @test result == [0.0 1.0; 2.0 3.0]
+            end
+
+            @testset "RHS StaticArray" begin
+                A = @SMatrix[1.0 2.0; 3.0 4.0]
+                result = -1.0*ones(2,2)
+                ret = add!!(result, A)
+                @test ret === result  # must be same object
+                @test result == [0.0 1.0; 2.0 3.0]
+            end
+
+            @testset "RHS Diagonal" begin
+                A = Diagonal([1.0, 2.0])
+                result = -1.0*ones(2,2)
+                ret = add!!(result, A)
+                @test ret === result  # must be same object
+                @test result == [0.0 -1.0; -1.0 1.0]
+            end
+        end
+
+        @testset "Unhappy Path" begin
+            # wrong length
+            @test_throws DimensionMismatch add!!(ones(4,4), ones(2,2))
+            # wrong shape
+            @test_throws DimensionMismatch add!!(ones(4,4), ones(16))
+            # wrong type (adding scalar to array)
+            @test_throws MethodError add!!(ones(4), 21.0)
+        end
+    end
+
+    @testset "InplaceableThunk" begin
+        A=[1.0 2.0; 3.0 4.0]
+        ithunk = InplaceableThunk(
+            @thunk(A*B),
+            x -> x.+=A
+        )
+
+        result = -1.0*ones(2,2)
+        ret = add!!(result, ithunk)
+        @test ret === result  # must be same object
+        @test result == [0.0 1.0; 2.0 3.0]
+    end
+end

test/rules.jl

@@ -1,6 +1,5 @@
 #######
 # Demo setup
-using StaticArrays: @SVector
 
 cool(x) = x + 1
 cool(x, y) = x + y + 1

test/runtests.jl

@@ -3,6 +3,7 @@ using BenchmarkTools
 using ChainRulesCore
 using LinearAlgebra: Diagonal, dot
 using FiniteDifferences
+using StaticArrays
 using Test
 
 @testset "ChainRulesCore" begin
@@ -13,6 +14,8 @@ using Test
         include("differentials/composite.jl")
     end
 
+    include("accumulation.jl")
+
     include("ruleset_loading.jl")
     include("rules.jl")
     include("rule_definition_tools.jl")