Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Port a few rules from Diffractor #463

Merged
merged 16 commits into from
Jul 30, 2021
Merged

Port a few rules from Diffractor #463

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
28b2084
Handle AbstractZero in a few LinearAlgebra rules
Keno Jul 1, 2021
c4baa5c
Mark eltype(::Type) nondiff
Keno Jul 1, 2021
2261929
Add unthunk in reshape rule
Keno Jul 1, 2021
823bb1c
Add rules for typeassert/ifelse
Keno Jul 1, 2021
8c33480
Apply suggestions from code review
simeonschaub Jul 28, 2021
d95337d
make `ifelse` pullback type stable for numbers
simeonschaub Jul 28, 2021
a3a16f1
Merge remote-tracking branch 'origin/master' into kf/diffractorbackport
simeonschaub Jul 28, 2021
fea23b6
Merge remote-tracking branch 'origin/master' into kf/diffractorbackport
simeonschaub Jul 29, 2021
014c4ff
:facepalm:
simeonschaub Jul 29, 2021
3d6827e
skip some inference tests
simeonschaub Jul 29, 2021
c005d8e
Update src/rulesets/Core/core.jl
simeonschaub Jul 29, 2021
87d1170
Apply suggestions from code review
simeonschaub Jul 30, 2021
c509549
add some comments
simeonschaub Jul 30, 2021
9a1ff23
bump version
simeonschaub Jul 30, 2021
16dcd5b
remove duplicate rules
simeonschaub Jul 30, 2021
998846f
Update Project.toml
simeonschaub Jul 30, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/rulesets/Base/nondiff.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -461,3 +461,11 @@ VERSION >= v"1.1" && @non_differentiable Sys.isopenbsd(::Symbol)
@non_differentiable Threads.nthreads()
@non_differentiable Threads.threadid()
@non_differentiable Threads.threadid(::Task)

@non_differentiable similar(::Any...)
@non_differentiable eltype(::Type)

@non_differentiable one(::Any)
@non_differentiable ones(::Any...)
@non_differentiable zero(::Any)
@non_differentiable zeros(::Any...)
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
@non_differentiable similar(::Any...)
@non_differentiable eltype(::Type)
@non_differentiable one(::Any)
@non_differentiable ones(::Any...)
@non_differentiable zero(::Any)
@non_differentiable zeros(::Any...)
@non_differentiable eltype(::Type)

could we also move the eltype to the right place alphabetically?

Copy link
Member

Choose a reason for hiding this comment

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

Do you mean

Suggested change
@non_differentiable similar(::Any...)
@non_differentiable eltype(::Type)
@non_differentiable one(::Any)
@non_differentiable ones(::Any...)
@non_differentiable zero(::Any)
@non_differentiable zeros(::Any...)
@non_differentiable eltype(::Type)
@non_differentiable similar(::Any...)
@non_differentiable one(::Any)
@non_differentiable ones(::Any...)
@non_differentiable zero(::Any)
@non_differentiable zeros(::Any...)

?

Copy link
Member

Choose a reason for hiding this comment

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

No, sorry, I did mean to delete the others, since they are already in base.jl on master, see e.g.

ChainRules.jl/src/rulesets/Base/nondiff.jl

Line 307 in 3f18cd6

@non_differentiable one(::Any) 

julia> @which ones
Base

Copy link
Member

Choose a reason for hiding this comment

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

actually eltype also seems to belong to Base

Copy link
Member

Choose a reason for hiding this comment

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

Ah, sorry, I misunderstood. eltype should still be in Base/nondiff.jl though, right?

Copy link
Member

Choose a reason for hiding this comment

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

Correct! My bad, forgot that Base/base.jl exists 😄

19 changes: 18 additions & 1 deletion src/rulesets/Core/core.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,5 +7,22 @@
@non_differentiable Core.isdefined(::Any, ::Any)
@non_differentiable Core.:(<:)(::Any, ::Any)

@non_differentiable Core.apply_type(::Any, ::Any)
@non_differentiable Core.apply_type(::Any, ::Any...)
@non_differentiable Core.typeof(::Any)

frule((_, ẋ, _), ::typeof(typeassert), x, T) = (typeassert(x, T), ẋ)
function rrule(::typeof(typeassert), x, T)
typeassert_pullback(Δ) = (NoTangent(), Δ, NoTangent())
return typeassert(x, T), typeassert_pullback
end

frule((_, _, ȧ, ḃ), ::typeof(ifelse), c, a, b) = (ifelse(c, a, b), ifelse(c, ȧ, ḃ))
function rrule(::typeof(ifelse), c, a, b)
ifelse_pullback(Δ) = (NoTangent(), NoTangent(), ifelse(c, Δ, ZeroTangent()), ifelse(c, ZeroTangent(), Δ))
return ifelse(c, a, b), ifelse_pullback
end
# ensure type stability for numbers
function rrule(::typeof(ifelse), c, a::Number, b::Number)
ifelse_pullback(Δ) = (NoTangent(), NoTangent(), ifelse(c, Δ, zero(Δ)), ifelse(c, zero(Δ), Δ))
return ifelse(c, a, b), ifelse_pullback
end
8 changes: 8 additions & 0 deletions src/rulesets/LinearAlgebra/structured.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@ end
_Adjoint_mat_pullback(ȳ::Tangent, proj) = (NoTangent(), proj(ȳ.parent))
_Adjoint_mat_pullback(ȳ::AbstractVecOrMat, proj) = (NoTangent(), proj(adjoint(ȳ)))
_Adjoint_mat_pullback(ȳ::AbstractThunk, proj) = return _Adjoint_mat_pullback(unthunk(ȳ), proj)
_Adjoint_mat_pullback(ȳ::AbstractZero, proj) = (NoTangent(), proj(ȳ))
function rrule(::Type{<:Adjoint}, A::AbstractMatrix{<:Number})
project_A = ProjectTo(A)
Adjoint_mat_pullback(ȳ) = _Adjoint_mat_pullback(ȳ, project_A)
Expand All @@ -116,13 +117,15 @@ end
_Adjoint_vec_pullback(ȳ::Tangent) = (NoTangent(), vec(ȳ.parent))
_Adjoint_vec_pullback(ȳ::AbstractMatrix) = (NoTangent(), vec(adjoint(ȳ)))
_Adjoint_vec_pullback(ȳ::AbstractThunk) = return _Adjoint_vec_pullback(unthunk(ȳ))
_Adjoint_vec_pullback(ȳ::AbstractZero) = (NoTangent(), ȳ)
function rrule(::Type{<:Adjoint}, A::AbstractVector{<:Number})
return Adjoint(A), _Adjoint_vec_pullback
end

_adjoint_mat_pullback(ȳ::Tangent, proj) = (NoTangent(), proj(ȳ.parent))
_adjoint_mat_pullback(ȳ::AbstractVecOrMat, proj) = (NoTangent(), proj(adjoint(ȳ)))
_adjoint_mat_pullback(ȳ::AbstractThunk, proj) = return _adjoint_mat_pullback(unthunk(ȳ), proj)
_adjoint_mat_pullback(ȳ::AbstractZero, proj) = (NoTangent(), proj(ȳ))
function rrule(::typeof(adjoint), A::AbstractMatrix{<:Number})
project_A = ProjectTo(A)
adjoint_mat_pullback(ȳ) = _adjoint_mat_pullback(ȳ, project_A)
Expand All @@ -132,6 +135,7 @@ end
_adjoint_vec_pullback(ȳ::Tangent) = (NoTangent(), vec(ȳ.parent))
_adjoint_vec_pullback(ȳ::AbstractMatrix) = (NoTangent(), vec(adjoint(ȳ)))
_adjoint_vec_pullback(ȳ::AbstractThunk) = return _adjoint_vec_pullback(unthunk(ȳ))
_adjoint_vec_pullback(ȳ::AbstractZero) = (NoTangent(), ȳ)
function rrule(::typeof(adjoint), A::AbstractVector{<:Number})
return adjoint(A), _adjoint_vec_pullback
end
Expand All @@ -145,6 +149,7 @@ end
_Transpose_mat_pullback(ȳ::Tangent, proj) = (NoTangent(), proj(ȳ.parent))
_Transpose_mat_pullback(ȳ::AbstractVecOrMat, proj) = (NoTangent(), proj(Transpose(ȳ)))
_Transpose_mat_pullback(ȳ::AbstractThunk, proj) = return _Transpose_mat_pullback(unthunk(ȳ), proj)
_Transpose_mat_pullback(ȳ::AbstractZero, proj) = (NoTangent(), proj(ȳ))
function rrule(::Type{<:Transpose}, A::AbstractMatrix{<:Number})
project_A = ProjectTo(A)
Transpose_mat_pullback(ȳ) = _Transpose_mat_pullback(ȳ, project_A)
Expand All @@ -154,13 +159,15 @@ end
_Transpose_vec_pullback(ȳ::Tangent) = (NoTangent(), vec(ȳ.parent))
_Transpose_vec_pullback(ȳ::AbstractMatrix) = (NoTangent(), vec(Transpose(ȳ)))
_Transpose_vec_pullback(ȳ::AbstractThunk) = return _Transpose_vec_pullback(unthunk(ȳ))
_Transpose_vec_pullback(ȳ::AbstractZero) = (NoTangent(), ȳ)
function rrule(::Type{<:Transpose}, A::AbstractVector{<:Number})
return Transpose(A), _Transpose_vec_pullback
end

_transpose_mat_pullback(ȳ::Tangent, proj) = (NoTangent(), proj(ȳ.parent))
_transpose_mat_pullback(ȳ::AbstractVecOrMat, proj) = (NoTangent(), proj(transpose(ȳ)))
_transpose_mat_pullback(ȳ::AbstractThunk, proj) = return _transpose_mat_pullback(unthunk(ȳ), proj)
_transpose_mat_pullback(ȳ::AbstractZero, proj) = (NoTangent(), proj(ȳ))
function rrule(::typeof(transpose), A::AbstractMatrix{<:Number})
project_A = ProjectTo(A)
transpose_mat_pullback(ȳ) = _transpose_mat_pullback(ȳ, project_A)
Expand All @@ -170,6 +177,7 @@ end
_transpose_vec_pullback(ȳ::Tangent) = (NoTangent(), vec(ȳ.parent))
_transpose_vec_pullback(ȳ::AbstractMatrix) = (NoTangent(), vec(transpose(ȳ)))
_transpose_vec_pullback(ȳ::AbstractThunk) = return _transpose_vec_pullback(unthunk(ȳ))
_transpose_vec_pullback(ȳ::AbstractZero) = (NoTangent(), ȳ)
function rrule(::typeof(transpose), A::AbstractVector{<:Number})
return transpose(A), _transpose_vec_pullback
end
Expand Down
12 changes: 12 additions & 0 deletions test/rulesets/Core/core.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
@testset "typeassert" begin
test_rrule(typeassert, 1.0, Float64 ⊢ NoTangent())
test_frule(typeassert, 1.0, Float64 ⊢ NoTangent())
simeonschaub marked this conversation as resolved.
Show resolved Hide resolved
end

@testset "ifelse" begin
test_rrule(ifelse, true ⊢ NoTangent(), 1.0, 2.0)
test_frule(ifelse, false ⊢ NoTangent(), 1.0, 2.0)

test_rrule(ifelse, true ⊢ NoTangent(), [1.0], [2.0]; check_inferred=false)
test_frule(ifelse, false ⊢ NoTangent(), [1.0], [2.0]; check_inferred=false)
simeonschaub marked this conversation as resolved.
Show resolved Hide resolved
end
4 changes: 4 additions & 0 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,10 @@ println("Testing ChainRules.jl")
include_test("test_helpers.jl")
println()
@testset "rulesets" begin
@testset "Core" begin
include_test("rulesets/Core/core.jl")
end

@testset "Base" begin
include_test("rulesets/Base/base.jl")
include_test("rulesets/Base/fastmath_able.jl")
Expand Down