Skip to content

Commit

Permalink
Add more methods on thunks to make ChainRules work on 1.0 (#381)
Browse files Browse the repository at this point in the history
  • Loading branch information
mzgubic authored Jun 23, 2021
1 parent 66bd0a1 commit 653cd7d
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 1 deletion.
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
6 changes: 6 additions & 0 deletions src/differentials/thunks.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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...)
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 8 additions & 0 deletions test/differentials/thunks.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down

2 comments on commit 653cd7d

@mzgubic
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/39487

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.10.9 -m "<description of version>" 653cd7d626183faf70107a6c639c04ae4a85f0c5
git push origin v0.10.9

Please sign in to comment.