Handle FIRK in-place Mooncake AD fallback

Route in-place FIRK Mooncake residual AD through ForwardDiff for the boundary-condition paths as well as collocation/full residual paths. This avoids the Julia pre-release Mooncake rule-generation failure while preserving sparse wrapper structure.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

lib/BoundaryValueDiffEqFIRK/src/BoundaryValueDiffEqFIRK.jl

@@ -1,6 +1,6 @@
 module BoundaryValueDiffEqFIRK
 
-using ADTypes: ADTypes, AutoSparse, AutoForwardDiff, AutoEnzyme
+using ADTypes: ADTypes, AutoSparse, AutoForwardDiff, AutoEnzyme, AutoMooncake
 using ArrayInterface: fast_scalar_indexing
 using BandedMatrices: BandedMatrix, Ones
 using BoundaryValueDiffEqCore: BoundaryValueDiffEqCore,

lib/BoundaryValueDiffEqFIRK/src/firk.jl

@@ -657,6 +657,14 @@ function __construct_problem(
 end
 
 __firk_iip_ad_diffmode(::Val{true}, ::AutoEnzyme) = AutoForwardDiff()
+__firk_iip_ad_diffmode(::Val{true}, ::AutoMooncake) = AutoForwardDiff()
+function __firk_iip_ad_diffmode(iip, diffmode::AutoSparse)
+    return AutoSparse(
+        __firk_iip_ad_diffmode(iip, get_dense_ad(diffmode));
+        sparsity_detector = __default_sparsity_detector(diffmode),
+        coloring_algorithm = __default_coloring_algorithm(diffmode)
+    )
+end
 __firk_iip_ad_diffmode(_, diffmode) = diffmode
 
 __firk_iip_needs_diffcache(iip, diffmode::AutoSparse) = __firk_iip_needs_diffcache(
@@ -754,7 +762,7 @@ function __construct_problem(
     ) where {iip, T, DC, tune_parameters, BC, C, LF}
     (; prob, alg, stage, bcresid_prototype, f_prototype) = cache
     (; jac_alg) = alg
-    (; bc_diffmode) = jac_alg
+    bc_diffmode = __firk_iip_ad_diffmode(Val(iip), jac_alg.bc_diffmode)
     N = length(cache.mesh)
 
     resid_bc = cache.bcresid_prototype
@@ -844,7 +852,7 @@ function __construct_problem(
     ) where {iip, T, DC, tune_parameters, BC, C, LF}
     (; prob, alg, stage, bcresid_prototype, f_prototype) = cache
     (; jac_alg) = alg
-    (; bc_diffmode) = jac_alg
+    bc_diffmode = __firk_iip_ad_diffmode(Val(iip), jac_alg.bc_diffmode)
     N = length(cache.mesh)
 
     resid_bc = cache.bcresid_prototype
@@ -1110,7 +1118,7 @@ function __construct_problem(
         loss::LF, ::StandardBVProblem, ::Val{true}
     ) where {iip, T, DC, tune_parameters, BC, C, LF}
     (; jac_alg) = cache.alg
-    (; bc_diffmode) = jac_alg
+    bc_diffmode = __firk_iip_ad_diffmode(Val(iip), jac_alg.bc_diffmode)
     (; bcresid_prototype, f_prototype) = cache
     N = length(cache.mesh)
     resid_bc = cache.bcresid_prototype
@@ -1196,7 +1204,7 @@ function __construct_problem(
         loss::LF, ::StandardBVProblem, ::Val{false}
     ) where {iip, T, DC, tune_parameters, BC, C, LF}
     (; jac_alg) = cache.alg
-    (; bc_diffmode) = jac_alg
+    bc_diffmode = __firk_iip_ad_diffmode(Val(iip), jac_alg.bc_diffmode)
     (; bcresid_prototype, f_prototype, prob) = cache
     N = length(cache.mesh)
     resid_bc = cache.bcresid_prototype