Port a few rules from Diffractor (#463)

* Handle AbstractZero in a few LinearAlgebra rules

* Mark eltype(::Type) nondiff

* Add rules for typeassert/ifelse

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

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.2.0"
+version = "1.3.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/nondiff.jl

@@ -128,6 +128,7 @@
 @non_differentiable eachline(::AbstractString)
 @non_differentiable eachline(::IO)
 @non_differentiable eachmatch(::Regex, ::AbstractString)
+@non_differentiable eltype(::Type)
 @non_differentiable endswith(::AbstractString, ::AbstractString)
 @non_differentiable endswith(::AbstractString, ::Regex)
 @non_differentiable eof(::Any)

src/rulesets/Core/core.jl

@@ -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

src/rulesets/LinearAlgebra/structured.jl

@@ -107,6 +107,8 @@ 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)
+# currently needed by Diffractor (ref https://github.com/JuliaDiff/Diffractor.jl/issues/25)
+_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)
@@ -116,13 +118,17 @@ end
 _Adjoint_vec_pullback(ȳ::Tangent) = (NoTangent(), vec(ȳ.parent))
 _Adjoint_vec_pullback(ȳ::AbstractMatrix) = (NoTangent(), vec(adjoint(ȳ)))
 _Adjoint_vec_pullback(ȳ::AbstractThunk) = return _Adjoint_vec_pullback(unthunk(ȳ))
+# currently needed by Diffractor (ref https://github.com/JuliaDiff/Diffractor.jl/issues/25)
+_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)
+# currently needed by Diffractor (ref https://github.com/JuliaDiff/Diffractor.jl/issues/25)
+_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)
@@ -132,6 +138,8 @@ end
 _adjoint_vec_pullback(ȳ::Tangent) = (NoTangent(), vec(ȳ.parent))
 _adjoint_vec_pullback(ȳ::AbstractMatrix) = (NoTangent(), vec(adjoint(ȳ)))
 _adjoint_vec_pullback(ȳ::AbstractThunk) = return _adjoint_vec_pullback(unthunk(ȳ))
+# currently needed by Diffractor (ref https://github.com/JuliaDiff/Diffractor.jl/issues/25)
+_adjoint_vec_pullback(ȳ::AbstractZero) = (NoTangent(), ȳ)
 function rrule(::typeof(adjoint), A::AbstractVector{<:Number})
     return adjoint(A), _adjoint_vec_pullback
 end
@@ -145,6 +153,8 @@ 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)
+# currently needed by Diffractor (ref https://github.com/JuliaDiff/Diffractor.jl/issues/25)
+_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)
@@ -154,13 +164,17 @@ end
 _Transpose_vec_pullback(ȳ::Tangent) = (NoTangent(), vec(ȳ.parent))
 _Transpose_vec_pullback(ȳ::AbstractMatrix) = (NoTangent(), vec(Transpose(ȳ)))
 _Transpose_vec_pullback(ȳ::AbstractThunk) = return _Transpose_vec_pullback(unthunk(ȳ))
+# currently needed by Diffractor (ref https://github.com/JuliaDiff/Diffractor.jl/issues/25)
+_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)
+# currently needed by Diffractor (ref https://github.com/JuliaDiff/Diffractor.jl/issues/25)
+_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)
@@ -170,6 +184,8 @@ end
 _transpose_vec_pullback(ȳ::Tangent) = (NoTangent(), vec(ȳ.parent))
 _transpose_vec_pullback(ȳ::AbstractMatrix) = (NoTangent(), vec(transpose(ȳ)))
 _transpose_vec_pullback(ȳ::AbstractThunk) = return _transpose_vec_pullback(unthunk(ȳ))
+# currently needed by Diffractor (ref https://github.com/JuliaDiff/Diffractor.jl/issues/25)
+_transpose_vec_pullback(ȳ::AbstractZero) = (NoTangent(), ȳ)
 function rrule(::typeof(transpose), A::AbstractVector{<:Number})
     return transpose(A), _transpose_vec_pullback
 end

test/rulesets/Core/core.jl

@@ -0,0 +1,12 @@
+@testset "typeassert" begin
+    test_rrule(typeassert, 1.1, Float64)
+    test_frule(typeassert, 1.1, Float64)
+end
+
+@testset "ifelse" begin
+    test_rrule(ifelse, true, 1.1, 2.0)
+    test_frule(ifelse, false, 1.1, 2.0)
+
+    test_rrule(ifelse, true, [1.1], [2.0]; check_inferred=false)
+    test_frule(ifelse, false, [1.1], [2.0]; check_inferred=false)
+end

test/runtests.jl

@@ -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")