make add!! responsible for deciding if to use InplaceThunk.add!

Author: oxinabox

Project.toml

@@ -1,10 +1,11 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.9.16"
+version = "0.9.17"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 BenchmarkTools = "0.5"

src/ChainRulesCore.jl

@@ -1,6 +1,7 @@
 module ChainRulesCore
 using Base.Broadcast: broadcasted, Broadcasted, broadcastable, materialize, materialize!
 using LinearAlgebra: LinearAlgebra
+using SparseArrays: SparseVector, SparseMatrixCSC
 using MuladdMacro: @muladd
 
 export on_new_rule, refresh_rules  # generation tools

src/accumulation.jl

@@ -3,13 +3,48 @@
 
 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)
+"""
+    add!!(x, t::ImplacableThunk)
+
+The specialization of `add!!` for [`InplaceableThunk`](@ref) promises to only call
+`t.add!` on `x` if `x` is suitably mutable; otherwise it will be out of place.
+"""
+function add!!(x, t::InplaceableThunk)
+    return if is_inplaceable_destination(x)
+        t.add!(x)
+    else
+        x + t
+    end
+end
 
-function add!!(x::Array{<:Any, N}, y::AbstractArray{<:Any, N}) where N
-    return x .+= y
+function add!!(x::AbstractArray{<:Any, N}, y::AbstractArray{<:Any, N}) where N
+    return if is_inplaceable_destination(x)
+        x .+= y
+    else
+        x + y
+    end
+end
+
+
+"""
+    is_inplaceable_destination(x)
+
+Returns true if `x` is suitable for for storing inplace accumulation of gradients.
+For arrays this boils down `x .= y` if will work to mutate `x`, if `y` is an appropriate
+differential.
+"""
+is_inplaceable_destination(::Any) = false
+is_inplaceable_destination(::Array) = true
+is_inplaceable_destination(::SparseVector) = true
+is_inplaceable_destination(::SparseMatrixCSC) = true
+is_inplaceable_destination(::BitArray) = true
+function is_inplaceable_destination(x::AbstractArray)
+    p = parent(x)
+    p === x && return false  # no parent
+    # basically all wrapper types delegate `setindex!` to their `parent` after some
+    # processing and so are mutable if their `parent` is.
+    return is_inplaceable_destination(p)
 end

test/accumulation.jl

@@ -1,4 +1,23 @@
 @testset "accumulation.jl" begin
+    @testset "is_inplaceable_destination" begin
+        is_inplaceable_destination = ChainRulesCore.is_inplaceable_destination
+
+        @test is_inplaceable_destination([1, 2, 3, 4])
+        @test !is_inplaceable_destination(1:4)
+
+        @test is_inplaceable_destination(Diagonal([1, 2, 3, 4]))
+        @test !is_inplaceable_destination(Diagonal(1:4))
+
+        @test is_inplaceable_destination(view([1, 2, 3, 4], :, :))
+        @test !is_inplaceable_destination(view(1:4, :, :))
+
+        @test !is_inplaceable_destination(1.3)
+        @test is_inplaceable_destination(falses(4))
+        @test is_inplaceable_destination(spzeros(4))
+        @test is_inplaceable_destination(spzeros(2, 2))
+        @test !is_inplaceable_destination(@SVector [1, 2, 3])
+    end
+
     @testset "add!!" begin
         @testset "scalar" begin
             @test 16 == add!!(12, 4)
@@ -11,33 +30,50 @@
             @test 16 == add!!(16, DoesNotExist())  # Should this be an error?
         end
 
-        @testset "Array" begin
-            @testset "Happy Path" begin
+        @testset "add!!(::AbstractArray, ::AbstractArray)" begin
+            @testset "LHS Array (inplace)" 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]
+                    accumuland = -1.0*ones(2,2)
+                    ret = add!!(accumuland, A)
+                    @test ret === accumuland  # must be same object
+                    @test accumuland == [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]
+                    accumuland = -1.0*ones(2,2)
+                    ret = add!!(accumuland, A)
+                    @test ret === accumuland  # must be same object
+                    @test accumuland == [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]
+                    accumuland = -1.0*ones(2,2)
+                    ret = add!!(accumuland, A)
+                    @test ret === accumuland  # must be same object
+                    @test accumuland == [0.0 -1.0; -1.0 1.0]
                 end
             end
 
+            @testset "add!!(::StaticArray, ::Array) (out of place)" begin
+                A = [1.0 2.0; 3.0 4.0]
+                accumuland = @SMatrix [-1.0 -1.0; -1.0 -1.0]
+                ret = add!!(accumuland, A)
+                @test ret == [0.0 1.0; 2.0 3.0]  # must return right answer
+                @test ret !== accumuland  # must not be same object
+                @test accumuland == [-1.0 -1.0; -1.0 -1.0]  # must not have changed
+            end
+
+            @testset "add!!(::Diagonal{<:Vector}, ::Diagonal{<:Vector}) (inplace)" begin
+                A = Diagonal([1.0, 2.0])
+                accumuland = Diagonal([-2.0, -2.0])
+                ret = add!!(accumuland, A)
+                @test ret === accumuland  # must be same object
+                @test accumuland == Diagonal([-1.0, 0.0])
+            end
+
             @testset "Unhappy Path" begin
                 # wrong length
                 @test_throws DimensionMismatch add!!(ones(4,4), ones(2,2))
@@ -49,16 +85,26 @@
         end
 
         @testset "InplaceableThunk" begin
-            A=[1.0 2.0; 3.0 4.0]
             ithunk = InplaceableThunk(
-                @thunk(A*B),
-                x -> x.+=A
+                @thunk(-1.0*ones(2, 2)),
+                x -> x .-= ones(2, 2)
             )
 
-            accumuland = -1.0*ones(2,2)
-            ret = add!!(accumuland, ithunk)
-            @test ret === accumuland  # must be same object
-            @test accumuland == [0.0 1.0; 2.0 3.0]
+            @testset "in place" begin
+                accumuland = [1.0 2.0; 3.0 4.0]
+                ret = add!!(accumuland, ithunk)
+                @test ret == [0.0 1.0; 2.0 3.0]  # must return right answer
+                @test ret === accumuland  # must be same object
+            end
+
+            @testset "out of place" begin
+                accumuland = @SMatrix [1.0 2.0; 3.0 4.0]
+
+                ret = add!!(accumuland, ithunk)
+                @test ret == [0.0 1.0; 2.0 3.0]  # must return right answer
+                @test ret !== accumuland  # must not be same object
+                @test accumuland == [1.0 2.0; 3.0 4.0]  # must not have mutated
+            end
         end
     end
 end

test/runtests.jl

@@ -3,6 +3,7 @@ using BenchmarkTools
 using ChainRulesCore
 using LinearAlgebra: Diagonal, dot
 using StaticArrays
+using SparseArrays
 using Test
 
 @testset "ChainRulesCore" begin