Odd behaviour regarding Adjoint Sensitivity Analysis

Specific Domains Numerics
, ,
8

You were just returning a vector of vectors. AD libraries do not know how to represent Jacobians of such outputs. Instead, just return a normal array and it’s fine:

using LinearAlgebra
using OrdinaryDiffEq
using SciMLSensitivity
using Interpolations
using Zygote
using ForwardDiff
using Plots

T = 4.0
N = 40
dt = T / N
np = N + 1 

A = [0 1 0 0; 0 0 1 0; 0 0 0 1;-2 0 1 0]
B = [0,0,0,1]

uk = rand(np)#zeros(np)

function f(x,u)
    return A*x + B*u
end

function S_x(u)
    function forward(x, p, t)
        return f(x,linear_interpolation(0.0:dt:T,p)(t))
    end
    u0 = [1.0, 0.0, -1.0, 0.0]
    tspan = (0.0, T)
    prob = ODEProblem(forward, u0, tspan, u)
    sol = solve(prob,Midpoint(),saveat = T/N)
    return Array(sol)
end

S_x(uk)
ForwardDiff.jacobian(S_x,uk)
Zygote.jacobian(S_x,uk)