Add more methods on thunks to make ChainRules work on 1.0 (#381)

Author: mzgubic

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.10.8"
+version = "0.10.9"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/differentials/thunks.jl

@@ -41,8 +41,10 @@ Base.mapreduce(f, op, a::AbstractThunk; kws...) = mapreduce(f, op, unthunk(a); k
 function Base.mapreduce(f, op, itr, a::AbstractThunk; kws...)
     return mapreduce(f, op, itr, unthunk(a); kws...)
 end
+Base.sum(a::AbstractThunk; kws...) = sum(unthunk(a); kws...)
 Base.sum!(r, A::AbstractThunk; kws...) = sum!(r, unthunk(A); kws...)
 
+Base.fill(a::AbstractThunk, b::Integer) = fill(unthunk(a), b)
 Base.vec(a::AbstractThunk) = vec(unthunk(a))
 Base.reshape(a::AbstractThunk, args...) = reshape(unthunk(a), args...)
 Base.getindex(a::AbstractThunk, args...) = getindex(unthunk(a), args...)
@@ -78,6 +80,7 @@ LinearAlgebra.dot(a::AbstractThunk, b::AbstractThunk) = dot(unthunk(a), unthunk(
 LinearAlgebra.ldiv!(a, b::AbstractThunk) = throw(MutateThunkException())
 LinearAlgebra.rdiv!(a::AbstractThunk, b) = throw(MutateThunkException())
 
+LinearAlgebra.mul!(A, B::AbstractThunk, C) = mul!(A, unthunk(B), C)
 LinearAlgebra.mul!(C::AbstractThunk, A, B, α, β) = throw(MutateThunkException())
 function LinearAlgebra.mul!(C::AbstractThunk, A::AbstractThunk, B, α, β)
     return throw(MutateThunkException())
@@ -190,6 +193,9 @@ end
 
 Base.show(io::IO, x::Thunk) = print(io, "Thunk($(repr(x.f)))")
 
+Base.convert(::Type{<:Thunk}, a::AbstractZero) = @thunk(a)
+
+
 """
     InplaceableThunk(val::Thunk, add!::Function)
 

test/differentials/thunks.jl

@@ -22,6 +22,10 @@
         @test occursin(r"Thunk\(.*rand.*\)", rep)
     end
 
+    @testset "convert" begin
+        @test convert(Thunk, ZeroTangent()) isa Thunk
+    end
+
     @testset "unthunk" begin
         @test unthunk(@thunk(3)) == 3
         @test unthunk(@thunk(@thunk(3))) isa Thunk
@@ -107,8 +111,10 @@
             @test 3 == mapreduce(_ -> 1, +, t)
             @test 3 == mapreduce((_, _) -> 1, +, v, t)
         end
+        @test 10 == sum(@thunk([1 2; 3 4]))
         @test [4 6] == sum!([1 1], @thunk([1 2; 3 4]))
 
+        @test fill(3.2, 3) == fill(@thunk(3.2), 3)
         @test v == vec(t)
         @test [1 2 3] == reshape(t, 1, 3)
         @test 1 == getindex(t, 1)
@@ -152,6 +158,8 @@
                                               rdiv!(deepcopy(a), 2.0)
         end
 
+        @test mul!(deepcopy(a), a, a) == mul!(deepcopy(a), t, a)
+
         res = mul!(deepcopy(a), a, a, true, true)
         @test_throws MutateThunkException mul!(deepcopy(t), a, a, true, true)
         @test_throws MutateThunkException mul!(deepcopy(t), t, a, true, true)