Take is_inplaceable_destination seriously (#577)

* take is_inplaceable_destination seriously

* fix tests on 1.9

* Update src/accumulation.jl

Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: Frames Catherine White <oxinabox@ucc.asn.au>

* three tweaks

Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com>
Co-authored-by: Frames Catherine White <oxinabox@ucc.asn.au>

Author: 3 people

src/ChainRulesCore.jl

@@ -12,7 +12,7 @@ export frule_via_ad, rrule_via_ad
 # definition helper macros
 export @non_differentiable, @opt_out, @scalar_rule, @thunk, @not_implemented
 export ProjectTo, canonicalize, unthunk  # tangent operations
-export add!!  # gradient accumulation operations
+export add!!, is_inplaceable_destination  # gradient accumulation operations
 export ignore_derivatives, @ignore_derivatives
 # tangents
 export Tangent, NoTangent, InplaceableThunk, Thunk, ZeroTangent, AbstractZero, AbstractThunk

src/accumulation.jl

@@ -38,30 +38,41 @@ end
     is_inplaceable_destination(x) -> Bool
 
 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
-tangent.
-Wrapper array types do not need to overload this if they overload `Base.parent`, and are
-`is_inplaceable_destination` if and only if their parent array is.
-Other types should overload this, as it defaults to `false`.
+For arrays this means `x .= y` will mutate `x`, if `y` is an appropriate tangent.
+
+Here "appropriate" means that `y` cannot be complex unless `x` is too,
+and that for structured matrices like `x isa Diagonal`, `y` shares this structure.
+
+!!! note "history"
+    Wrapper array types should overload this function if they can be written into.
+    Before ChainRulesCore 1.16, it would guess `true` for most wrappers based on `parent`,
+    but this is not safe, e.g. it will lead to an error with ReadOnlyArrays.jl. 
+
+There must always be a correct non-mutating path, so in uncertain cases,
+this function returns `false`.
 """
 is_inplaceable_destination(::Any) = false
+
 is_inplaceable_destination(::Array) = true
+is_inplaceable_destination(:: Array{<:Integer}) = false
+
 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)
+
+function is_inplaceable_destination(x::SubArray)
+    alpha = is_inplaceable_destination(parent(x))
+    beta = x.indices isa Tuple{Vararg{Union{Integer, Base.Slice, UnitRange}}}
+    return alpha && beta
 end
 
+for T in [:PermutedDimsArray, :ReshapedArray]
+    @eval is_inplaceable_destination(x::Base.$T) = is_inplaceable_destination(parent(x))
+end
+for T in [:Adjoint, :Transpose, :Diagonal, :UpperTriangular, :LowerTriangular]
+    @eval is_inplaceable_destination(x::LinearAlgebra.$T) = is_inplaceable_destination(parent(x))
+end
 # Hermitian and Symmetric are too fussy to deal with right now
 # https://github.com/JuliaLang/julia/issues/38056
-# TODO: https://github.com/JuliaDiff/ChainRulesCore.jl/issues/236
-is_inplaceable_destination(::LinearAlgebra.Hermitian) = false
-is_inplaceable_destination(::LinearAlgebra.Symmetric) = false
 
 function debug_add!(accumuland, t::InplaceableThunk)
     returned_value = t.add!(accumuland)

test/accumulation.jl

@@ -2,23 +2,26 @@
     @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([1.0, 2.0, 3.0])
+        @test !is_inplaceable_destination([1, 2, 3, 4])  # gradients cannot reliably be written into integer arrays
+        @test !is_inplaceable_destination(1:4.0)
 
-        @test is_inplaceable_destination(Diagonal([1, 2, 3, 4]))
-        @test !is_inplaceable_destination(Diagonal(1:4))
+        @test is_inplaceable_destination(Diagonal([1.0, 2.0, 3.0]))
+        @test !is_inplaceable_destination(Diagonal(1:4.0))
 
-        @test is_inplaceable_destination(view([1, 2, 3, 4], :, :))
-        @test !is_inplaceable_destination(view(1:4, :, :))
+        @test is_inplaceable_destination(view([1.0, 2.0, 3.0], :, :))
+        @test is_inplaceable_destination(view([1.0 2.0; 3.0 4.0], :, 2))
+        @test !is_inplaceable_destination(view(1:4.0, :, :))
+        mat = view([1.0, 2.0, 3.0], :, fill(1, 10))
+        @test !is_inplaceable_destination(mat)  # The concern is that `mat .+= x` is unsafe on GPU / parallel.
 
-        @test is_inplaceable_destination(falses(4))
+        @test !is_inplaceable_destination(falses(4))  # gradients can never be written into boolean
         @test is_inplaceable_destination(spzeros(4))
         @test is_inplaceable_destination(spzeros(2, 2))
 
-        @test !is_inplaceable_destination(1.3)
-        @test !is_inplaceable_destination(@SVector [1, 2, 3])
-        @test !is_inplaceable_destination(Hermitian([1 2; 2 4]))
-        @test !is_inplaceable_destination(Symmetric([1 2; 2 4]))
+        @test !is_inplaceable_destination(1:3.0)
+        @test !is_inplaceable_destination(@SVector [1.0, 2.0, 3.0])
+        @test !is_inplaceable_destination(Hermitian([1.0 2.0; 2.0 4.0]))
     end
 
     @testset "add!!" begin

test/tangent_types/tangent.jl

@@ -109,11 +109,11 @@ end
         @test NoTangent() === @inferred Base.tail(ntang1)
 
         # TODO: uncomment this once https://github.com/JuliaLang/julia/issues/35516
-        if VERSION >= v"1.8-"
-            @test haskey(Tangent{Tuple{Float64}}(2.0), 1) == true
-        else
-            @test_broken haskey(Tangent{Tuple{Float64}}(2.0), 1) == true
-        end
+        # if VERSION >= v"1.8-"
+        #     @test haskey(Tangent{Tuple{Float64}}(2.0), 1) == true
+        # else
+        #     @test_broken haskey(Tangent{Tuple{Float64}}(2.0), 1) == true
+        # end
         @test_broken hasproperty(Tangent{Tuple{Float64}}(2.0), 2) == false
 
         @test length(Tangent{Foo}(; x=2.5)) == 1
@@ -148,12 +148,16 @@ end
         cr = Tangent{Tuple{String,Int,Int}}("something", 2, 1)
         @test reverse(c) === cr
 
-        # can't reverse a named tuple or a dict
-        @test_throws MethodError reverse(Tangent{Foo}(; x=1.0, y=2.0))
+        if VERSION < v"1.9-"
+            # can't reverse a named tuple or a dict
+            @test_throws MethodError reverse(Tangent{Foo}(; x=1.0, y=2.0))
 
-        d = Dict(:x => 1, :y => 2.0)
-        cdict = Tangent{typeof(d),typeof(d)}(d)
-        @test_throws MethodError reverse(Tangent{Foo}())
+            d = Dict(:x => 1, :y => 2.0)
+            cdict = Tangent{typeof(d),typeof(d)}(d)
+            @test_throws MethodError reverse(Tangent{Foo}())
+        else
+            # These now work but do we care?
+        end
     end
 
     @testset "unset properties" begin