control on chunk_size in chunk_size

[JianghuiDu]

There doesn't seem to be an argument for chunk_size in JacVec or JacVecOperator. My code has some temporary variables inside so I use dualcache. But my choice of chunk_size in dualcache is not compatible with JacVec unless it is set to 1.

Here's an example using the example shown in the dualcache tutorial:

using LinearAlgebra, OrdinaryDiffEq, PreallocationTools
function foo(du, u, (A, tmp), t)
    tmp = get_tmp(tmp, u)
    mul!(tmp, A, u)
    @. du = u + tmp
    nothing
end
chunk_size = 1
# chunk_size = 5  doesn't work!

(A, tmp) = (ones(5,5), dualcache(zeros(5), Val{chunk_size}))
ff = (y,x)->foo(y,x,(A,tmp),nothing)

u0 = rand(5);
du0 = similar(u0);
ff(du0,u0)

using SparseDiffTools

J = JacVec(ff,u0;autodiff=true);
res = similar(u0);
mul!(res,J,u0) # Does 1 f evaluation

[ChrisRackauckas]

JacVec operations only ever need 1 chunk maximum, so it's kept to one.

[JianghuiDu]

But in dualcache we may need larger chunk size, does it mean that have to be limited to 1 too? Why should they be related at all?

[ChrisRackauckas]

That's a dualcache issue. We could just cut the size of the dualcache if it's generated for a larger chunk size. In theory, you could always use a smaller chunk. In practice, it doesn't give the right size so we should make a view.

[JianghuiDu]

Where should we add the view?

[ChrisRackauckas]

https://github.com/SciML/PreallocationTools.jl/blob/master/src/PreallocationTools.jl#L26-L35

If the chunk size of the reinterpret T is greater than the chunk size of the dual_du, then we should make a view. For example, if dual_du is a length 11 vector with chunk size 2 and T has chunk size one, then reinterpret will return a length 22 vector of chunk size 1 duals. We just need to make that be a view to the first 11 terms and it would work. And then we would just have to error if the chunk size of T is greater than the chunk size of dual_du.

[JianghuiDu]

You mean like

function get_tmp(dc::DiffCache, u::T) where T<:ForwardDiff.Dual
  x = reinterpret(T, dc.dual_du)

@view x[1:length(dc.dual_du)]
end

Also, how is the chunk size of T is controlled in this case? In the example I gave, I never specify the chunk size of u. Is it automatically chosen when ForwardDiff is called within JacVec? Can we specify that?

[ChrisRackauckas]

You mean like

Yup, exactly. If we throw that on each of those this case should be solved. The view needs to be multidimensional though to handle when x is a higher dimensional array. You'll want to splat axes(dc.dual_du) instead.

Also, how is the chunk size of T is controlled in this case?

It's from the algorithm itself. JacVec knows it only needs a single chunk to ever do J*v calculations, so it hardcodes a chunk 1. Rosenbrock has a chunk 1 calculation as well. Etc. ForwardDiff has an automatic chunk size for its Jacobian calculations, so DiffCache matches that. That will make things work for any Jacobian calculations internal to SciML. The only other case to handle then is that some things are not computing full Jacobians and so T can be a chunk size smaller than the chosen default, and the view should then make that fine.

[JianghuiDu]

This has been solved by SciML/PreallocationTools.jl@090adf9