ForwardDiff yields incorrect results if one misues `DiffCache`

I’d like to report a mistake that some new users like me might make when using PreallocationTools. One should not use DiffCache to store non-variables such as parameters. It makes perfect sense in a hindsight, because the code will not yield correct residuals when Dual variables are not multiplied by the proper parameter constants.

I’m posting here just for others’ reference. I prefer it to fail explicitly, if there is a way.

using PreallocationTools
using ForwardDiff
using ForwardDiff: Dual
using FiniteDiff

import PreallocationTools: get_tmp


get_tmp(dc::AbstractArray, u) = dc;
get_tmp(dc::Float64, u) = dc;


struct InputCache{T}
    x::T
end


struct OutputCache{T}
    v::T
end


function f!(out, xy, p, input_cache, output_cache)
    # Step 1: copy variables to an input cache
    input_cache = get_tmp(input_cache.x, xy)
    input_cache .= xy

    # Step 2: calculate residuals and save to output cache
    p = get_tmp(p, xy)  # `p` is for parameters. Probably `get_tmp` should never be applied to parameters?

    output_cache = get_tmp(output_cache.v, xy)
    output_cache[1] = 2 * input_cache[1] * p[1]
    output_cache[2] = 2 * input_cache[1] * input_cache[2]

    # Step 3: copy from output cache to the final `out` vector
    out .= output_cache

    # The three steps above demonstrate the calculation phases in a large
    # program, where they can be separated into three functions.

    nothing
end

output_cache = OutputCache(dualcache(zeros(2)));
input_cache = InputCache(dualcache(zeros(2)));

xy = [1., 2];
p = dualcache([1.], 2);  # <- if one mistakenly made `p` a DiffCache

wrap_f! = (out, xy) -> f!(out, xy, p, input_cache, output_cache)

# FiniteDiff works ok
jout = zeros(2, 2);
FiniteDiff.finite_difference_jacobian!(jout, wrap_f!, xy)
@show jout


# incorrect if `p` is DiffCache
out = [0, 0]
ForwardDiff.jacobian(wrap_f!, out, xy)


# correct if `p` is a Vector
p = [1.]  # <- This fixes it, but are there ways to make it failsafe, or at least issue some warning?
wrap_f! = (out, xy) -> f!(out, xy, p, input_cache, output_cache)
ForwardDiff.jacobian(wrap_f!, out, xy)
2 Likes

Yeah, open an issue, but I’m not sure this can be handled. By definition values in a tmp_cache are made to be undef until written to: it’s a cache not a parameter.

1 Like