From 93a82877af2dade388a0040344f01749c61712a0 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 12 Dec 2021 08:19:53 -0500 Subject: [PATCH 1/4] Allow users to set the tag for the configs --- src/differentiation/compute_jacobian_ad.jl | 24 ++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index d2ef9154..4eb97cb0 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -11,11 +11,13 @@ end getsize(::Val{N}) where N = N getsize(N::Integer) = N void_setindex!(args...) = (setindex!(args...); return) +gettag(::Type{ForwardDiff.Dual{T}}) where {T} = T const default_chunk_size = ForwardDiff.pickchunksize function ForwardColorJacCache(f::F,x,_chunksize = nothing; dx = nothing, + tag = nothing, colorvec=1:length(x), sparsity::Union{AbstractArray,Nothing}=nothing) where {F} @@ -25,15 +27,21 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; chunksize = _chunksize end + if tag === nothing + T = typeof(ForwardDiff.Tag(f,eltype(vec(x)))) + else + T = tag + end + if x isa Array p = generate_chunked_partials(x,colorvec,chunksize) - t = similar(x,Dual{typeof(ForwardDiff.Tag(f,eltype(vec(x)))),eltype(x),length(first(first(p)))}) + t = similar(x,Dual{T}) for i in eachindex(t) - t[i] = Dual{typeof(ForwardDiff.Tag(f,eltype(vec(x)))),eltype(x),length(first(first(p)))}(x[i],ForwardDiff.Partials(first(p)[i])) + t[i] = Dual{T,eltype(x),length(first(first(p)))}(x[i],ForwardDiff.Partials(first(p)[i])) end else p = adapt.(parameterless_type(x),generate_chunked_partials(x,colorvec,chunksize)) - _t = Dual{typeof(ForwardDiff.Tag(f,eltype(vec(x))))}.(vec(x),first(p)) + _t = Dual{T}.(vec(x),first(p)) t = ArrayInterface.restructure(x,_t) end @@ -44,7 +52,7 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; else tup = ArrayInterface.allowed_getindex(ArrayInterface.allowed_getindex(p,1),1) .* false _pi = adapt(parameterless_type(dx),[tup for i in 1:length(dx)]) - fx = reshape(Dual{typeof(ForwardDiff.Tag(f,eltype(vec(x))))}.(vec(dx),_pi),size(dx)...) + fx = reshape(Dual{T}.(vec(dx),_pi),size(dx)...) _dx = dx end @@ -162,7 +170,7 @@ function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number},f::F,x::Abstract for i in eachindex(p) partial_i = p[i] - t = reshape(Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}.(vecx, partial_i),size(t)) + t = reshape(Dual{gettag(eltype(t))}.(vecx, partial_i),size(t)) fx = f(t) if !(sparsity isa Nothing) for j in 1:chunksize @@ -230,7 +238,7 @@ function forwarddiff_color_jacobian_immutable(f,x::AbstractArray{<:Number},jac_c for i in eachindex(p) partial_i = p[i] - t = reshape(Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}.(vecx, partial_i),size(t)) + t = reshape(Dual{gettag(eltype(t))}.(vecx, partial_i),size(t)) fx = f(t) if !(sparsity isa Nothing) for j in 1:chunksize @@ -311,10 +319,10 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, if vect isa Array @inbounds @simd ivdep for j in eachindex(vect) - vect[j] = Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}(vecx[j], partial_i[j]) + vect[j] = Dual{gettag(eltype(t))}(vecx[j], partial_i[j]) end else - vect .= Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}.(vecx, partial_i) + vect .= Dual{gettag(eltype(t))}.(vecx, partial_i) end f(fx,t) From 7879092a475dd198fe86c914c3d851c58a37790a Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 12 Dec 2021 08:52:54 -0500 Subject: [PATCH 2/4] always fully make the dual --- src/differentiation/compute_jacobian_ad.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index 4eb97cb0..f7233d35 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -41,7 +41,7 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; end else p = adapt.(parameterless_type(x),generate_chunked_partials(x,colorvec,chunksize)) - _t = Dual{T}.(vec(x),first(p)) + _t = Dual{T,eltype(x),length(first(first(p)))}.(vec(x),first(p)) t = ArrayInterface.restructure(x,_t) end @@ -52,7 +52,7 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; else tup = ArrayInterface.allowed_getindex(ArrayInterface.allowed_getindex(p,1),1) .* false _pi = adapt(parameterless_type(dx),[tup for i in 1:length(dx)]) - fx = reshape(Dual{T}.(vec(dx),_pi),size(dx)...) + fx = reshape(Dual{T,eltype(dx),length(tup)}.(vec(dx),_pi),size(dx)...) _dx = dx end From f9fd755a26845017888c7b6468000de2695e9b57 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Tue, 14 Dec 2021 17:48:35 -0500 Subject: [PATCH 3/4] fix tag construction --- src/differentiation/compute_jacobian_ad.jl | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index f7233d35..e205cad0 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -14,6 +14,7 @@ void_setindex!(args...) = (setindex!(args...); return) gettag(::Type{ForwardDiff.Dual{T}}) where {T} = T const default_chunk_size = ForwardDiff.pickchunksize +const SMALLTAG = ForwardDiff.Tag(missing,Float64) function ForwardColorJacCache(f::F,x,_chunksize = nothing; dx = nothing, @@ -41,7 +42,7 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; end else p = adapt.(parameterless_type(x),generate_chunked_partials(x,colorvec,chunksize)) - _t = Dual{T,eltype(x),length(first(first(p)))}.(vec(x),first(p)) + _t = Dual{T,eltype(x),length(first(first(p)))}.(vec(x),ForwardDiff.Partials.(first(p))) t = ArrayInterface.restructure(x,_t) end @@ -52,7 +53,7 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; else tup = ArrayInterface.allowed_getindex(ArrayInterface.allowed_getindex(p,1),1) .* false _pi = adapt(parameterless_type(dx),[tup for i in 1:length(dx)]) - fx = reshape(Dual{T,eltype(dx),length(tup)}.(vec(dx),_pi),size(dx)...) + fx = reshape(Dual{T,eltype(dx),length(tup)}.(vec(dx),ForwardDiff.Partials.(_pi)),size(dx)...) _dx = dx end @@ -170,7 +171,7 @@ function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number},f::F,x::Abstract for i in eachindex(p) partial_i = p[i] - t = reshape(Dual{gettag(eltype(t))}.(vecx, partial_i),size(t)) + t = reshape(eltype(t).(vecx, ForwardDiff.Partials.(partial_i)),size(t)) fx = f(t) if !(sparsity isa Nothing) for j in 1:chunksize @@ -238,7 +239,7 @@ function forwarddiff_color_jacobian_immutable(f,x::AbstractArray{<:Number},jac_c for i in eachindex(p) partial_i = p[i] - t = reshape(Dual{gettag(eltype(t))}.(vecx, partial_i),size(t)) + t = reshape(eltype(t).(vecx, ForwardDiff.Partials.(partial_i)),size(t)) fx = f(t) if !(sparsity isa Nothing) for j in 1:chunksize @@ -319,10 +320,10 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, if vect isa Array @inbounds @simd ivdep for j in eachindex(vect) - vect[j] = Dual{gettag(eltype(t))}(vecx[j], partial_i[j]) + vect[j] = eltype(t)(vecx[j], ForwardDiff.Partials(partial_i[j])) end else - vect .= Dual{gettag(eltype(t))}.(vecx, partial_i) + vect .= eltype(t).(vecx, ForwardDiff.Partials.(partial_i)) end f(fx,t) From aa9fc34a308566c86255d2204cd92cbfb8ef478f Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Tue, 14 Dec 2021 17:59:26 -0500 Subject: [PATCH 4/4] fix for GPU? --- src/differentiation/compute_jacobian_ad.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index e205cad0..782ae35a 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -42,7 +42,7 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing; end else p = adapt.(parameterless_type(x),generate_chunked_partials(x,colorvec,chunksize)) - _t = Dual{T,eltype(x),length(first(first(p)))}.(vec(x),ForwardDiff.Partials.(first(p))) + _t = Dual{T,eltype(x),getsize(chunksize)}.(vec(x),ForwardDiff.Partials.(first(p))) t = ArrayInterface.restructure(x,_t) end