Merge branch 'ox/type_stable_zero' into ox/mutation2

Author: oxinabox

src/codegen/forward.jl

@@ -34,7 +34,7 @@ function fwd_transform!(ci, mi, nargs, N)
             args = map(stmt.args) do stmt
                 emit!(mapstmt!(stmt))
             end
-            return Expr(:call, Core._apply_iterate, FwdIterate(ZeroBundle{N}(iterate)), ∂☆new{N}(), emit!(Expr(:call, tuple, args[1])), args[2:end]...)
+            return Expr(:call, Core._apply_iterate, FwdIterate(DNEBundle{N}(iterate)), ∂☆new{N}(), emit!(Expr(:call, tuple, args[1])), args[2:end]...)
         elseif isa(stmt, SSAValue)
             return SSAValue(ssa_mapping[stmt.id])
         elseif isa(stmt, Core.SlotNumber)
@@ -62,7 +62,7 @@ function fwd_transform!(ci, mi, nargs, N)
         # Always disable `@inbounds`, as we don't actually know if the AD'd
         # code is truly `@inbounds` or not.
         elseif isexpr(stmt, :boundscheck)
-            return ZeroBundle{N}(true)
+            return DNEBundle{N}(true)
         else
             # Fallback case, for literals.
             # If it is an Expr, then it is not a literal

src/codegen/forward_demand.jl

@@ -363,6 +363,6 @@ function forward_diff!(interp::ADInterpreter, ir::IRCode, src::CodeInfo, mi::Met
     rt = CC._ir_abstract_constant_propagation(interp, irsv)
 
     ir = compact!(ir)
-
+    
     return ir
 end

src/higher_fwd_rules.jl

@@ -19,27 +19,27 @@ end
 
 jeval(j, x) = j(x)
 for f in (sin, cos, exp)
-    function (∂☆ₙ::∂☆{N})(fb::ZeroBundle{N, typeof(f)}, x::TaylorBundle{N}) where {N}
+    function (∂☆ₙ::∂☆{N})(fb::AbstractZeroBundle{N, typeof(f)}, x::TaylorBundle{N}) where {N}
         njet(Val{N}(), primal(fb), primal(x))(x)
     end
-    function (∂⃖ₙ::∂⃖{N})(∂☆ₘ::∂☆{M}, fb::ZeroBundle{M, typeof(f)}, x::TaylorBundle{M}) where {N, M}
+    function (∂⃖ₙ::∂⃖{N})(∂☆ₘ::∂☆{M}, fb::AbstractZeroBundle{M, typeof(f)}, x::TaylorBundle{M}) where {N, M}
         ∂⃖ₙ(jeval, njet(Val{N+M}(), primal(fb), primal(x)), x)
     end
 end
 
 # TODO: It's a bit embarassing that we need to write these out, but currently the
 # compiler is not strong enough to automatically lift the frule. Let's hope we
 # can delete these in the near future.
-function (∂☆ₙ::∂☆{N})(fb::ZeroBundle{N, typeof(+)}, a::TaylorBundle{N}, b::TaylorBundle{N}) where {N}
+function (∂☆ₙ::∂☆{N})(fb::AbstractZeroBundle{N, typeof(+)}, a::TaylorBundle{N}, b::TaylorBundle{N}) where {N}
     TaylorBundle{N}(primal(a) + primal(b),
         map(+, a.tangent.coeffs, b.tangent.coeffs))
 end
 
-function (∂☆ₙ::∂☆{N})(fb::ZeroBundle{N, typeof(+)}, a::TaylorBundle{N}, b::ZeroBundle{N}) where {N}
+function (∂☆ₙ::∂☆{N})(fb::AbstractZeroBundle{N, typeof(+)}, a::TaylorBundle{N}, b::AbstractZeroBundle{N}) where {N}
     TaylorBundle{N}(primal(a) + primal(b), a.tangent.coeffs)
 end
 
-function (∂☆ₙ::∂☆{N})(fb::ZeroBundle{N, typeof(-)}, a::TaylorBundle{N}, b::TaylorBundle{N}) where {N}
+function (∂☆ₙ::∂☆{N})(fb::AbstractZeroBundle{N, typeof(-)}, a::TaylorBundle{N}, b::TaylorBundle{N}) where {N}
     TaylorBundle{N}(primal(a) - primal(b),
         map(-, a.tangent.coeffs, b.tangent.coeffs))
 end

src/stage1/broadcast.jl

@@ -9,7 +9,7 @@ end
 
 n_getfield(∂ₙ::∂☆{N}, b::ATB{N}, x::Union{Symbol, Int}) where {N} = ∂ₙ(ZeroBundle{N}(getfield), b, ZeroBundle{N}(x))
 
-function (∂ₙ::∂☆{N})(zc::ZeroBundle{N, typeof(copy)},
+function (∂ₙ::∂☆{N})(zc::AbstractZeroBundle{N, typeof(copy)},
                      bc::ATB{N, <:Broadcasted}) where {N}
   bc = ∂ₙ(ZeroBundle{N}(Broadcast.flatten), bc)
   args = n_getfield(∂ₙ, bc, :args)

src/stage1/forward.jl

@@ -98,9 +98,12 @@ struct ∂☆shuffle{N}; end
 
 function shuffle_base(r)
     (primal, dual) = r
-    if isa(dual, Union{NoTangent, ZeroTangent})
+    if dual isa NoTangent
         UniformBundle{1}(primal, dual)
     else
+        if dual isa ZeroTangent  # Normalize zero for type-stability reasons
+            dual = zero_tangent(primal)
+        end
         TaylorBundle{1}(primal, (dual,))
     end
 end
@@ -193,25 +196,25 @@ struct FwdMap{N, T<:AbstractTangentBundle{N}}
 end
 (f::FwdMap{N})(args::AbstractTangentBundle{N}...) where {N} = ∂☆{N}()(f.f, args...)
 
-function (::∂☆{N})(::ZeroBundle{N, typeof(map)}, f::ATB{N}, tup::TaylorBundle{N, <:Tuple}) where {N}
+function (::∂☆{N})(::AbstractZeroBundle{N, typeof(map)}, f::ATB{N}, tup::TaylorBundle{N, <:Tuple}) where {N}
     ∂vararg{N}()(map(FwdMap(f), destructure(tup))...)
 end
 
-function (::∂☆{N})(::ZeroBundle{N, typeof(map)}, f::ATB{N}, args::ATB{N, <:AbstractArray}...) where {N}
+function (::∂☆{N})(::AbstractZeroBundle{N, typeof(map)}, f::ATB{N}, args::ATB{N, <:AbstractArray}...) where {N}
     # TODO: This could do an inplace map! to avoid the extra rebundling
     rebundle(map(FwdMap(f), map(unbundle, args)...))
 end
 
-function (::∂☆{N})(::ZeroBundle{N, typeof(map)}, f::ATB{N}, args::ATB{N}...) where {N}
+function (::∂☆{N})(::AbstractZeroBundle{N, typeof(map)}, f::ATB{N}, args::ATB{N}...) where {N}
     ∂☆recurse{N}()(ZeroBundle{N, typeof(map)}(map), f, args...)
 end
 
 
-function (::∂☆{N})(f::ZeroBundle{N, typeof(ifelse)}, arg::ATB{N, Bool}, args::ATB{N}...) where {N}
+function (::∂☆{N})(f::AbstractZeroBundle{N, typeof(ifelse)}, arg::ATB{N, Bool}, args::ATB{N}...) where {N}
     ifelse(arg.primal, args...)
 end
 
-function (::∂☆{N})(f::ZeroBundle{N, typeof(Core.ifelse)}, arg::ATB{N, Bool}, args::ATB{N}...) where {N}
+function (::∂☆{N})(f::AbstractZeroBundle{N, typeof(Core.ifelse)}, arg::ATB{N, Bool}, args::ATB{N}...) where {N}
     Core.ifelse(arg.primal, args...)
 end
 
@@ -233,48 +236,48 @@ end
      primal(∂☆{N}()(ZeroBundle{N}(getindex), r, ZeroBundle{N}(2))))
 end
 
-function (this::∂☆{N})(::ZeroBundle{N, typeof(Core._apply_iterate)}, iterate::ATB{N}, f::ATB{N}, args::ATB{N}...) where {N}
+function (this::∂☆{N})(::AbstractZeroBundle{N, typeof(Core._apply_iterate)}, iterate::ATB{N}, f::ATB{N}, args::ATB{N}...) where {N}
     Core._apply_iterate(FwdIterate(iterate), this, (f,), args...)
 end
 
 
-function (this::∂☆{N})(::ZeroBundle{N, typeof(iterate)}, t::TaylorBundle{N, <:Tuple}) where {N}
+function (this::∂☆{N})(::AbstractZeroBundle{N, typeof(iterate)}, t::TaylorBundle{N, <:Tuple}) where {N}
     r = iterate(destructure(t))
     r === nothing && return ZeroBundle{N}(nothing)
     ∂vararg{N}()(r[1], ZeroBundle{N}(r[2]))
 end
 
-function (this::∂☆{N})(::ZeroBundle{N, typeof(iterate)}, t::TaylorBundle{N, <:Tuple}, a::ATB{N}, args::ATB{N}...) where {N}
+function (this::∂☆{N})(::AbstractZeroBundle{N, typeof(iterate)}, t::TaylorBundle{N, <:Tuple}, a::ATB{N}, args::ATB{N}...) where {N}
     r = iterate(destructure(t), primal(a), map(primal, args)...)
     r === nothing && return ZeroBundle{N}(nothing)
     ∂vararg{N}()(r[1], ZeroBundle{N}(r[2]))
 end
 
-function (this::∂☆{N})(::ZeroBundle{N, typeof(Base.indexed_iterate)}, t::TaylorBundle{N, <:Tuple}, i::ATB{N}) where {N}
+function (this::∂☆{N})(::AbstractZeroBundle{N, typeof(Base.indexed_iterate)}, t::TaylorBundle{N, <:Tuple}, i::ATB{N}) where {N}
     r = Base.indexed_iterate(destructure(t), primal(i))
     ∂vararg{N}()(r[1], ZeroBundle{N}(r[2]))
 end
 
-function (this::∂☆{N})(::ZeroBundle{N, typeof(Base.indexed_iterate)}, t::TaylorBundle{N, <:Tuple}, i::ATB{N}, st1::ATB{N}, st::ATB{N}...) where {N}
+function (this::∂☆{N})(::AbstractZeroBundle{N, typeof(Base.indexed_iterate)}, t::TaylorBundle{N, <:Tuple}, i::ATB{N}, st1::ATB{N}, st::ATB{N}...) where {N}
     r = Base.indexed_iterate(destructure(t), primal(i), primal(st1), map(primal, st)...)
     ∂vararg{N}()(r[1], ZeroBundle{N}(r[2]))
 end
 
-function (this::∂☆{N})(::ZeroBundle{N, typeof(Base.indexed_iterate)}, t::TangentBundle{N, <:Tuple}, i::ATB{N}, st::ATB{N}...) where {N}
+function (this::∂☆{N})(::AbstractZeroBundle{N, typeof(Base.indexed_iterate)}, t::TangentBundle{N, <:Tuple}, i::ATB{N}, st::ATB{N}...) where {N}
     ∂vararg{N}()(this(ZeroBundle{N}(getfield), t, i), ZeroBundle{N}(primal(i) + 1))
 end
 
-function (this::∂☆{N})(::ZeroBundle{N, typeof(getindex)}, t::TaylorBundle{N, <:Tuple}, i::ZeroBundle) where {N}
+function (this::∂☆{N})(::AbstractZeroBundle{N, typeof(getindex)}, t::TaylorBundle{N, <:Tuple}, i::AbstractZeroBundle) where {N}
     field_ind = primal(i)
     the_partials = ntuple(order_ind->partial(t, order_ind)[field_ind], N)
     TaylorBundle{N}(primal(t)[field_ind], the_partials)
 end
 
-function (this::∂☆{N})(::ZeroBundle{N, typeof(typeof)}, x::ATB{N}) where {N}
+function (this::∂☆{N})(::AbstractZeroBundle{N, typeof(typeof)}, x::ATB{N}) where {N}
     DNEBundle{N}(typeof(primal(x)))
 end
 
-function (this::∂☆{N})(f::ZeroBundle{N, Core.IntrinsicFunction}, args::ATB{N}...) where {N}
+function (this::∂☆{N})(f::AbstractZeroBundle{N, Core.IntrinsicFunction}, args::ATB{N}...) where {N}
     ff = primal(f)
     if ff === Base.not_int
         DNEBundle{N}(ff(map(primal, args)...))

src/stage1/mixed.jl

@@ -70,26 +70,26 @@ function (f::FwdIterate)(arg::ATB{N}, st) where {N}
      primal(∂☆{N}()(ZeroBundle{N}(getindex), r, ZeroBundle{N}(2))))
 end
 
-function (this::∂☆{N})(::ZeroBundle{N, typeof(Core._apply_iterate)}, iterate::ATB{N}, f::ATB{N}, args::ATB{N}...) where {N}
+function (this::∂☆{N})(::AbstractZeroBundle{N, typeof(Core._apply_iterate)}, iterate::ATB{N}, f::ATB{N}, args::ATB{N}...) where {N}
     Core._apply_iterate(FwdIterate(iterate), this, (f,), args...)
 end
 =#
 
-function (this::∂⃖{N})(that::∂☆{M}, ::ZeroBundle{M, typeof(Core._apply_iterate)},
+function (this::∂⃖{N})(that::∂☆{M}, ::AbstractZeroBundle{M, typeof(Core._apply_iterate)},
         iterate, f, args::ATB{M, <:Tuple}...) where {N, M}
     @assert primal(iterate) === Base.iterate
     x, ∂⃖f = Core._apply_iterate(FwdIterate(iterate), this, (that, f), args...)
     return x, ApplyOdd{1, c_order(N)}(UnApply{map(x->length(primal(x)), args)}(), ∂⃖f)
 end
 
 
-function ChainRules.rrule(∂::∂☆{N}, m::ZeroBundle{N, typeof(map)}, p::ZeroBundle{N, typeof(+)}, A::ATB{N}, B::ATB{N}) where {N}
+function ChainRules.rrule(∂::∂☆{N}, m::AbstractZeroBundle{N, typeof(map)}, p::AbstractZeroBundle{N, typeof(+)}, A::ATB{N}, B::ATB{N}) where {N}
     ∂(m, p, A, B), Δ->(NoTangent(), NoTangent(), NoTangent(), Δ, Δ)
 end
 
 mapev_unbundled(_, js, a) = rebundle(mapev(js, unbundle(a)))
-function (∂⃖ₙ::∂⃖{N})(∂☆ₘ::∂☆{M}, ::ZeroBundle{M, typeof(map)},
-                    f::ZeroBundle{M}, a::ATB{M, <:Array}) where {N, M}
+function (∂⃖ₙ::∂⃖{N})(∂☆ₘ::∂☆{M}, ::AbstractZeroBundle{M, typeof(map)},
+                    f::AbstractZeroBundle{M}, a::ATB{M, <:Array}) where {N, M}
     @assert Base.issingletontype(typeof(primal(f)))
     js = map(primal(a)) do x
         ∂f = ∂☆{N+M}()(ZeroBundle{N+M}(primal(f)),

src/stage1/recurse_fwd.jl

@@ -15,7 +15,6 @@ struct ∂☆new{N}; end
 function (::∂☆new{1})(B::Type, xs::AbstractTangentBundle{1}...)
     primal_args = map(primal, xs)
     the_primal = _construct(B, primal_args)
-    @info "∂☆new{1}"
     tangent_tup = map(first_partial, xs)
     the_partial = if B<:Tuple
         Tangent{B, typeof(tangent_tup)}(tangent_tup)
@@ -24,13 +23,14 @@ function (::∂☆new{1})(B::Type, xs::AbstractTangentBundle{1}...)
         tangent_nt = NamedTuple{names}(tangent_tup)
         StructuralTangent{B}(tangent_nt)
     end
+    @show typeof(the_partial)
+    #TODO: I think we need https://github.com/JuliaDiff/Diffractor.jl/pull/236/files here
     return TaylorBundle{1, B}(the_primal, (the_partial,))
 end
 
 function (::∂☆new{N})(B::Type, xs::AbstractTangentBundle{N}...) where {N}
     primal_args = map(primal, xs)
     the_primal = _construct(B, primal_args)
-    @info "∂☆new{N}"
     the_partials = ntuple(Val{N}()) do ii
         tangent_tup = map(x->partial(x, ii), xs)
         tangent = if B<:Tuple

src/tangent.jl

@@ -427,9 +427,11 @@ Creates a bundle with a zero tangent.
 """
 struct zero_bundle{N} end
 function (::zero_bundle{N})(primal) where N
-    if zero_tangent(primal) isa ZeroTangent
-        return ZeroBundle{N}(primal)
+    if zero_tangent(primal)  isa AbstractZero
+        return UniformBundle{N}(primal, zero_tangent(primal) )
     else
+        # Note: it is important that zero_tangent(primal) is called in ntuple
+        # so it gets distrinct values for each order, so it doesn't alias if mutated.
         return TaylorBundle{N}(primal, ntuple(_->zero_tangent(primal), N))
     end
 end

test/stage2_fwd.jl

@@ -6,32 +6,37 @@ module stage2_fwd
         @test sin′(1.0) == cos(1.0)
     end
     let sin′′ = Diffractor.dontuse_nth_order_forward_stage2(Tuple{typeof(mysin), Float64}, 2)
-        @test isa(sin′′, Core.OpaqueClosure{Tuple{Float64}, Float64})
+        # This broke some time between 1.10 and 1.11-DEV.10001
+        @test_broken isa(sin′′, Core.OpaqueClosure{Tuple{Float64}, Float64})
         @test sin′′(1.0) == -sin(1.0)
     end
 
     myminus(a, b) = a - b
     self_minus(a) = myminus(a, a)
     ChainRulesCore.@scalar_rule myminus(x, y) (true, -1)
     let self_minus′ = Diffractor.dontuse_nth_order_forward_stage2(Tuple{typeof(self_minus), Float64})
-        @test isa(self_minus′, Core.OpaqueClosure{Tuple{Float64}, Float64})
+        # This broke some time between 1.10 and 1.11-DEV.10001
+        @test_broken isa(self_minus′, Core.OpaqueClosure{Tuple{Float64}, Float64})
         @test self_minus′(1.0) == 0.
     end
     ChainRulesCore.@scalar_rule myminus(x, y) (true, true) # frule for `x - y`
     let self_minus′ = Diffractor.dontuse_nth_order_forward_stage2(Tuple{typeof(self_minus), Float64})
-        @test isa(self_minus′, Core.OpaqueClosure{Tuple{Float64}, Float64})
+        # This broke some time between 1.10 and 1.11-DEV.10001
+        @test_broken isa(self_minus′, Core.OpaqueClosure{Tuple{Float64}, Float64})
         @test self_minus′(1.0) == 2.
     end
 
     myminus2(a, b) = a - b
     self_minus2(a) = myminus2(a, a)
     let self_minus2′ = Diffractor.dontuse_nth_order_forward_stage2(Tuple{typeof(self_minus2), Float64})
-        @test isa(self_minus2′, Core.OpaqueClosure{Tuple{Float64}, Float64})
+        # This broke some time between 1.10 and 1.11-DEV.10001
+        @test_broken isa(self_minus2′, Core.OpaqueClosure{Tuple{Float64}, Float64})
         @test self_minus2′(1.0) == 0.
     end
     ChainRulesCore.@scalar_rule myminus2(x, y) (true, true) # frule for `x - y`
     let self_minus2′ = Diffractor.dontuse_nth_order_forward_stage2(Tuple{typeof(self_minus2), Float64})
-        @test isa(self_minus2′, Core.OpaqueClosure{Tuple{Float64}, Float64})
+        # This broke some time between 1.10 and 1.11-DEV.10001
+        @test_broken isa(self_minus2′, Core.OpaqueClosure{Tuple{Float64}, Float64})
         @test self_minus2′(1.0) == 2.
     end
 

test/tangent.jl

@@ -126,7 +126,7 @@ end
     zero_bundle = Diffractor.zero_bundle
 
     tup_zb = zero_bundle{2}()((1, 0))
-    @test tup_zb isa ZeroBundle{2}
+    @test tup_zb isa Diffractor.AbstractTangentBundle{2}
     @test iszero(tup_zb[TaylorTangentIndex(1)])
     @test iszero(tup_zb[TaylorTangentIndex(2)])