reshape to automatically handle smaller sizes

Author: ChrisRackauckas

src/PreallocationTools.jl

@@ -23,17 +23,34 @@ forward-mode automatic differentiation.
 """
 dualcache(u::AbstractArray, N=Val{ForwardDiff.pickchunksize(length(u))}) = DiffCache(u, size(u), N)
 
+chunksize(::Type{ForwardDiff.Dual{T,V,N}}) where {T,V,N} = N
+
 function get_tmp(dc::DiffCache, u::T) where T<:ForwardDiff.Dual
   x = reinterpret(T, dc.dual_du)
+  if chunksize(T) === chunksize(eltype(dc.dual_du))
+      x
+  else
+      @view x[axes(dc.du)...]
+  end
 end
 
 function get_tmp(dc::DiffCache, u::AbstractArray{T}) where T<:ForwardDiff.Dual
   x = reinterpret(T, dc.dual_du)
+  if chunksize(T) === chunksize(eltype(dc.dual_du))
+      x
+  else
+      @view x[axes(dc.du)...]
+  end
 end
 
 function get_tmp(dc::DiffCache, u::LabelledArrays.LArray{T,N,D,Syms}) where {T,N,D,Syms}
   x = reinterpret(T, dc.dual_du.__x)
-  LabelledArrays.LArray{T,N,D,Syms}(x)
+  _x = if chunksize(T) === chunksize(eltype(dc.dual_du))
+      x
+  else
+      @view x[axes(dc.du)...]
+  end
+  LabelledArrays.LArray{T,N,D,Syms}(_x)
 end
 
 get_tmp(dc::DiffCache, u::Number) = dc.du

test/runtests.jl

@@ -28,3 +28,25 @@ function foo(du, u, (A, lbc), t)
 end
 prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), LazyBufferCache()))
 solve(prob, TRBDF2())
+
+## Check resizing
+
+randmat = rand(10, 2)
+sto = similar(randmat)
+stod = dualcache(sto)
+
+function claytonsample!(sto, τ; randmat=randmat)
+    sto = get_tmp(sto, τ)
+    sto .= randmat
+    τ == 0 && return sto
+
+    n = size(sto, 1)
+    for i in 1:n
+        v = sto[i, 2]
+        u = sto[i, 1]
+        sto[i, 2] = (1 - u^(-τ) + u^(-τ)*v^(-(τ/(1 + τ))))^(-1/τ)
+    end
+    return sto
+end
+
+ForwardDiff.derivative(τ -> claytonsample!(stod, τ), 0.3)