Propagate multiple shooting jac_alg

Author: avik-pal

src/solve/multiple_shooting.jl

@@ -66,9 +66,17 @@ function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;),
         end
     end
 
+    if prob.problem_type isa TwoPointBVProblem
+        diffmode_shooting = __get_non_sparse_ad(alg.jac_alg.diffmode)
+    else
+        diffmode_shooting = __get_non_sparse_ad(alg.jac_alg.bc_diffmode)
+    end
+    shooting_alg = Shooting(alg.ode_alg, alg.nlsolve,
+        BVPJacobianAlgorithm(diffmode_shooting))
+
     single_shooting_prob = remake(prob; u0 = reshape(u_at_nodes[1:N], u0_size))
-    return __solve(single_shooting_prob, Shooting(alg.ode_alg; alg.nlsolve);
-        odesolve_kwargs, nlsolve_kwargs, verbose, kwargs...)
+    return __solve(single_shooting_prob, shooting_alg; odesolve_kwargs, nlsolve_kwargs,
+        verbose, kwargs...)
 end
 
 # TODO: We can save even more memory by hoisting the preallocated caches for the ODEs
@@ -150,10 +158,8 @@ function __solve_nlproblem!(::StandardBVProblem, alg::MultipleShooting, bcresid_
         internal_ode_kwargs...)
 
     # BC Part
-    if alg.jac_alg.bc_diffmode isa AbstractSparseADType
-        error("Multiple Shooting doesn't support sparse AD for Boundary Conditions yet!")
-    end
-    sd_bc = NoSparsityDetection()
+    sd_bc = alg.jac_alg.bc_diffmode isa AbstractSparseADType ?
+            SymbolicsSparsityDetection() : NoSparsityDetection()
     bc_jac_cache = sparse_jacobian_cache(alg.jac_alg.bc_diffmode,
         sd_bc, nothing, similar(bcresid_prototype), u_at_nodes)
     ode_cache_bc_jac_fn = __multiple_shooting_init_jacobian_odecache(ensemblealg, prob,

src/utils.jl

@@ -203,6 +203,25 @@ end
 __vec_bc(sol, p, t, bc, u_size) = vec(bc(__restructure_sol(sol, u_size), p, t))
 __vec_bc(sol, p, bc, u_size) = vec(bc(reshape(sol, u_size), p))
 
+__get_non_sparse_ad(ad::AbstractADType) = ad
+function __get_non_sparse_ad(ad::AbstractSparseADType)
+    if ad isa AutoSparseForwardDiff
+        return AutoForwardDiff{__get_chunksize(ad), typeof(ad.tag)}(ad.tag)
+    elseif ad isa AutoSparseEnzyme
+        return AutoEnzyme()
+    elseif ad isa AutoSparseFiniteDiff
+        return AutoFiniteDiff()
+    elseif ad isa AutoSparseReverseDiff
+        return AutoReverseDiff(ad.compile)
+    elseif ad isa AutoSparseZygote
+        return AutoZygote()
+    else
+        throw(ArgumentError("Unknown AD Type"))
+    end
+end
+
+__get_chunksize(::AutoSparseForwardDiff{CK}) where {CK} = CK
+
 # Restructure Solution
 function __restructure_sol(sol::Vector{<:AbstractArray}, u_size)
     return map(Base.Fix2(reshape, u_size), sol)