Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix chunksize issue #595

Merged
merged 1 commit into from
Sep 21, 2023
Merged

Fix chunksize issue #595

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 18 additions & 6 deletions ext/OptimizationReverseDiffExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,20 @@

struct OptimizationReverseDiffTag end

function default_chunk_size(len)
if len < ForwardDiff.DEFAULT_CHUNK_THRESHOLD
len

Check warning on line 14 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L12-L14 

Added lines #L12 - L14 were not covered by tests
else
ForwardDiff.DEFAULT_CHUNK_THRESHOLD

Check warning on line 16 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L16 

Added line #L16 was not covered by tests
end
end

function Optimization.instantiate_function(f, x, adtype::AutoReverseDiff,
p = SciMLBase.NullParameters(),
num_cons = 0)
_f = (θ, args...) -> first(f.f(θ, p, args...))

chunksize = default_chunk_size(length(x))

Check warning on line 25 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L25 

Added line #L25 was not covered by tests

if f.grad === nothing
if adtype.compile
Expand All @@ -32,14 +42,14 @@
if f.hess === nothing
if adtype.compile
T = ForwardDiff.Tag(OptimizationReverseDiffTag(),eltype(x))
xdual = ForwardDiff.Dual{typeof(T),eltype(x),length(x)}.(x, Ref(ForwardDiff.Partials((ones(eltype(x), length(x))...,))))
xdual = ForwardDiff.Dual{typeof(T),eltype(x),chunksize}.(x, Ref(ForwardDiff.Partials((ones(eltype(x), chunksize)...,))))

Check warning on line 45 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L45 

Added line #L45 was not covered by tests
h_tape = ReverseDiff.GradientTape(_f, xdual)
htape = ReverseDiff.compile(h_tape)
function g(θ)
res1 = zeros(eltype(θ), length(θ))
ReverseDiff.gradient!(res1, htape, θ)
end
jaccfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk(x), T)
jaccfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk{chunksize}(), T)

Check warning on line 52 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L52 

Added line #L52 was not covered by tests
hess = function (res, θ, args...)
ForwardDiff.jacobian!(res, g, θ, jaccfg, Val{false}())
end
Expand Down Expand Up @@ -100,7 +110,7 @@
ReverseDiff.gradient!(res1, htape, θ)
end
gs = [x -> grad_cons(x, conshtapes[i]) for i in 1:num_cons]
jaccfgs = [ForwardDiff.JacobianConfig(gs[i], x, ForwardDiff.Chunk(x), T) for i in 1:num_cons]
jaccfgs = [ForwardDiff.JacobianConfig(gs[i], x, ForwardDiff.Chunk{chunksize}(), T) for i in 1:num_cons]

Check warning on line 113 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L113 

Added line #L113 was not covered by tests
cons_h = function (res, θ)
for i in 1:num_cons
ForwardDiff.jacobian!(res[i], gs[i], θ, jaccfgs[i], Val{false}())
Expand Down Expand Up @@ -134,6 +144,8 @@
adtype::AutoReverseDiff, num_cons = 0)
_f = (θ, args...) -> first(f.f(θ, cache.p, args...))

chunksize = default_chunk_size(length(cache.u0))

Check warning on line 147 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L147 

Added line #L147 was not covered by tests

if f.grad === nothing
if adtype.compile
_tape = ReverseDiff.GradientTape(_f, cache.u0)
Expand All @@ -152,14 +164,14 @@
if f.hess === nothing
if adtype.compile
T = ForwardDiff.Tag(OptimizationReverseDiffTag(),eltype(cache.u0))
xdual = ForwardDiff.Dual{typeof(T),eltype(cache.u0),length(cache.u0)}.(cache.u0, Ref(ForwardDiff.Partials((ones(eltype(cache.u0), length(cache.u0))...,))))
xdual = ForwardDiff.Dual{typeof(T),eltype(cache.u0),chunksize}.(cache.u0, Ref(ForwardDiff.Partials((ones(eltype(cache.u0), chunksize)...,))))

Check warning on line 167 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L167 

Added line #L167 was not covered by tests
h_tape = ReverseDiff.GradientTape(_f, xdual)
htape = ReverseDiff.compile(h_tape)
function g(θ)
res1 = zeros(eltype(θ), length(θ))
ReverseDiff.gradient!(res1, htape, θ)
end
jaccfg = ForwardDiff.JacobianConfig(g, cache.u0, ForwardDiff.Chunk(cache.u0), T)
jaccfg = ForwardDiff.JacobianConfig(g, cache.u0, ForwardDiff.Chunk{chunksize}(), T)

Check warning on line 174 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L174 

Added line #L174 was not covered by tests
hess = function (res, θ, args...)
ForwardDiff.jacobian!(res, g, θ, jaccfg, Val{false}())
end
Expand Down Expand Up @@ -220,7 +232,7 @@
ReverseDiff.gradient!(res1, htape, θ)
end
gs = [x -> grad_cons(x, conshtapes[i]) for i in 1:num_cons]
jaccfgs = [ForwardDiff.JacobianConfig(gs[i], cache.u0, ForwardDiff.Chunk(cache.u0), T) for i in 1:num_cons]
jaccfgs = [ForwardDiff.JacobianConfig(gs[i], cache.u0, ForwardDiff.Chunk{chunksize}(), T) for i in 1:num_cons]

Check warning on line 235 in ext/OptimizationReverseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationReverseDiffExt.jl#L235 

Added line #L235 was not covered by tests
cons_h = function (res, θ)
for i in 1:num_cons
ForwardDiff.jacobian!(res[i], gs[i], θ, jaccfgs[i], Val{false}())
Expand Down
14 changes: 8 additions & 6 deletions ext/OptimizationSparseDiffExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -492,6 +492,8 @@
num_cons = 0)
_f = (θ, args...) -> first(f.f(θ, p, args...))

chunksize = default_chunk_size(length(x))

Check warning on line 495 in ext/OptimizationSparseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationSparseDiffExt.jl#L495 

Added line #L495 was not covered by tests

if f.grad === nothing
if adtype.compile
_tape = ReverseDiff.GradientTape(_f, x)
Expand All @@ -514,7 +516,7 @@
hess_colors = SparseDiffTools.matrix_colors(tril(hess_sparsity))
if adtype.compile
T = ForwardDiff.Tag(OptimizationSparseReverseTag(),eltype(x))
xdual = ForwardDiff.Dual{typeof(T),eltype(x),length(x)}.(x, Ref(ForwardDiff.Partials((ones(eltype(x), length(x))...,))))
xdual = ForwardDiff.Dual{typeof(T),eltype(x),min(chunksize, maximum(hess_colors))}.(x, Ref(ForwardDiff.Partials((ones(eltype(x), min(chunksize, maximum(hess_colors)))...,))))

Check warning on line 519 in ext/OptimizationSparseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationSparseDiffExt.jl#L519 

Added line #L519 was not covered by tests
h_tape = ReverseDiff.GradientTape(_f, xdual)
htape = ReverseDiff.compile(h_tape)
function g(res1, θ)
Expand Down Expand Up @@ -582,15 +584,14 @@
conshess_colors = SparseDiffTools.matrix_colors.(conshess_sparsity)
if adtype.compile
T = ForwardDiff.Tag(OptimizationSparseReverseTag(),eltype(x))
xduals = [ForwardDiff.Dual{typeof(T),eltype(x),maximum(conshess_colors[i])}.(x, Ref(ForwardDiff.Partials((ones(eltype(x), maximum(conshess_colors[i]))...,)))) for i in 1:num_cons]
xduals = [ForwardDiff.Dual{typeof(T),eltype(x),min(chunksize, maximum(conshess_colors[i]))}.(x, Ref(ForwardDiff.Partials((ones(eltype(x), min(chunksize, maximum(conshess_colors[i])))...,)))) for i in 1:num_cons]

Check warning on line 587 in ext/OptimizationSparseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationSparseDiffExt.jl#L587 

Added line #L587 was not covered by tests
consh_tapes = [ReverseDiff.GradientTape(fncs[i], xduals[i]) for i in 1:num_cons]
conshtapes = ReverseDiff.compile.(consh_tapes)
function grad_cons(res1, θ, htape)
ReverseDiff.gradient!(res1, htape, θ)
end
gs = [(res1, x) -> grad_cons(res1, x, conshtapes[i]) for i in 1:num_cons]
jaccfgs = [ForwardColorJacCache(gs[i], x; tag = typeof(T), colorvec = conshess_colors[i], sparsity = conshess_sparsity[i]) for i in 1:num_cons]
println(jaccfgs)
cons_h = function (res, θ)
for i in 1:num_cons
SparseDiffTools.forwarddiff_color_jacobian!(res[i], gs[i], θ, jaccfgs[i])
Expand Down Expand Up @@ -629,6 +630,8 @@
adtype::AutoSparseReverseDiff, num_cons = 0)
_f = (θ, args...) -> first(f.f(θ, cache.p, args...))

chunksize = default_chunk_size(length(cache.u0))

Check warning on line 633 in ext/OptimizationSparseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationSparseDiffExt.jl#L633 

Added line #L633 was not covered by tests

if f.grad === nothing
if adtype.compile
_tape = ReverseDiff.GradientTape(_f, cache.u0)
Expand All @@ -651,7 +654,7 @@
hess_colors = SparseDiffTools.matrix_colors(tril(hess_sparsity))
if adtype.compile
T = ForwardDiff.Tag(OptimizationSparseReverseTag(),eltype(cache.u0))
xdual = ForwardDiff.Dual{typeof(T),eltype(cache.u0),length(cache.u0)}.(cache.u0, Ref(ForwardDiff.Partials((ones(eltype(cache.u0), length(cache.u0))...,))))
xdual = ForwardDiff.Dual{typeof(T),eltype(cache.u0),min(chunksize, maximum(hess_colors))}.(cache.u0, Ref(ForwardDiff.Partials((ones(eltype(cache.u0), min(chunksize, maximum(hess_colors)))...,))))

Check warning on line 657 in ext/OptimizationSparseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationSparseDiffExt.jl#L657 

Added line #L657 was not covered by tests
h_tape = ReverseDiff.GradientTape(_f, xdual)
htape = ReverseDiff.compile(h_tape)
function g(res1, θ)
Expand Down Expand Up @@ -719,15 +722,14 @@
conshess_colors = SparseDiffTools.matrix_colors.(conshess_sparsity)
if adtype.compile
T = ForwardDiff.Tag(OptimizationSparseReverseTag(),eltype(cache.u0))
xduals = [ForwardDiff.Dual{typeof(T),eltype(cache.u0),maximum(conshess_colors[i])}.(cache.u0, Ref(ForwardDiff.Partials((ones(eltype(cache.u0), maximum(conshess_colors[i]))...,)))) for i in 1:num_cons]
xduals = [ForwardDiff.Dual{typeof(T),eltype(cache.u0),min(chunksize, maximum(conshess_colors[i]))}.(cache.u0, Ref(ForwardDiff.Partials((ones(eltype(cache.u0), min(chunksize, maximum(conshess_colors[i])))...,)))) for i in 1:num_cons]

Check warning on line 725 in ext/OptimizationSparseDiffExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/OptimizationSparseDiffExt.jl#L725 

Added line #L725 was not covered by tests
consh_tapes = [ReverseDiff.GradientTape(fncs[i], xduals[i]) for i in 1:num_cons]
conshtapes = ReverseDiff.compile.(consh_tapes)
function grad_cons(res1, θ, htape)
ReverseDiff.gradient!(res1, htape, θ)
end
gs = [(res1, x) -> grad_cons(res1, x, conshtapes[i]) for i in 1:num_cons]
jaccfgs = [ForwardColorJacCache(gs[i], cache.u0; tag = typeof(T), colorvec = conshess_colors[i], sparsity = conshess_sparsity[i]) for i in 1:num_cons]
println(jaccfgs)
cons_h = function (res, θ)
for i in 1:num_cons
SparseDiffTools.forwarddiff_color_jacobian!(res[i], gs[i], θ, jaccfgs[i])
Expand Down
Loading