Accumulated control flow fixes

Author: Keno

Manifest.toml

@@ -18,9 +18,7 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
 [[ChainRules]]
 deps = ["ChainRulesCore", "Compat", "LinearAlgebra", "Random", "Statistics"]
-git-tree-sha1 = "948d680b9dfcedcbe58e4d380f3737b4f3e5ac7f"
-repo-rev = "kf/diffractorbackport"
-repo-url = "https://github.com/Keno/ChainRules.jl.git"
+path = "/home/keno/.julia/dev/ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 version = "0.8.18"
 

src/extra_rules.jl

@@ -145,6 +145,9 @@ function ChainRulesCore.rrule(::typeof(Core.tuple), args...)
     Core.tuple(args...), Δ->Core.tuple(NoTangent(), Δ...)
 end
 
+# TODO: What to do about these integer rules
+@ChainRulesCore.non_differentiable Base.rem(a::Integer, b::Type)
+
 ChainRulesCore.canonicalize(::ChainRulesCore.ZeroTangent) = ChainRulesCore.ZeroTangent()
 
 # Skip AD'ing through the axis computation

src/stage1/generated.jl

@@ -189,6 +189,18 @@ function (::∂⃖rrule{N})(z, z̄) where {N}
     y, ∂⃖rruleA{term_depth(N), 1}(∂⃖{minus1(N)}(), ȳ, z̄)
 end
 
+function (::∂⃖{N})(f::Core.IntrinsicFunction, args...) where {N}
+    # A few intrinsic functions are inserted by the compiler, so they need to
+    # be handled here. Otherwise, we just throw an appropriate error.
+    if f === Core.Intrinsics.not_int && length(args) == 1
+        return f(args...), EvenOddOdd{1, c_order(N)}(
+            Δ->(NoTangent(), NoTangent()),
+            Δ->NoTangent())
+    end
+
+    error("Rewrite reached intrinsic function $f. Missing rule?")
+end
+
 # The static parameter on `f` disables the compileable_sig heuristic
 function (::∂⃖{N})(f::T, args...) where {T, N}
     if N == 1
@@ -297,7 +309,7 @@ end
 
 struct tuple_back{M}; end
 (::tuple_back)(Δ::Tuple) = Core.tuple(NoTangent(), Δ...)
-(::tuple_back{N})(Δ::NoTangent) where {N} = Core.tuple(NoTangent(), ntuple(i->NoTangent(), N)...)
+(::tuple_back{N})(Δ::AbstractZero) where {N} = Core.tuple(NoTangent(), ntuple(i->Δ, N)...)
 
 function (::∂⃖{N})(::typeof(Core.tuple), args::Vararg{Any, M}) where {N, M}
     Core.tuple(args...),
@@ -333,7 +345,7 @@ end
     a.u(r)
 end
 
-function (this::∂⃖{N})(::typeof(Core._apply_iterate), iterate, f, args::Tuple...) where {N}
+function (this::∂⃖{N})(::typeof(Core._apply_iterate), iterate, f, args::Union{Tuple, NamedTuple}...) where {N}
     @assert iterate === Base.iterate
     x, ∂⃖f = Core._apply_iterate(iterate, this, (f,), args...)
     return x, ApplyOdd{1, c_order(N)}(UnApply{map(length, args)}(), ∂⃖f)

src/stage1/hacks.jl

@@ -1,7 +1,15 @@
 # Updated copy of the same code in Base, but with bugs fixed
-using Core.Compiler: count_added_node!, add!, NewSSAValue
+using Core.Compiler: count_added_node!, add!, NewSSAValue, add_pending!,
+    StmtRange, BasicBlock
 
+Base.length(c::Core.Compiler.NewNodeStream) = Core.Compiler.length(c)
 Base.setindex!(i::Instruction, args...) = Core.Compiler.setindex!(i, args...)
+Core.Compiler.BasicBlock(x::UnitRange) =
+    BasicBlock(StmtRange(first(x), last(x)))
+Core.Compiler.BasicBlock(x::UnitRange, preds::Vector{Int}, succs::Vector{Int}) =
+    BasicBlock(StmtRange(first(x), last(x)), preds, succs)
+Base.size(x::Core.Compiler.UnitRange) = Core.Compiler.size(x)
+
 function my_insert_node!(compact::IncrementalCompact, before, inst::NewInstruction, attach_after::Bool=false)
     @assert inst.effect_free_computed
     if isa(before, SSAValue)