implement in-place tuple derivative function

Author: jrevels

src/config.jl

@@ -2,7 +2,21 @@
 # Tag #
 #######
 
-struct Tag{F,M} end
+struct Tag{F,H} end
+
+# Here, we could've just as easily used `hash`; however, this
+# is unsafe/undefined behavior if `hash(::Type{V})` is overloaded
+# in a module loaded after ForwardDiff. Thus, we instead use
+# `hash(Symbol(V))`, which is somewhat safer since it's far less
+# likely that somebody would overwrite the Base definition for
+# `Symbol(::DataType)` or `hash(::Symbol)`.
+@generated function Tag(::Type{F}, ::Type{V}) where {F,V}
+    H = hash(Symbol(V))
+    return quote
+        $(Expr(:meta, :inline))
+        Tag{F,$H}()
+    end
+end
 
 #########
 # Chunk #
@@ -37,9 +51,9 @@ end
 
 abstract type AbstractConfig{T<:Tag,N} end
 
-struct ConfigMismatchError{F,G,M} <: Exception
+struct ConfigMismatchError{F,G,H} <: Exception
     f::F
-    cfg::AbstractConfig{Tag{G,M}}
+    cfg::AbstractConfig{Tag{G,H}}
 end
 
 function Base.showerror{F,G}(io::IO, e::ConfigMismatchError{F,G})
@@ -67,7 +81,7 @@ end
 function GradientConfig{V,N,F,T}(::F,
                                  x::AbstractArray{V},
                                  ::Chunk{N} = Chunk(x),
-                                 ::T = Tag{F,order(V)}())
+                                 ::T = Tag(F, V))
     seeds = construct_seeds(Partials{N,V})
     duals = similar(x, Dual{T,V,N})
     return GradientConfig{T,V,N,typeof(duals)}(seeds, duals)
@@ -85,7 +99,7 @@ end
 function JacobianConfig{V,N,F,T}(::F,
                                  x::AbstractArray{V},
                                  ::Chunk{N} = Chunk(x),
-                                 ::T = Tag{F,order(V)}())
+                                 ::T = Tag(F, V))
     seeds = construct_seeds(Partials{N,V})
     duals = similar(x, Dual{T,V,N})
     return JacobianConfig{T,V,N,typeof(duals)}(seeds, duals)
@@ -95,7 +109,7 @@ function JacobianConfig{Y,X,N,F,T}(::F,
                                    y::AbstractArray{Y},
                                    x::AbstractArray{X},
                                    ::Chunk{N} = Chunk(x),
-                                   ::T = Tag{F,order(X)}())
+                                   ::T = Tag(F, X))
     seeds = construct_seeds(Partials{N,X})
     yduals = similar(y, Dual{T,Y,N})
     xduals = similar(x, Dual{T,X,N})
@@ -107,15 +121,15 @@ end
 # HessianConfig #
 #################
 
-struct HessianConfig{T,V,N,D,MJ,DJ} <: AbstractConfig{T,N}
-    jacobian_config::JacobianConfig{Tag{Void,MJ},V,N,DJ}
-    gradient_config::GradientConfig{T,Dual{Tag{Void,MJ},V,N},D}
+struct HessianConfig{T,V,N,D,H,DJ} <: AbstractConfig{T,N}
+    jacobian_config::JacobianConfig{Tag{Void,H},V,N,DJ}
+    gradient_config::GradientConfig{T,Dual{Tag{Void,H},V,N},D}
 end
 
 function HessianConfig{F,V}(f::F,
                             x::AbstractArray{V},
                             chunk::Chunk = Chunk(x),
-                            tag::Tag = Tag{F,order(Dual{Void,V,0})}())
+                            tag::Tag = Tag(F, Dual{Void,V,0}))
     jacobian_config = JacobianConfig(nothing, x, chunk)
     gradient_config = GradientConfig(f, jacobian_config.duals, chunk, tag)
     return HessianConfig(jacobian_config, gradient_config)
@@ -125,7 +139,7 @@ function HessianConfig{F,V}(result::DiffResult,
                             f::F,
                             x::AbstractArray{V},
                             chunk::Chunk = Chunk(x),
-                            tag::Tag = Tag{F,order(Dual{Void,V,0})}())
+                            tag::Tag = Tag(F, Dual{Void,V,0}))
     jacobian_config = JacobianConfig(nothing, DiffBase.gradient(result), x, chunk)
     gradient_config = GradientConfig(f, jacobian_config.duals[2], chunk, tag)
     return HessianConfig(jacobian_config, gradient_config)

src/derivative.jl

@@ -2,37 +2,43 @@
 # API methods #
 ###############
 
-@generated function derivative{F,R<:Real}(f::F, x::R)
-    T = Tag{F,order(R)}
-    return quote
-        $(Expr(:meta, :inline))
-        return extract_derivative(f(Dual{$T}(x, one(x))))
-    end
+@inline function derivative(f::F, x::R) where {F,R<:Real}
+    T = Tag(F, R)
+    return extract_derivative(f(Dual{T}(x, one(x))))
 end
 
-@generated function derivative{F,N}(f::F, x::NTuple{N,Real})
-    T = Tag{F,maximum(order(R) for R in x.parameters)}
-    args = [:(Dual{$T}(x[$i], Val{N}, Val{$i})) for i in 1:N]
+@generated function derivative(f::F, x::NTuple{N,Real}) where {F,N}
+    args = [:(Dual{T}(x[$i], Val{N}, Val{$i})) for i in 1:N]
     return quote
         $(Expr(:meta, :inline))
+        T = Tag(F, typeof(x))
         extract_derivative(f($(args...)))
     end
 end
 
-@generated function derivative!{F,R<:Real}(out, f::F, x::R)
-    T = Tag{F,order(R)}
+@inline function derivative!(out, f::F, x::R) where {F,R<:Real}
+    T = Tag(F, typeof(x))
+    extract_derivative!(out, f(Dual{T}(x, one(x))))
+    return out
+end
+
+@generated function derivative!(out::NTuple{N,Any}, f::F, x::NTuple{N,Real}) where {F,N}
+    args = [:(Dual{T}(x[$i], Val{N}, Val{$i})) for i in 1:N]
     return quote
         $(Expr(:meta, :inline))
-        extract_derivative!(out, f(Dual{$T}(x, one(x))))
-        return out
+        T = Tag(F, typeof(x))
+        extract_derivative!(out, f($(args...)))
     end
 end
 
 #####################
 # result extraction #
 #####################
 
-@generated function extract_derivative{T,V,N}(y::Dual{T,V,N})
+# non-mutating #
+#--------------#
+
+@generated function extract_derivative(y::Dual{T,V,N}) where {T,V,N}
     return quote
         $(Expr(:meta, :inline))
         $(Expr(:tuple, [:(partials(y, $i)) for i in 1:N]...))
@@ -43,10 +49,36 @@ end
 @inline extract_derivative(y::Real) = zero(y)
 @inline extract_derivative(y::AbstractArray) = extract_derivative!(similar(y, valtype(eltype(y))), y)
 
+# mutating #
+#----------#
+
+@generated function extract_derivative!(out::NTuple{N,Any}, y::Dual{T,V,N}) where {T,V,N}
+    return quote
+        $(Expr(:meta, :inline))
+        $(Expr(:block, [:(out[$i][] = partials(y, $i)) for i in 1:N]...))
+        return out
+    end
+end
+
+@generated function extract_derivative!(out::NTuple{N,Any}, y::AbstractArray) where {N}
+    return quote
+        $(Expr(:meta, :inline))
+        $(Expr(:block, [:(extract_derivative!(out[$i], y, $i)) for i in 1:N]...))
+        return out
+    end
+end
+
 extract_derivative!(out::AbstractArray, y::AbstractArray) = map!(extract_derivative, out, y)
+extract_derivative!(out::AbstractArray, y::AbstractArray, p) = map!(x -> partials(x, p), out, y)
 
 function extract_derivative!(out::DiffResult, y)
     DiffBase.value!(value, out, y)
     DiffBase.derivative!(extract_derivative, out, y)
     return out
 end
+
+function extract_derivative!(out::DiffResult, y::AbstractArray, p)
+    DiffBase.value!(value, out, y)
+    DiffBase.derivative!(x -> partials(x, p), out, y)
+    return out
+end