add repack/canonicalize in vec_pjac! to support SciMLStructs

Author: albangossard

src/gauss_adjoint.jl

@@ -495,14 +495,15 @@ function vec_pjac!(out, λ, y, t, S::GaussIntegrand)
         vtmp4 .= λ
         Enzyme.remake_zero!(tmp3)
         Enzyme.remake_zero!(out)
-
+        
+        dp = isscimlstructure(p) ? repack(out) : out
         if SciMLBase.isinplace(sol.prob.f)
             Enzyme.remake_zero!(tmp6)
 
             Enzyme.autodiff(
                 Enzyme.Reverse, Enzyme.Duplicated(pf, tmp6), Enzyme.Const,
                 Enzyme.Duplicated(tmp3, tmp4),
-                Enzyme.Const(y), Enzyme.Duplicated(p, out), Enzyme.Const(t))
+                Enzyme.Const(y), Enzyme.Duplicated(p, dp), Enzyme.Const(t))
         else
             function g(du, u, p, t)
                 du .= f(u, p, t)
@@ -512,7 +513,10 @@ function vec_pjac!(out, λ, y, t, S::GaussIntegrand)
             Enzyme.autodiff(
                 Enzyme.Reverse, Enzyme.Duplicated(g, tmp6), Enzyme.Const,
                 Enzyme.Duplicated(tmp3, tmp4),
-                Enzyme.Const(y), Enzyme.Duplicated(p, out), Enzyme.Const(t))
+                Enzyme.Const(y), Enzyme.Duplicated(p, dp), Enzyme.Const(t))
+        end
+        if isscimlstructure(p)
+            out .+= canonicalize(Tunable(), dp)[1]
         end
     elseif sensealg.autojacvec isa MooncakeVJP
         _, _, p_grad = mooncake_run_ad(paramjac_config, y, p, t, λ)