fix all Dual definitions

Author: ChrisRackauckas

src/differentiation/jaches_products.jl

@@ -6,21 +6,22 @@ function auto_jacvec!(
     f,
     x,
     v,
-    cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(reshape(v, size(x)))),
+    cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x))))),
     cache2 = similar(cache1),
 )
-    cache1 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(reshape(v, size(x))))
+    cache1 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x)))))
     f(cache2, cache1)
     vecdy = _vec(dy)
-    vecdy .= partials.(vec(cache2), 1)
+    vecdy .= partials.(_vec(cache2), 1)
 end
 
 _vec(v) = vec(v)
 _vec(v::AbstractVector) = v
 
 function auto_jacvec(f, x, v)
     vv = reshape(v, axes(x))
-    vec(partials.(vec(f(ForwardDiff.Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(vv)))), 1))
+    y = ForwardDiff.Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(vv)))
+    vec(partials.(vec(f(y)), 1))
 end
 
 function num_jacvec!(
@@ -126,12 +127,12 @@ function autonum_hesvec!(
     f,
     x,
     v,
-    cache1 = ForwardDiff.Dual{DeivVecTag}.(x, v),
-    cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v),
+    cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x))))),
+    cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x))))),
 )
     cache = FiniteDiff.GradientCache(v[1], cache1, Val{:central})
     g = (dx, x) -> FiniteDiff.finite_difference_gradient!(dx, f, x, cache)
-    cache1 .= Dual{DeivVecTag}.(x, v)
+    cache1 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x)))))
     g(cache2, cache1)
     dy .= partials.(cache2, 1)
 end
@@ -168,16 +169,17 @@ function auto_hesvecgrad!(
     g,
     x,
     v,
-    cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v),
-    cache3 = ForwardDiff.Dual{DeivVecTag}.(x, v),
+    cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x))))),
+    cache3 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x))))),
 )
-    cache2 .= Dual{DeivVecTag}.(x, v)
+    cache2 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x)))))
     g(cache3, cache2)
     dy .= partials.(cache3, 1)
 end
 
 function auto_hesvecgrad(g, x, v)
-    partials.(g(Dual{DeivVecTag}.(x, v)), 1)
+    y = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x)))))
+    partials.(g(y), 1)
 end
 
 ### Operator Forms
@@ -192,8 +194,8 @@ end
 
 function JacVec(f, x::AbstractArray; autodiff = true)
     if autodiff
-        cache1 = ForwardDiff.Dual{DeivVecTag}.(x, x)
-        cache2 = ForwardDiff.Dual{DeivVecTag}.(x, x)
+        cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(x)))
+        cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(x)))
     else
         cache1 = similar(x)
         cache2 = similar(x)
@@ -261,8 +263,8 @@ end
 
 function HesVecGrad(g, x::AbstractArray; autodiff = false)
     if autodiff
-        cache1 = ForwardDiff.Dual{DeivVecTag}.(x, x)
-        cache2 = ForwardDiff.Dual{DeivVecTag}.(x, x)
+        cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(x)))
+        cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(x)))
     else
         cache1 = similar(x)
         cache2 = similar(x)

src/differentiation/jaches_products_zygote.jl

@@ -24,17 +24,19 @@ function numback_hesvec(f, x, v)
     (gxp - gxm)/(2ϵ)
 end
 
-function autoback_hesvec!(dy, f, x, v, cache2 = ForwardDiff.Dual{Nothing}.(x, v),
-                          cache3 = ForwardDiff.Dual{Nothing}.(x, v))
+function autoback_hesvec!(dy, f, x, v,
+                          cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x))))),
+                          cache3 = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x))))))
     g = let f=f
         (dx, x) -> dx .= first(Zygote.gradient(f,x))
     end
-    cache2 .= Dual{Nothing}.(x, v)
+    cache2 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x)))))
     g(cache3,cache2)
     dy .= partials.(cache3, 1)
 end
 
 function autoback_hesvec(f, x, v)
     g = x -> first(Zygote.gradient(f,x))
-    ForwardDiff.partials.(g(ForwardDiff.Dual{Nothing}.(x, v)), 1)
+    y = Dual{typeof(ForwardDiff.Tag(DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x)))))
+    ForwardDiff.partials.(g(y), 1)
 end

test/test_jaches_products.jl

@@ -18,8 +18,8 @@ function h(dy,x)
       FiniteDiff.finite_difference_gradient!(dy,g,x)
 end
 
-cache1 = ForwardDiff.Dual{SparseDiffTools.DeivVecTag}.(x, v)
-cache2 = ForwardDiff.Dual{SparseDiffTools.DeivVecTag}.(x, v)
+cache1 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(SparseDiffTools.DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(v)))
+cache2 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(SparseDiffTools.DeivVecTag,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(v)))
 @test num_jacvec!(dy, f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6
 @test num_jacvec!(dy, f, x, v, similar(v), similar(v)) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6
 @test num_jacvec(f, x, v) ≈ ForwardDiff.jacobian(f,similar(x),x)*v rtol=1e-6
@@ -44,8 +44,8 @@ cache2 = ForwardDiff.Dual{SparseDiffTools.DeivVecTag}.(x, v)
 @test numback_hesvec!(dy, g, x, v, similar(v), similar(v)) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
 @test numback_hesvec(g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
 
-cache3 = ForwardDiff.Dual{Nothing}.(x, v)
-cache4 = ForwardDiff.Dual{Nothing}.(x, v)
+cache3 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(Nothing,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(v)))
+cache4 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(Nothing,eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(Tuple.(v)))
 @test autoback_hesvec!(dy, g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
 @test autoback_hesvec!(dy, g, x, v, cache3, cache4) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8
 @test autoback_hesvec(g, x, v) ≈ ForwardDiff.hessian(g,x)*v rtol=1e-8