fix: fixes checkpointing issues for GaussKronrod Adjoint

maxesit · maxesit · commit ea2ba092dd66 · 2025-07-27T16:48:35.000+02:00
diff --git a/src/gauss_adjoint.jl b/src/gauss_adjoint.jl
@@ -1,6 +1,3 @@
-AbstractGAdjoint = Union{GaussAdjoint, GaussKronrodAdjoint} # needs to be supertype of GaussAdjoint && GaussKronrodAdjoint !!!
-
-
 mutable struct GaussIntegrand{pType, uType, lType, rateType, S, PF, PJC, PJT, DGP,
     G, SAlg <: AbstractGAdjoint}
     sol::S
@@ -579,8 +576,6 @@ function _adjoint_sensitivities(sol, sensealg::AbstractGAdjoint, alg; t = nothin
     elseif sensealg isa GaussKronrodAdjoint
         cb = IntegratingGKSumCallback((out, u, t, integrator) -> integrand(out, t, u), 
             integrand_values, allocate_vjp(tunables))
-    else
-        print("\n\nerror\n\n  typeof(sensealg) = ", typeof(sensealg),"\n\n")
     end
     rcb = nothing
     cb2 = nothing
diff --git a/src/gauss_adjoint.jl~ b/src/gauss_adjoint.jl~
@@ -1,5 +1,4 @@
-AbstractGAdjoint = Union{GaussAdjoint, GaussKronrodAdjoint} # needs to be supertype of GaussAdjoint && GaussKronrodAdjoint !!!
-
+# AbstractGAdjoint = Union{GaussAdjoint, GaussKronrodAdjoint} # needs to be supertype of GaussAdjoint && GaussKronrodAdjoint !!!
 
 mutable struct GaussIntegrand{pType, uType, lType, rateType, S, PF, PJC, PJT, DGP,
     G, SAlg <: AbstractGAdjoint}
@@ -573,14 +572,14 @@ function _adjoint_sensitivities(sol, sensealg::AbstractGAdjoint, alg; t = nothin
     end
     integrand = GaussIntegrand(sol, sensealg, checkpoints, dgdp_continuous)
     integrand_values = IntegrandValuesSum(allocate_zeros(tunables))
-    if typeof(sensealg) == GaussAdjoint
+    if sensealg isa GaussAdjoint
 	cb = IntegratingSumCallback((out, u, t, integrator) -> integrand(out, t, u),
             integrand_values, allocate_vjp(tunables))
-    elseif typeof(sensealg) == GaussKronrodAdjoint
+    elseif sensealg isa GaussKronrodAdjoint
         cb = IntegratingGKSumCallback((out, u, t, integrator) -> integrand(out, t, u), 
             integrand_values, allocate_vjp(tunables))
     else
-        print("\n\nerror\n\n  typeof(sensealg) = ", typeof(sensealg))
+        print("\n\nerror\n\n  typeof(sensealg) = ", typeof(sensealg),"\n\n")
     end
     rcb = nothing
     cb2 = nothing
diff --git a/src/sensitivity_algorithms.jl b/src/sensitivity_algorithms.jl
@@ -667,6 +667,8 @@ function setvjp(sensealg::GaussKronrodAdjoint{CS, AD, FDT, Nothing}, vjp) where
     GaussKronrodAdjoint{CS, AD, FDT, typeof(vjp)}(vjp, sensealg.checkpointing)
 end
 
+AbstractGAdjoint = Union{GaussAdjoint, GaussKronrodAdjoint}
+
 """
 ```julia
 TrackerAdjoint <: AbstractAdjointSensitivityAlgorithm{nothing, true, nothing}
@@ -1384,6 +1386,8 @@ end
 @inline ischeckpointing(alg::InterpolatingAdjoint, sol) = alg.checkpointing || !sol.dense
 @inline ischeckpointing(alg::GaussAdjoint) = alg.checkpointing
 @inline ischeckpointing(alg::GaussAdjoint, sol) = alg.checkpointing || !sol.dense
+@inline ischeckpointing(alg::GaussKronrodAdjoint) = alg.checkpointing
+@inline ischeckpointing(alg::GaussKronrodAdjoint, sol) = alg.checkpointing || !sol.dense
 @inline ischeckpointing(alg::BacksolveAdjoint, sol = nothing) = alg.checkpointing
 
 @inline isnoisemixing(alg::AbstractSensitivityAlgorithm) = false
diff --git a/src/sensitivity_algorithms.jl~ b/src/sensitivity_algorithms.jl~
