Merge pull request #159 from oscardssmith/small-improvements

separate minor changes from #158

Author: oscardssmith

src/interface.jl

@@ -59,17 +59,14 @@ dx(x::Complex) = error("Tried to take the gradient of a complex-valued function.
 dx(x) = error("Cotangent space not defined for `$(typeof(x))`. Try a real-valued function.")
 
 """
-    ∂x(x)
+    ∂xⁿ{N}(x)
 
-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`
+For `x` in a one dimensional manifold, map x to the trivial, unital, Nth order
+tangent bundle. It should hold that `∀x ⟨∂ⁿ{1}x(x), dx(x)⟩ = 1`
 """
-∂x(x::Real) = ExplicitTangentBundle{1}(x, (one(x),))
-∂x(x) = error("Tangent space not defined for `$(typeof(x)).")
-
 struct ∂xⁿ{N}; end
 
-(::∂xⁿ{N})(x::Real) where {N} = TaylorBundle{N}(x, (one(x), (zero(x) for i = 1:(N-1))...,))
+(::∂xⁿ{N})(x::Real) where {N} = TaylorBundle{N}(x, ntuple(i->i==1 ? one(x) : zero(x), N))
 (::∂xⁿ)(x) = error("Tangent space not defined for `$(typeof(x)).")
 
 function ChainRules.rrule(∂::∂xⁿ, x)
@@ -173,11 +170,6 @@ raise_pd(f::PrimeDerivativeFwd{N,T}) where {N,T} = PrimeDerivativeFwd{N+1,T}(get
 
 (f::PrimeDerivativeFwd{0})(x) = getfield(f, :f)(x)
 
-function (f::PrimeDerivativeFwd{1})(x)
-    z = ∂☆¹(ZeroBundle{1}(getfield(f, :f)), ∂x(x))
-    z[TaylorTangentIndex(1)]
-end
-
 function (f::PrimeDerivativeFwd{N})(x) where N
     z = ∂☆{N}()(ZeroBundle{N}(getfield(f, :f)), ∂xⁿ{N}()(x))
     z[TaylorTangentIndex(N)]

src/tangent.jl

@@ -334,10 +334,6 @@ function unbundle(atb::TaylorBundle{Order, A}) where {Order, Dim, T, A<:Abstract
     StructArray{TaylorBundle{Order, T}}((atb.primal, atb.tangent.coeffs...))
 end
 
-function ChainRulesCore.rrule(::typeof(unbundle), atb::TaylorBundle)
-    unbundle(atb), Δ->throw(Δ)
-end
-
 function StructArrays.staticschema(::Type{<:TaylorBundle{N, B, T}}) where {N, B, T}
     Tuple{B, T.parameters...}
 end
@@ -355,11 +351,11 @@ function StructArrays.createinstance(T::Type{<:TaylorBundle}, args...)
     T(first(args), Base.tail(args))
 end
 
-function unbundle(zb::ZeroBundle{N, A}) where {N,T,Dim,A<:AbstractArray{T, Dim}}
-    StructArray{ZeroBundle{N, T}}((zb.primal, fill(zb.tangent.val, size(zb.primal)...)))
+function unbundle(u::UniformBundle{N, A}) where {N,T,Dim,A<:AbstractArray{T, Dim}}
+    StructArray{UniformBundle{N, T}}((u.primal, fill(u.tangent.val, size(u.primal)...)))
 end
 
-function ChainRulesCore.rrule(::typeof(unbundle), atb::ZeroBundle)
+function ChainRulesCore.rrule(::typeof(unbundle), atb::AbstractTangentBundle)
     unbundle(atb), Δ->throw(Δ)
 end
 
@@ -383,6 +379,11 @@ function rebundle(A::AbstractArray{<:TaylorBundle{N}}) where {N}
         end)
 end
 
+function rebundle(A::AbstractArray{<:UniformBundle{N}}) where {N}
+    @assert all(x->getfield(x, :tangent)==(first(A).tangent), A)
+    UniformBundle{N}(map(x->x.primal, A), first(A).tangent.val)
+end
+
 function ChainRulesCore.rrule(::typeof(rebundle), atb)
     rebundle(atb), Δ->throw(Δ)
 end