SparseDiffTools.jl icon indicating copy to clipboard operation
SparseDiffTools.jl copied to clipboard

Threaded Jacobian calls

Open ChrisRackauckas opened this issue 4 years ago • 6 comments

It would be good to port over some of the PolyesterForwardDiff threading support onto the loops here. It wouldn't be too hard but the caches would need to be extended so they don't overlap.

Pinging @chriselrod for reference.

ChrisRackauckas avatar Dec 29 '21 22:12 ChrisRackauckas

So I'm multithreading the same jacobian function with different inputs and I'm seeing incorrect behavior as compared to running on a single thread. I'm preallocating a single cache for the jacobian. I'm assuming this is also due to cache overlap. Is there a better workaround other than preallocating a different cache for each function call? (possible a different cache on each thread?)

aarontrowbridge avatar Sep 09 '23 22:09 aarontrowbridge

Maybe we should consider cacheless forward mode AD. The idea is, if you don't need to mutate x and want derivatives with respect to x, you can add a lazy wrapper that defines getindex to load from x while also initializing the Dual correctly.

chriselrod avatar Sep 09 '23 23:09 chriselrod

This is indeed my situation and I'm very interested in this potential solution! I'm parsing compute_jacobian_ad.jl atm; if you can point me to the right place I'm down to work on implementing this :)

aarontrowbridge avatar Sep 10 '23 00:09 aarontrowbridge

This is indeed my situation and I'm very interested in this potential solution! I'm parsing compute_jacobian_ad.jl atm; if you can point me to the right place I'm down to work on implementing this :)

FWIW, I just implemented that approach in C++ https://github.com/LoopModels/Math/blob/aab7313183efc5655a001866dc84de43ee73daab/include/Math/Dual.hpp#L190-L275 with a basic test here https://github.com/LoopModels/Math/blob/aab7313183efc5655a001866dc84de43ee73daab/test/dual_test.cpp#L16-L44

I'd suggest ignoring ForwardDiff.jl's code, using only ForwardDiff.Dual. Write your own array wrapper like the DualVector there:

# T is the Dual tag to use
# V is the eltype of the wrapped array
# N is the chunk size
struct DualVector{T,V,N,A<:AbstractVector{V}} <: AbstractVector{ForwardDiff.Dual{T,V,N}}
    data::A
    offset::Int
end
Base.@propagate_inbounds function Base.getindex(A::DualVector{T,V,N}, i::Int) where {T,V,N}
    x = A.data[i]
    # NOTE: if you modify this, avoid capturing `V` in a closure
    # because closures do not specialize on types!
    p = ntuple(V∘==(i-offset), Val(N))
    ForwardDiff.Dual{T}(x, p)
end
# TODO: rest of array interface, e.g. `size`, `IndexStyle`, etc should forward

Basically, you pick a chunk size (N), and then loop over blocks of size N, incrementing the offset (which starts at 0), calling f each time. Each call will give you the partials of the N x result_dimension block of the Jacobian. Copy those from the returned value into your result matrix.

chriselrod avatar Sep 10 '23 01:09 chriselrod

Is this issue still active? From the looks of the tests, it looks like PolyesterForwardDiff is supported? Is it still under dev?

Cardoza2 avatar Jun 28 '25 04:06 Cardoza2

This repository is inactive. We recommend users update to using the DifferentiationInterface libraries.

ChrisRackauckas avatar Jun 30 '25 13:06 ChrisRackauckas