WIP: begin switching forward mode over to zero_bundles for mutation support

Author: oxinabox

src/codegen/forward.jl

@@ -69,7 +69,7 @@ function fwd_transform!(ci, mi, nargs, N)
             if isa(stmt, Expr)
                 error("Unexprected statement encountered. This is a bug in Diffractor. stmt=$stmt")
             end
-            return Expr(:call, ZeroBundle{N}, stmt)
+            return Expr(:call, zero_bundle{order}(), stmt)
         end
     end
 

src/codegen/forward_demand.jl

@@ -264,12 +264,12 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
             return transform!(ir, arg, order, maparg)
         elseif isa(arg, GlobalRef)
             @assert isconst(arg)
-            return ZeroBundle{order}(getfield(arg.mod, arg.name))
+            return zero_bundle{order}()(getfield(arg.mod, arg.name))
         elseif isa(arg, QuoteNode)
-            return ZeroBundle{order}(arg.value)
+            return zero_bundle{order}(){order}(arg.value)
         end
         @assert !isa(arg, Expr)
-        return ZeroBundle{order}(arg)
+        return zero_bundle{order}()(arg)
     end
 
     for (ssa, (order, custom)) in enumerate(ssa_orders)
@@ -309,7 +309,7 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
                     stmt = insert_node!(ir, ssa, NewInstruction(inst))
                 end
 
-                replace_call!(ir, SSAValue(ssa), Expr(:call, ZeroBundle{order}, stmt))
+                replace_call!(ir, SSAValue(ssa), Expr(:call, zero_bundle{order}(), stmt))
             elseif isa(stmt, SSAValue) || isa(stmt, QuoteNode)
                 inst[:inst] = maparg(stmt, SSAValue(ssa), order)
                 inst[:type] = Any
@@ -329,7 +329,7 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
                 inst[:type] = Any
                 inst[:flag] |= CC.IR_FLAG_REFINED
             else
-                val = ZeroBundle{order}(inst[:inst])
+                val = zero_bundle{order}()(inst[:inst])
                 inst[:inst] = val
                 inst[:type] = Const(val)
             end

src/stage1/recurse_fwd.jl

@@ -50,7 +50,7 @@ _construct(::Type{B}, args) where B<:Tuple = B(args)
 # Hack for making things that do not have public constructors constructable:
 @generated _construct(B::Type, args) = Expr(:splatnew, :B, :args)
 
-@generated (::∂☆new{N})(B::Type) where {N} = return :(ZeroBundle{$N}($(Expr(:new, :B))))
+@generated (::∂☆new{N})(B::Type) where {N} = return :(zero_bundle{order}()($(Expr(:new, :B))))
 
 # Sometimes we don't know whether or not we need to the ZeroBundle when doing
 # the transform, so this can happen - allow it for now.

src/tangent.jl

@@ -418,3 +418,18 @@ end
 function ChainRulesCore.rrule(::typeof(rebundle), atb)
     rebundle(atb), Δ->throw(Δ)
 end
+
+
+"""
+    (::zero_bundle{N})(primal)
+
+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)
+    else
+        return TaylorBundle{N}(primal, ntuple(_->zero_tangent(primal), N))
+    end
+end

test/forward_mutation.jl

@@ -0,0 +1,16 @@
+using Diffractor
+using Diffractor: ∂☆
+using Diffractor: bundle
+
+mutable struct MDemo1
+    x::Float64
+end
+function double!(val::MDemo1)
+    val.x *= 2.0
+    return val
+end
+function wrap_and_double(x)
+    val = MDemo1(x)
+    double!(val)
+end
+∂☆{1}()(ZeroBundle{1}(wrap_and_double), TaylorBundle{1}(1.5, (1.0,)))

test/runtests.jl

@@ -18,6 +18,7 @@ const bwd = Diffractor.PrimeDerivativeBack
     "tangent.jl",
     "forward_diff_no_inf.jl",
     "forward.jl",
+    "forward_mutation.jl",
     "reverse.jl",
     "regression.jl",
     "AbstractDifferentiationTests.jl"

test/tangent.jl

@@ -122,4 +122,19 @@ end
     @test truncate(et, Val(1)) == TaylorTangent((1.0,))
 end
 
+@testset "zero_bundle" begin
+    zero_bundle = Diffractor.zero_bundle
+
+    tup_zb = zero_bundle{2}()((1, 0))
+    @test tup_zb isa ZeroBundle{2}
+    @test iszero(tup_zb[TaylorTangentIndex(1)])
+    @test iszero(tup_zb[TaylorTangentIndex(2)])
+
+
+    ref_zb = zero_bundle{2}()(Ref(1.5))
+    @test ref_zb isa TaylorBundle{2}
+    @test iszero(ref_zb[TaylorTangentIndex(1)])
+    @test iszero(ref_zb[TaylorTangentIndex(2)])
+end
+
 end  # module