added kwarg levels for standard chunking

Author: thomvet

src/PreallocationTools.jl

@@ -7,27 +7,23 @@ struct DiffCache{T<:AbstractArray, S<:AbstractArray}
     dual_du::S
 end
 
-function DiffCache(u::AbstractArray{T}, siz, chunk_size::Int) where {T}
-    x = adapt(ArrayInterface.parameterless_type(u), zeros(T, (chunk_size+1)*prod(siz)))
-    DiffCache(u, x)
-end
-
-function DiffCache(u::AbstractArray{T}, siz, chunk_sizes::AbstractArray{V}) where {T,V<:Int}
-    clamp!(chunk_sizes,1,ForwardDiff.DEFAULT_CHUNK_THRESHOLD) 
+function DiffCache(u::AbstractArray{T}, siz, chunk_sizes) where {T}
     x = adapt(ArrayInterface.parameterless_type(u), zeros(T, prod(chunk_sizes .+ 1)*prod(siz)))
     DiffCache(u, x)
 end
 
 """
 
-`dualcache(u::AbstractArray, N = default_cache_size(length(u)))`
+`dualcache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u)); levels::Int = 1)`
+`dualcache(u::AbstractArray; N::AbstractArray{<:Int})`
 
 Builds a `DualCache` object that stores both a version of the cache for `u`
 and for the `Dual` version of `u`, allowing use of pre-cached vectors with
-forward-mode automatic differentiation.
+forward-mode automatic differentiation. Supports nested AD.
 
 """
-dualcache(u::AbstractArray, N=ForwardDiff.pickchunksize(length(u))) = DiffCache(u, size(u), N)
+dualcache(u::AbstractArray, N::Int=ForwardDiff.pickchunksize(length(u)); levels::Int = 1) = DiffCache(u, size(u), N^levels)
+dualcache(u::AbstractArray, N::AbstractArray{<:Int}) = DiffCache(u, size(u), N)
 
 """
 

test/core_nesteddual.jl

@@ -42,9 +42,9 @@ function foo(du, u, p, t)
     nothing
 end
 
-ps = 2 #use to specify problem size (ps ∈ {1,2})
-coeffs = rand(ps^2)
-cache = dualcache(zeros(ps,ps), [4, 4, 4])
+ps = 3 #use to specify problem size; don't go crazy on this, because compilation time...
+coeffs = -rand(ps^2)
+cache = dualcache(zeros(ps,ps), [9, 9, 9])
 prob = ODEProblem(foo, ones(ps, ps), (0., 1.0), (coeffs, cache))
 realsol = solve(prob, TRBDF2(), saveat = 0.0:0.01:1.0, reltol = 1e-8)
 
@@ -66,7 +66,7 @@ fn(x,p) = objfun(x, p[1], p[2], p[3])
 optfun = OptimizationFunction(fn, GalacticOptim.AutoForwardDiff())
 optprob = OptimizationProblem(optfun, rand(size(coeffs)...), (prob, realsol, cache))
 newtonsol = solve(optprob, Newton())
-bfgssol = solve(optprob, BFGS()) #since only gradients are used here, we could go with a slim dualcache(zeros(ps,ps), [4,4]) as well.
+bfgssol = solve(optprob, BFGS()) #since only gradients are used here, we could go with a smaller dualcache(zeros(ps,ps), [9,9]) as well.
 
 @test all(abs.(coeffs .- newtonsol.u) .< 1e-3)
 @test all(abs.(coeffs .- bfgssol.u) .< 1e-3)
@@ -81,7 +81,6 @@ function foo(du, u, p, t)
     nothing
 end
 
-ps = 2 #use to specify problem size (ps ∈ {1,2})
 coeffs = rand(1)
 cache = dualcache(zeros(ps,ps), [1, 1, 4])
 prob = ODEProblem(foo, ones(ps, ps), (0., 1.0), (coeffs, cache))