Fix GaussAdjoint with callbacks

[ChrisRackauckas]

No description provided.

[Vaibhavdixit02]

Ah, this was what #1034 was trying to address lol

[ChrisRackauckas]

That would've been good to know 😅

[ChrisRackauckas]

Bump?

[frankschae]

hmm ... it's probably close but it currently fails for the callback where the correction should be 0:

g(sol) = sum(sol)
function dg!(out, u, p, t, i)
    (out .= 1)
end
@testset "callbacks with no effect" begin
    condition(u, t, integrator) = t == 5
    affect!(integrator) = integrator.u[1] += 0.0
    cb = DiscreteCallback(condition, affect!, save_positions = (false, false))
    tstops = [5.0]
    test_discrete_callback(cb, tstops, g, dg!)
end

I think it's the line:

if sensealg isa GaussAdjoint
            @assert integrator.f.f isa ODEGaussAdjointSensitivityFunction
            @show integrator.f.f.integrating_cb.affect!.integrand_values.integrand dgrad
            integrator.f.f.integrating_cb.affect!.integrand_values.integrand .= dgrad

in callback_tracking.jl, whch for the example above prints:

integrator.f.f.integrating_cb.affect!.integrand_values.integrand = [4.318834073956617, -23.50595577533703, 3.507883353046969, -26.70725798446173]
dgrad = [-0.0, -0.0, -0.0, -0.0]

[ChrisRackauckas]

Is that with the VJPs as Enzyme or ReverseDiff? IIUC it always defaults to ReverseDiff right now?

[ChrisRackauckas]

The incorrect values I think stem from missing make_zero!s, but it's currently waiting on an Enzyme tag from @wsmoses before #1067 finishes and then this can get retested.

[ChrisRackauckas]

@jClugstor did your PR look at GaussAdjoint?

[jClugstor]

No, but there is a '@test_broken' for 'GaussAdjoint' with a callback in that PR, so if callbacks get fixed that will need to change.