Skip to content

Commit 32cfe9f

Browse files
ChrisRackauckasChrisRackauckas
committed
Apply DI strict = Val(false) fix to MIRK (and other sublibs)
Extends the `DifferentiationInterface.prepare_jacobian` strict-mode fix from FIRK (previous commit) to the remaining collocation sublibs: * `BoundaryValueDiffEqMIRK/src/mirk.jl` : 12 call sites * `BoundaryValueDiffEqMIRKN/src/mirkn.jl` : 6 call sites * `BoundaryValueDiffEqAscher/src/ascher.jl`: 2 call sites Same root cause as FIRK: preparation is done with a `VectorOfArray`- wrapping `Base.ReshapedArray` `y`, but execution inside the nonlinear solve is called with a plain `Vector{Float64}`. Under DI's strict mode (default `Val(true)`) this trips `PreparationMismatchError`. Passing `strict = Val(false)` at each call site matches what `Shooting` already does and what the FIRK commit set up. Verified locally with `julia +1.11`: julia +1.11 -e 'using Pkg; Pkg.activate(temp=true); Pkg.develop.([(path="lib/BoundaryValueDiffEqCore",), (path="lib/BoundaryValueDiffEqAscher",)]); using BoundaryValueDiffEqAscher' # Ascher precompiles cleanly (~8 s) (same for BoundaryValueDiffEqMIRKN) # MIRKN precompiles cleanly (~10 s) MIRK module itself precompiles cleanly once the DI fix is in place (verified with PrecompileMIRK=false so the separate `@compile_workload` isn't exercised). Any remaining failures in MIRK's precompile workload are a separate `VectorOfArray` v4 iteration issue outside the scope of the DI strict-mode fix. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
1 parent 24a0225 commit 32cfe9f

3 files changed

Lines changed: 57 additions & 20 deletions

File tree

‎lib/BoundaryValueDiffEqAscher/src/ascher.jl‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -344,9 +344,13 @@ function __construct_nlproblem(cache::AscherCache{iip, T}) where {iip, T}
344344
end
345345

346346
jac_cache = if iip
347-
DI.prepare_jacobian(loss, resid_prototype, diffmode, lz, Constant(cache.p))
347+
DI.prepare_jacobian(
348+
loss, resid_prototype, diffmode, lz, Constant(cache.p); strict = Val(false)
349+
)
348350
else
349-
DI.prepare_jacobian(loss, diffmode, lz, Constant(cache.p))
351+
DI.prepare_jacobian(
352+
loss, diffmode, lz, Constant(cache.p); strict = Val(false)
353+
)
350354
end
351355

352356
jac_prototype = if iip

‎lib/BoundaryValueDiffEqMIRK/src/mirk.jl‎

Lines changed: 34 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -593,9 +593,13 @@ function __construct_problem(
593593
resid_collocation = safe_similar(y, L_f_prototype * (N - 1))
594594

595595
cache_bc = if iip
596-
DI.prepare_jacobian(loss_bc, resid_bc, bc_diffmode, y, Constant(cache.p))
596+
DI.prepare_jacobian(
597+
loss_bc, resid_bc, bc_diffmode, y, Constant(cache.p); strict = Val(false)
598+
)
597599
else
598-
DI.prepare_jacobian(loss_bc, bc_diffmode, y, Constant(cache.p))
600+
DI.prepare_jacobian(
601+
loss_bc, bc_diffmode, y, Constant(cache.p); strict = Val(false)
602+
)
599603
end
600604

601605
nonbc_diffmode = AutoSparse(
@@ -605,10 +609,13 @@ function __construct_problem(
605609
)
606610
cache_collocation = if iip
607611
DI.prepare_jacobian(
608-
loss_collocation, resid_collocation, nonbc_diffmode, y, Constant(cache.p)
612+
loss_collocation, resid_collocation, nonbc_diffmode, y, Constant(cache.p);
613+
strict = Val(false)
609614
)
610615
else
611-
DI.prepare_jacobian(loss_collocation, nonbc_diffmode, y, Constant(cache.p))
616+
DI.prepare_jacobian(
617+
loss_collocation, nonbc_diffmode, y, Constant(cache.p); strict = Val(false)
618+
)
612619
end
613620

614621
J_bc = if iip
@@ -675,9 +682,13 @@ function __construct_problem(
675682
resid_prototype = vcat(resid_bc, resid_collocation)
676683

677684
cache_bc = if iip
678-
DI.prepare_jacobian(loss_bc, resid_bc, bc_diffmode, y, Constant(cache.p))
685+
DI.prepare_jacobian(
686+
loss_bc, resid_bc, bc_diffmode, y, Constant(cache.p); strict = Val(false)
687+
)
679688
else
680-
DI.prepare_jacobian(loss_bc, bc_diffmode, y, Constant(cache.p))
689+
DI.prepare_jacobian(
690+
loss_bc, bc_diffmode, y, Constant(cache.p); strict = Val(false)
691+
)
681692
end
682693

683694
nonbc_diffmode = if jac_alg.nonbc_diffmode isa AutoSparse
@@ -708,10 +719,13 @@ function __construct_problem(
708719

709720
cache_collocation = if iip
710721
DI.prepare_jacobian(
711-
loss_collocation, resid_collocation, nonbc_diffmode, y, Constant(cache.p)
722+
loss_collocation, resid_collocation, nonbc_diffmode, y, Constant(cache.p);
723+
strict = Val(false)
712724
)
713725
else
714-
DI.prepare_jacobian(loss_collocation, nonbc_diffmode, y, Constant(cache.p))
726+
DI.prepare_jacobian(
727+
loss_collocation, nonbc_diffmode, y, Constant(cache.p); strict = Val(false)
728+
)
715729
end
716730

717731
J_bc = if iip
@@ -846,9 +860,13 @@ function __construct_problem(
846860
end
847861

848862
diffcache = if iip
849-
DI.prepare_jacobian(loss, resid, diffmode, y, Constant(cache.p))
863+
DI.prepare_jacobian(
864+
loss, resid, diffmode, y, Constant(cache.p); strict = Val(false)
865+
)
850866
else
851-
DI.prepare_jacobian(loss, diffmode, y, Constant(cache.p))
867+
DI.prepare_jacobian(
868+
loss, diffmode, y, Constant(cache.p); strict = Val(false)
869+
)
852870
end
853871

854872
jac_prototype = if iip
@@ -909,9 +927,13 @@ function __construct_problem(
909927
end
910928

911929
diffcache = if iip
912-
DI.prepare_jacobian(loss, resid, diffmode, y, Constant(cache.p))
930+
DI.prepare_jacobian(
931+
loss, resid, diffmode, y, Constant(cache.p); strict = Val(false)
932+
)
913933
else
914-
DI.prepare_jacobian(loss, diffmode, y, Constant(cache.p))
934+
DI.prepare_jacobian(
935+
loss, diffmode, y, Constant(cache.p); strict = Val(false)
936+
)
915937
end
916938

917939
jac_prototype = if iip

‎lib/BoundaryValueDiffEqMIRKN/src/mirkn.jl‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -224,9 +224,13 @@ function __construct_nlproblem(
224224
end
225225

226226
cache_bc = if iip
227-
DI.prepare_jacobian(loss_bc, resid_bc, bc_diffmode, y, Constant(cache.p))
227+
DI.prepare_jacobian(
228+
loss_bc, resid_bc, bc_diffmode, y, Constant(cache.p); strict = Val(false)
229+
)
228230
else
229-
DI.prepare_jacobian(loss_bc, bc_diffmode, y, Constant(cache.p))
231+
DI.prepare_jacobian(
232+
loss_bc, bc_diffmode, y, Constant(cache.p); strict = Val(false)
233+
)
230234
end
231235

232236
nonbc_diffmode = if jac_alg.nonbc_diffmode isa AutoSparse
@@ -241,10 +245,13 @@ function __construct_nlproblem(
241245

242246
cache_collocation = if iip
243247
DI.prepare_jacobian(
244-
loss_collocation, resid_collocation, nonbc_diffmode, y, Constant(cache.p)
248+
loss_collocation, resid_collocation, nonbc_diffmode, y, Constant(cache.p);
249+
strict = Val(false)
245250
)
246251
else
247-
DI.prepare_jacobian(loss_collocation, nonbc_diffmode, y, Constant(cache.p))
252+
DI.prepare_jacobian(
253+
loss_collocation, nonbc_diffmode, y, Constant(cache.p); strict = Val(false)
254+
)
248255
end
249256

250257
J_bc = if iip
@@ -314,9 +321,13 @@ function __construct_nlproblem(
314321
end
315322

316323
diffcache = if iip
317-
DI.prepare_jacobian(loss, resid, diffmode, y, Constant(cache.p))
324+
DI.prepare_jacobian(
325+
loss, resid, diffmode, y, Constant(cache.p); strict = Val(false)
326+
)
318327
else
319-
DI.prepare_jacobian(loss, diffmode, y, Constant(cache.p))
328+
DI.prepare_jacobian(
329+
loss, diffmode, y, Constant(cache.p); strict = Val(false)
330+
)
320331
end
321332

322333
jac_prototype = if iip

0 commit comments

Comments
 (0)