How to make this function compatible with ForwardDiff?

Well, yeah, you didn’t match the chunk sizes. If you’re differentiation has a single chunk you need a cache of a single chunk as well. Its default is to assume Jacobians w.r.t. a sized variable matching the cache, but if you’re differentiating w.r.t. a scalar you need to have a single chunk. So the following works:

using ForwardDiff
using PreallocationTools

randmat = rand(10, 2)
sto = similar(randmat)
stod = dualcache(sto, Val{1})

function claytonsample!(sto, τ; randmat=randmat)
    sto = get_tmp(sto, τ)
    @show size(sto), size(randmat)
    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)
1 Like