WIP: Refactor forward mode data structures (#54)

To allow chunking like ForwardDiff.

import `∂⃖¹`

some fixes

Author: Keno

src/interface.jl

@@ -64,7 +64,7 @@ dx(x) = error("Cotangent space not defined for `$(typeof(x))`. Try a real-valued
 For `x` in a one dimensional manifold, map x to the trivial, unital, 1st order
 tangent bundle. It should hold that `∀x ⟨∂x(x), dx(x)⟩ = 1`
 """
-∂x(x::Real) = TangentBundle{1}(x, (one(x),))
+∂x(x::Real) = ExplicitTangentBundle{1}(x, (one(x),))
 ∂x(x) = error("Tangent space not defined for `$(typeof(x)).")
 
 struct ∂xⁿ{N}; end
@@ -177,7 +177,7 @@ raise_pd(f::PrimeDerivativeFwd{N,T}) where {N,T} = PrimeDerivativeFwd{plus1(N),T
 
 function (f::PrimeDerivativeFwd{1})(x)
     z = ∂☆¹(ZeroBundle{1}(getfield(f, :f)), ∂x(x))
-    z.partials[1]
+    z.tangent.partials[1]
 end
 
 function (f::PrimeDerivativeFwd{N})(x) where N

src/jet.jl

@@ -187,9 +187,9 @@ function (∂⃖ₙ::∂⃖{N})(::typeof(map), f, a::Array) where {N}
         ∂f = ∂☆{N}()(ZeroBundle{N}(f),
                      TaylorBundle{N}(x,
                        (one(x), (zero(x) for i = 1:(N-1))...,)))
-        @assert isa(∂f, TaylorBundle) || isa(∂f, TangentBundle{1})
+        @assert isa(∂f, TaylorBundle) || isa(∂f, ExplicitTangentBundle{1})
         Jet{typeof(x), N}(x, ∂f.primal,
-            isa(∂f, TangentBundle) ? ∂f.partials : ∂f.coeffs)
+            isa(∂f, ExplicitTangentBundle) ? ∂f.tangent.partials : ∂f.tangent.coeffs)
     end
     ∂⃖ₙ(mapev, js, a)
 end
@@ -243,18 +243,18 @@ end
     O = min(M,N)
     quote
         domain_check(j, x.primal)
-        coeffs = x.coeffs
+        coeffs = x.tangent.coeffs
         TaylorBundle{$O}(j[0],
             ($((:(jet_taylor_ev(Val{$i}(), coeffs, j)) for i = 1:O)...),))
     end
 end
 
-function (j::Jet{T, 1} where T)(x::TangentBundle{1})
+function (j::Jet{T, 1} where T)(x::ExplicitTangentBundle{1})
     domain_check(j, x.primal)
-    coeffs = x.partials
-    TangentBundle{1}(j[0], (jet_taylor_ev(Val{1}(), coeffs, j),))
+    coeffs = x.tangent.partials
+    ExplicitTangentBundle{1}(j[0], (jet_taylor_ev(Val{1}(), coeffs, j),))
 end
 
-function (j::Jet{T, N} where T)(x::TangentBundle{N, M}) where {N, M}
+function (j::Jet{T, N} where T)(x::ExplicitTangentBundle{N, M}) where {N, M}
     error("TODO")
 end

src/stage1/forward.jl

@@ -1,15 +1,14 @@
-partial(x::TangentBundle, i) = x.partials[i]
-partial(x::TaylorBundle{1}, i) = x.coeffs[i]
-partial(x::UniformBundle, i) = x.partial
-partial(x::CompositeBundle{N, B}, i) where {N, B} = Tangent{B}(map(x->partial(x, i), x.tup)...)
-partial(x::ZeroTangent, i) = ZeroTangent()
+partial(x::TangentBundle, i) = partial(getfield(x, :tangent), i)
+partial(x::ExplicitTangent, i) = getfield(getfield(x, :partials), i)
+partial(x::TaylorTangent, i) = getfield(getfield(x, :coeffs), i)
+partial(x::UniformTangent, i) = getfield(x, :val)
+partial(x::ProductTangent, i) = ProductTangent(map(x->partial(x, i), getfield(x, :factors)))
+partial(x::AbstractZero, i) = x
+partial(x::CompositeBundle{N, B}, i) where {N, B} = Tangent{B}(map(x->partial(x, i), getfield(x, :tup))...)
 primal(x::AbstractTangentBundle) = x.primal
 primal(z::ZeroTangent) = ZeroTangent()
 
-first_partial(x::TangentBundle{1}) = getfield(getfield(x, :partials), 1)
-first_partial(x::TaylorBundle{1}) = getfield(getfield(x, :coeffs), 1)
-first_partial(x::UniformBundle) = getfield(x, :partial)
-first_partial(x::CompositeBundle) = map(first_partial, getfield(x, :tup))
+first_partial(x) = partial(x, 1)
 
 # TODO: Which version do we want in ChainRules?
 function my_frule(args::ATB{1}...)
@@ -24,22 +23,22 @@ my_frule(::ZeroBundle{1, typeof(my_frule)}, args::ATB{1}...) = nothing
 (::∂☆{N})(::ZeroBundle{N, typeof(my_frule)}, ::ZeroBundle{N, ZeroBundle{1, typeof(my_frule)}}, args::ATB{N}...) where {N} = ZeroBundle{N}(nothing)
 
 shuffle_down(b::UniformBundle{N, B, U}) where {N, B, U} =
-    UniformBundle{minus1(N), <:Any, U}(UniformBundle{1, B, U}(b.primal, b.partial), b.partial)
+    UniformBundle{minus1(N), <:Any, U}(UniformBundle{1, B, U}(b.primal, b.tangent.val), b.tangent.val)
 
-function shuffle_down(b::TangentBundle{N, B}) where {N, B}
+function shuffle_down(b::ExplicitTangentBundle{N, B}) where {N, B}
     # N.B: This depends on the special properties of the canonical tangent index order
-    TangentBundle{N-1}(
-        TangentBundle{1}(b.primal, (partial(b, 1),)),
+    ExplicitTangentBundle{N-1}(
+        ExplicitTangentBundle{1}(b.primal, (partial(b, 1),)),
         ntuple(2^(N-1)-1) do i
-            TangentBundle{1}(partial(b, 2*i), (partial(b, 2*i+1),))
+            ExplicitTangentBundle{1}(partial(b, 2*i), (partial(b, 2*i+1),))
         end)
 end
 
 function shuffle_down(b::TaylorBundle{N, B}) where {N, B}
     TaylorBundle{N-1}(
-        TangentBundle{1}(b.primal, (b.coeffs[1],)),
+        ExplicitTangentBundle{1}(b.primal, (b.tangent.coeffs[1],)),
         ntuple(N-1) do i
-            TangentBundle{1}(b.coeffs[i], (b.coeffs[i+1],))
+            ExplicitTangentBundle{1}(b.tangent.coeffs[i], (b.tangent.coeffs[i+1],))
         end)
 end
 
@@ -60,7 +59,7 @@ function shuffle_up(r::CompositeBundle{1})
     if z₁ == z₂
         return TaylorBundle{2}(z₀, (z₁, z₁₂))
     else
-        return TangentBundle{2}(z₀, (z₁, z₂, z₁₂))
+        return ExplicitTangentBundle{2}(z₀, (z₁, z₂, z₁₂))
     end
 end
 
@@ -86,14 +85,14 @@ function shuffle_up(r::CompositeBundle{N}) where {N}
             N+1))
     else
         return TangentBundle{N+1}(r.tup[1].primal,
-            (r.tup[1].partials..., primal(b),
+            (r.tup[1].tangent.partials..., primal(b),
             ntuple(i->partial(b,i), 2^(N+1)-1)...))
     end
 end
 
 function shuffle_up(r::UniformBundle{N, B, U}) where {N, B, U}
     (a, b) = primal(r)
-    if r.partial === b
+    if r.tangent.val === b
         u = b
     elseif b == NoTangent() && U === ZeroTangent
         u = b
@@ -107,7 +106,7 @@ end
 struct ∂☆internal{N}; end
 struct ∂☆shuffle{N}; end
 
-shuffle_base(r) = TangentBundle{1}(r[1], (r[2],))
+shuffle_base(r) = ExplicitTangentBundle{1}(r[1], (r[2],))
 
 function (::∂☆internal{1})(args::AbstractTangentBundle{1}...)
     r = my_frule(args...)
@@ -119,7 +118,7 @@ function (::∂☆internal{1})(args::AbstractTangentBundle{1}...)
 end
 
 function ChainRulesCore.frule_via_ad(::DiffractorRuleConfig, partials, args...)
-    bundles = map((p,a) -> TangentBundle{1}(a, (p,)), partials, args)
+    bundles = map((p,a) -> ExplicitTangentBundle{1}(a, (p,)), partials, args)
     result = ∂☆internal{1}()(bundles...)
     primal(result), first_partial(result)
 end
@@ -142,14 +141,14 @@ end
 # Special case rules for performance
 @Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::TangentBundle{N}, s::AbstractTangentBundle{N}) where {N}
     s = primal(s)
-    TangentBundle{N}(getfield(primal(x), s),
-        map(x->lifted_getfield(x, s), x.partials))
+    ExplicitTangentBundle{N}(getfield(primal(x), s),
+        map(x->lifted_getfield(x, s), x.tangent.partials))
 end
 
 @Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::TaylorBundle{N}, s::AbstractTangentBundle{N}) where {N}
     s = primal(s)
     TaylorBundle{N}(getfield(primal(x), s),
-        map(y->lifted_getfield(y, s), x.coeffs))
+        map(y->lifted_getfield(y, s), x.tangent.coeffs))
 end
 
 @Base.constprop :aggressive function (::∂☆{N})(::ATB{N, typeof(getfield)}, x::CompositeBundle{N}, s::AbstractTangentBundle{N, Int}) where {N}
@@ -162,16 +161,16 @@ end
 
 @Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::ATB{N}, s::ATB{N}, inbounds::ATB{N}) where {N}
     s = primal(s)
-    TangentBundle{N}(getfield(primal(x), s, primal(inbounds)),
-        map(x->lifted_getfield(x, s), x.partials))
+    ExplicitTangentBundle{N}(getfield(primal(x), s, primal(inbounds)),
+        map(x->lifted_getfield(x, s), x.tangent.partials))
 end
 
 @Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::UniformBundle{N, <:Any, U}, s::AbstractTangentBundle{N}) where {N, U}
-    UniformBundle{N,<:Any,U}(getfield(primal(x), primal(s)), x.partial)
+    UniformBundle{N,<:Any,U}(getfield(primal(x), primal(s)), x.tangent.val)
 end
 
 @Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::UniformBundle{N, <:Any, U}, s::AbstractTangentBundle{N}, inbounds::AbstractTangentBundle{N}) where {N, U}
-    UniformBundle{N,<:Any,U}(getfield(primal(x), primal(s), primal(inbounds)), x.partial)
+    UniformBundle{N,<:Any,U}(getfield(primal(x), primal(s), primal(inbounds)), x.tangent.val)
 end
 
 function (::∂☆{N})(f::ATB{N, typeof(tuple)}, args::AbstractTangentBundle{N}...) where {N}

src/stage1/mixed.jl

@@ -95,9 +95,9 @@ function (∂⃖ₙ::∂⃖{N})(∂☆ₘ::∂☆{M}, ::ZeroBundle{M, typeof(map
         ∂f = ∂☆{N+M}()(ZeroBundle{N+M}(primal(f)),
                      TaylorBundle{N+M}(x,
                        (one(x), (zero(x) for i = 1:(N+M-1))...,)))
-        @assert isa(∂f, TaylorBundle) || isa(∂f, TangentBundle{1})
+        @assert isa(∂f, TaylorBundle) || isa(∂f, ExplicitTangentBundle{1})
         Jet{typeof(x), N+M}(x, ∂f.primal,
-            isa(∂f, TangentBundle) ? ∂f.partials : ∂f.coeffs)
+            isa(∂f, ExplicitTangentBundle) ? ∂f.tangent.partials : ∂f.tangent.coeffs)
     end
     ∂⃖ₙ(mapev_unbundled, ∂☆ₘ, js, a)
 end