BoundsError with `get_tmp`

[dafred94]

I am facing BoundsErrors when calling get_tmp and can't quite figure out what's going wrong.
My package versions on Julia 1.6.2:

[d236fae5] PreallocationTools v0.2.0
[f6369f11] ForwardDiff v0.10.23

An MWE that produces this error:

using PreallocationTools
import ForwardDiff
using LinearAlgebra

function rhs!(du, u, p, t)
    A = p
    du = get_tmp(du, A)
    mul!(du, A, u)
end

u = [1.0, 0.0]

A = [0.0 1.0;
     1.0 0.0]

du = similar(u)

_du = dualcache(du)

# this one works
rhs!(_du, u, A, 0.0)

# throws BoundsError
ForwardDiff.gradient(A -> rhs!(_du, u, A, 0.0), A)

Am I doing something wrong here?

[toollu]

Yeah i discovered a similar (same?) yesterday, iirc this line is giving an incorrect nelem causing the bounds error.

When I fix

chunksz  = 12
_du = dualcache(du,chunksz)

it works for my case. I can post a running example later when I have access to a machine with julia again.

[ChrisRackauckas]

Does that div need to round up instead of down?

[thomvet]

I’ll check it out in more detail in the morning when at a box with Julia, but my initial thought is that the div shouldn’t need to round at all!

This is a case where you initialize a cache with two elements and then take the gradient with respect to 4. Not sure that this can be automatically resolved, because cache use comes after cache allocation. Little bit of a hen-egg problem.

V0.2 handles the (more common?) opposite case correctly (i.e. downsizing the cache) and flexibly.

One way around the issue is of course to just throw more memory at it (setting chunksize to threshold of 12), which you have already found.

So, preliminarily I would say it’s not automatically resolvable at the code level (I’ll think about it some more) - if you need an automatic solution lbc() might work.

Maybe it’s also just a matter of documenting this limitation?

[dafred94]

With your help, I adjusted my MWE:

using PreallocationTools
import ForwardDiff
using LinearAlgebra


function rhs!(du, u, p, t)
    A = p
    mul!(du, A, u)
end


function loss(du, u, p, t)
    _du = get_tmp(du, p)
    rhs!(_du, u, p, t)
    l = _du[1]
    return l
end


u = [3.0, 0.0]

A = [1.0 1.0;
     1.0 1.0]

du = similar(u)

_du_12 = dualcache(du, 12)
_du_default = dualcache(du)

# works fine!
ForwardDiff.gradient(A -> loss(_du_12, u, A, 0.0), A)

# throws BoundsError
ForwardDiff.gradient(A -> loss(_du_default, u, A, 0.0), A)

[thomvet]

I think, I have an idea how to fix this, but want to check performance on a few real examples, before making a PR.

Stay tuned! :-)

[thomvet]

Made a PR addressing this (see above). Essentially, it just enlarges the cache automatically, which allocates on the first get_tmp call with this specific variable. If the code consists of many, many calls that are the same this is of course negligible. A warning telling the user what the right N (chunk size) would have been also gets thrown.

The updated MWE is below. Note that in your example a chunk size of 4 is actually enough.

using PreallocationTools
using BenchmarkTools
using ForwardDiff
using LinearAlgebra

function rhs!(du, u, p, t)
    A = p
    mul!(du, A, u)
end

function loss(du, u, p, t)
    _du = get_tmp(du, p)
    rhs!(_du, u, p, t)
    l = _du[1]
    return l
end

u = [3.0, 0.0]
A = ones(2,2)

du = similar(u)
_du_adjusted = dualcache(du, 4)
_du_default = dualcache(du)
f1 = A -> loss(_du_adjusted, u, A, 0.0)
f2 = A -> loss(_du_default, u, A, 0.0)

# works fine, no warning
@btime ForwardDiff.gradient(f1, A)

# also works, but throws warning
@btime ForwardDiff.gradient(f2, A)

[thomvet]

@dafred94 and @toollu

Can you confirm that this is working for you now? If yes, we could close this :)

[dafred94]

I finally had some time to test my code with the adjustments you made in the new release. It works!
Thank you very much for your quick fix!