prep for fearless (Taylor's Version)

Author: oscardssmith

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) = ExplicitTangentBundle{1}(x, (one(x),))
+∂x(x::Real) = TaylorBundle{1}(x, (one(x),))
 ∂x(x) = error("Tangent space not defined for `$(typeof(x)).")
 
 struct ∂xⁿ{N}; end

src/jet.jl

@@ -1,5 +1,5 @@
 """
-    struct Jet{T, N}
+    struct Jet{S, T, N}
 
 Represents the truncated (N-1)-th order Taylor series
 
@@ -15,8 +15,8 @@ For a jet `j`, several operations are supported:
    derivatives. Mathematically this corresponds to an infinitessimal ball
    around `a`.
 """
-struct Jet{T, N}
-    a::T
+struct Jet{S, T, N}
+    a::S
     f₀::T
     fₙ::NTuple{N, T}
 end
@@ -25,13 +25,13 @@ function ChainRulesCore.rrule(::typeof(Base.getproperty), j::Jet, s)
     error("Raw getproperty not allowed in AD code")
 end
 
-function Base.:+(j1::Jet{T, N}, j2::Jet{T, N}) where {T, N}
+function Base.:+(j1::Jet{S, T, N}, j2::Jet{S, T, N}) where {S, T, N}
     @assert j1.a === j2.a
-    Jet{T, N}(j1.a, j1.f₀ + j2.f₀, map(+, j1.fₙ, j2.fₙ))
+    Jet{S, T, N}(j1.a, j1.f₀ + j2.f₀, map(+, j1.fₙ, j2.fₙ))
 end
 
-function Base.:+(j::Jet{T, N}, x::T) where {T, N}
-    Jet{T, N}(j.a, j.f₀+x, j.fₙ)
+function Base.:+(j::Jet{S, T, N}, x::T) where {S, T, N}
+    Jet{S, T, N}(j.a, j.f₀+x, j.fₙ)
 end
 
 struct One; end
@@ -44,28 +44,28 @@ function ChainRulesCore.rrule(::typeof(+), j::Jet, x::One)
     j + x, Δ->(NoTangent(), One(), ZeroTangent())
 end
 
-function Base.zero(j::Jet{T, N}) where {T, N}
+function Base.zero(j::Jet{S, T, N}) where {S, T, N}
     let z = zero(j[0])
-        Jet{T, N}(j.a, z,
+        Jet{S, T, N}(j.a, z,
             ntuple(_->z, N))
     end
 end
 function ChainRulesCore.rrule(::typeof(Base.zero), j::Jet)
     zero(j), Δ->(NoTangent(), ZeroTangent())
 end
 
-function Base.getindex(j::Jet{T, N}, i::Integer) where {T, N}
+function Base.getindex(j::Jet{S, T, N}, i::Integer) where {S, T, N}
     (0 <= i <= N) || throw(BoundsError(j, i))
     i == 0 && return j.f₀
     return j.fₙ[i]
 end
 
-function deriv(j::Jet{T, N}) where {T, N}
-    Jet{T, N-1}(j.a, j.fₙ[1], Base.tail(j.fₙ))
+function deriv(j::Jet{S, T, N}) where {S, T, N}
+    Jet{S, T, N-1}(j.a, j.fₙ[1], Base.tail(j.fₙ))
 end
 
-function integrate(j::Jet{T, N}) where {T, N}
-    Jet{T, N+1}(j.a, zero(j.f₀), tuple(j.f₀, j.fₙ...))
+function integrate(j::Jet{S, T, N}) where {S, T, N}
+    Jet{S, T, N+1}(j.a, zero(j.f₀), tuple(j.f₀, j.fₙ...))
 end
 
 deriv(::NoTangent) = NoTangent()
@@ -188,7 +188,7 @@ function (∂⃖ₙ::∂⃖{N})(::typeof(map), f, a::Array) where {N}
                      TaylorBundle{N}(x,
                        (one(x), (zero(x) for i = 1:(N-1))...,)))
         @assert isa(∂f, TaylorBundle) || isa(∂f, ExplicitTangentBundle{1})
-        Jet{typeof(x), N}(x, ∂f.primal,
+        Jet{typeof(x), typeof(x), N}(x, ∂f.primal,
             isa(∂f, ExplicitTangentBundle) ? ∂f.tangent.partials : ∂f.tangent.coeffs)
     end
     ∂⃖ₙ(mapev, js, a)
@@ -239,7 +239,7 @@ expressions for the t′ᵢ that are hopefully easier on the compiler.
     end...)
 end
 
-@generated function (j::Jet{T, N} where T)(x::TaylorBundle{M}) where {N, M}
+@generated function (j::Jet{S, T, N} where {S, T})(x::TaylorBundle{M}) where {N, M}
     O = min(M,N)
     quote
         domain_check(j, x.primal)
@@ -249,12 +249,12 @@ end
     end
 end
 
-function (j::Jet{T, 1} where T)(x::ExplicitTangentBundle{1})
+function (j::Jet{S, T, 1} where {S,T})(x::ExplicitTangentBundle{1})
     domain_check(j, x.primal)
     coeffs = x.tangent.partials
     ExplicitTangentBundle{1}(j[0], (jet_taylor_ev(Val{1}(), coeffs, j),))
 end
 
-function (j::Jet{T, N} where T)(x::ExplicitTangentBundle{N, M}) where {N, M}
+function (j::Jet{S, T, N} where T)(x::ExplicitTangentBundle{N, M}) where {S, N, M}
     error("TODO")
 end

src/stage1/forward.jl

@@ -36,9 +36,9 @@ end
 
 function shuffle_down(b::TaylorBundle{N, B}) where {N, B}
     TaylorBundle{N-1}(
-        ExplicitTangentBundle{1}(b.primal, (b.tangent.coeffs[1],)),
+        TaylorBundle{1}(b.primal, (b.tangent.coeffs[1],)),
         ntuple(N-1) do i
-            ExplicitTangentBundle{1}(b.tangent.coeffs[i], (b.tangent.coeffs[i+1],))
+            TaylorBundle{1}(b.tangent.coeffs[i], (b.tangent.coeffs[i+1],))
         end)
 end
 
@@ -54,13 +54,8 @@ end
 function shuffle_up(r::CompositeBundle{1})
     z₀ = primal(r.tup[1])
     z₁ = partial(r.tup[1], 1)
-    z₂ = primal(r.tup[2])
     z₁₂ = partial(r.tup[2], 1)
-    if z₁ == z₂
-        return TaylorBundle{2}(z₀, (z₁, z₁₂))
-    else
-        return ExplicitTangentBundle{2}(z₀, (z₁, z₂, z₁₂))
-    end
+    return TaylorBundle{2}(z₀, (z₁, z₁₂))
 end
 
 function taylor_compatible(a::ATB{N}, b::ATB{N}) where {N}
@@ -118,7 +113,7 @@ function (::∂☆internal{1})(args::AbstractTangentBundle{1}...)
 end
 
 function ChainRulesCore.frule_via_ad(::DiffractorRuleConfig, partials, args...)
-    bundles = map((p,a) -> ExplicitTangentBundle{1}(a, (p,)), partials, args)
+    bundles = map((p,a) -> TaylorBundle{1}(a, (p,)), partials, args)
     result = ∂☆internal{1}()(bundles...)
     primal(result), first_partial(result)
 end
@@ -139,13 +134,13 @@ end
 (::∂☆{N})(args::AbstractTangentBundle{N}...) where {N} = ∂☆internal{N}()(args...)
 
 # Special case rules for performance
-@Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::TangentBundle{N}, s::AbstractTangentBundle{N}) where {N}
+@Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::ExplicitTangentBundle{N}, s::AbstractTangentBundle{N}) where {N}
     s = primal(s)
     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::TangentBundle{N}, s::ATB{N}, inbounds::ATB{N}) where {N}
+@Base.constprop :aggressive function (::∂☆{N})(f::ATB{N, typeof(getfield)}, x::ExplicitTangentBundle{N}, s::ATB{N}, inbounds::ATB{N}) where {N}
     s = primal(s)
     ExplicitTangentBundle{N}(getfield(primal(x), s, primal(inbounds)),
         map(x->lifted_getfield(x, s), x.tangent.partials))

src/tangent.jl

@@ -90,9 +90,8 @@ struct ExplicitTangent{P <: Tuple} <: AbstractTangentSpace
     partials::P
 end
 
-@eval struct TaylorTangent{C <: Tuple} <: AbstractTangentSpace
+struct TaylorTangent{C <: Tuple} <: AbstractTangentSpace
     coeffs::C
-    TaylorTangent(coeffs) = $(Expr(:new, :(TaylorTangent{typeof(coeffs)}), :coeffs))
 end
 
 """
@@ -140,24 +139,17 @@ struct UniformTangent{U} <: AbstractTangentSpace
     val::U
 end
 
-function _TangentBundle end
-
-@eval struct TangentBundle{N, B, P <: AbstractTangentSpace} <: AbstractTangentBundle{N, B}
-    primal::B
-    tangent::P
-    global _TangentBundle(::Val{N}, primal::B, tangent::P) where {N, B, P} = $(Expr(:new, :(TangentBundle{N, Core.Typeof(primal), typeof(tangent)}), :primal, :tangent))
-end
-
 """
     struct TangentBundle{N, B, P}
 
 Represents a tangent bundle as an explicit primal together
 with some representation of (potentially a product of) the tangent space.
 """
-TangentBundle
-
-TangentBundle{N}(primal::B, tangent::P) where {N, B, P<:AbstractTangentSpace} =
-    _TangentBundle(Val{N}(), primal, tangent)
+struct TangentBundle{N, B, P <: AbstractTangentSpace} <: AbstractTangentBundle{N, B}
+    primal::B
+    tangent::P
+    TangentBundle{N}(B, P) where {N} = new{N, typeof(B), typeof(P)}(B,P)
+end
 
 const ExplicitTangentBundle{N, B, P} = TangentBundle{N, B, ExplicitTangent{P}}
 
@@ -166,17 +158,17 @@ check_tangent_invariant(lp, N) = @assert lp == 2^N - 1
 
 function ExplicitTangentBundle{N}(primal::B, partials::P) where {N, B, P}
     check_tangent_invariant(length(partials), N)
-    _TangentBundle(Val{N}(), primal, ExplicitTangent{P}(partials))
+    TangentBundle{N}(primal, ExplicitTangent{P}(partials))
 end
 
 function ExplicitTangentBundle{N,B}(primal::B, partials::P) where {N, B, P}
     check_tangent_invariant(length(partials), N)
-    _TangentBundle(Val{N}(), primal, ExplicitTangent{P}(partials))
+    TangentBundle{N}(primal, ExplicitTangent{P}(partials))
 end
 
 function ExplicitTangentBundle{N,B,P}(primal::B, partials::P) where {N, B, P}
     check_tangent_invariant(length(partials), N)
-    _TangentBundle(Val{N}(), primal, ExplicitTangent{P}(partials))
+    TangentBundle{N}(primal, ExplicitTangent{P}(partials))
 end
 
 function Base.show(io::IO, x::ExplicitTangentBundle)
@@ -203,7 +195,7 @@ const TaylorBundle{N, B, P} = TangentBundle{N, B, TaylorTangent{P}}
 
 function TaylorBundle{N, B}(primal::B, coeffs) where {N, B}
     check_taylor_invariants(coeffs, primal, N)
-    _TangentBundle(Val{N}(), primal, TaylorTangent(coeffs))
+    TangentBundle{N}(primal, TaylorTangent(coeffs))
 end
 
 function check_taylor_invariants(coeffs, primal, N)
@@ -215,7 +207,7 @@ end
 @ChainRulesCore.non_differentiable check_taylor_invariants(coeffs, primal, N)
 
 function TaylorBundle{N}(primal, coeffs) where {N}
-    _TangentBundle(Val{N}(), primal, TaylorTangent(coeffs))
+    TangentBundle{N}(primal, TaylorTangent(coeffs))
 end
 
 function Base.show(io::IO, x::TaylorBundle{1})
@@ -230,25 +222,13 @@ function Base.getindex(tb::TaylorBundle, tti::CanonicalTangentIndex)
     tb.tangent.coeffs[count_ones(tti.i)]
 end
 
-function truncate(tt::TaylorTangent, order::Val{N}) where {N}
-    TaylorTangent(tt.coeffs[1:N])
-end
-
-function truncate(ut::UniformTangent, order::Val)
-    ut
-end
-
-function truncate(tb::TangentBundle, order::Val)
-    _TangentBundle(order, tb.primal, truncate(tb.tangent, order))
-end
-
 const UniformBundle{N, B, U} = TangentBundle{N, B, UniformTangent{U}}
-UniformBundle{N, B, U}(primal::B, partial::U) where {N,B,U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(partial))
-UniformBundle{N, B, U}(primal::B) where {N,B,U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(U.instance))
-UniformBundle{N, B}(primal::B, partial::U) where {N,B,U} = _TangentBundle(Val{N}(),primal, UniformTangent{U}(partial))
-UniformBundle{N}(primal, partial::U) where {N,U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(partial))
-UniformBundle{N, <:Any, U}(primal, partial::U) where {N, U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(U.instance))
-UniformBundle{N, <:Any, U}(primal) where {N, U} = _TangentBundle(Val{N}(), primal, UniformTangent{U}(U.instance))
+UniformBundle{N, B, U}(primal::B, partial::U) where {N,B,U} = TangentBundle{N}(primal, UniformTangent{U}(partial))
+UniformBundle{N, B, U}(primal::B) where {N,B,U} = TangentBundle{N}(primal, UniformTangent{U}(U.instance))
+UniformBundle{N, B}(primal::B, partial::U) where {N,B,U} = TangentBundle{N}(primal, UniformTangent{U}(partial))
+UniformBundle{N}(primal, partial::U) where {N,U} = TangentBundle{N}(primal, UniformTangent{U}(partial))
+UniformBundle{N, <:Any, U}(primal, partial::U) where {N, U} = TangentBundle{N}(primal, UniformTangent{U}(U.instance))
+UniformBundle{N, <:Any, U}(primal) where {N, U} = TangentBundle{N}(primal, UniformTangent{U}(U.instance))
 
 const ZeroBundle{N, B} = UniformBundle{N, B, ZeroTangent}
 const DNEBundle{N, B} = UniformBundle{N, B, NoTangent}

test/runtests.jl

@@ -50,7 +50,7 @@ ChainRules.rrule(::typeof(my_tuple), args...) = args, Δ->Core.tuple(NoTangent()
 
 # Minimal 2-nd order forward smoke test
 @test Diffractor.∂☆{2}()(Diffractor.ZeroBundle{2}(sin),
-    Diffractor.ExplicitTangentBundle{2}(1.0, (1.0, 1.0, 0.0)))[Diffractor.CanonicalTangentIndex(1)] == sin'(1.0)
+    Diffractor.TaylorBundle{2}(1.0, (1.0 0.0)))[Diffractor.CanonicalTangentIndex(1)] == sin'(1.0)
 
 function simple_control_flow(b, x)
     if b