ForwardDiff.dual TypeError while Solving ODE with 3 by 3 Matrix Input and ProximalOperators Package which Keep the Output Positive Definite

Hi! I write a Turing ODE code as follows:

The data is u_account = [953.0; 1075.0; 1223.0] which represents three states [A,B,C] being observed.

The parameter p consists of variables [v1,v2,v4,u1,u3,w1]

Vector u_mean denotes the mean of initial states of [A,B,C]

The ODE function is using to calculate the variance part for multivariable distribution.

The ODE gets input 1e-16*Identity matrix and then the output of ODE called phi_new is modified with ProximalOperators package in order to force to be positive definite.

Finally use MvNormal distribution(u_initial,phi_new) to infer the observed data.

This leads to the typeassert error message :expected Float64, got a value of type ForwardDiff.Dual{Nothing,Float64,9}

The issue might causes by the line
phi_new, _ = prox(IndPSD(), Symmetric(convert(Array{Float64},phi_new)))
since the varible phi_new becomes ForwardDiff.dual type after the Turing runs and fucntion prox() from ProximalOperators package requires it input matrix to be float64 type. Here, even I explictly apply the code convert(Array{Float64,2},phi_new)), phi_new could not be changed to Float64 from ForwardDiff.dual type.

Could you help me with this? Thanks a lot in advance!

using Random
using LinearAlgebra
using DifferentialEquations
using StatsPlots
using Plots
using DiffEqSensitivity
using LinearAlgebra
using Statistics
using Distributions
using Turing
using PDMats
using ProximalOperators

u_account = [953.0; 1075.0; 1223.0];

@model bayes_lna(y) = begin
    logv1 ~ Normal(0.0,1.0) 
    logv2 ~ Normal(0.0,1.0) 
    logv4 ~ Normal(0.0,1.0)
    logu1 ~ Normal(0.0,1.0)
    logu3 ~ Normal(0.0,1.0)
    logw1 ~ Normal(0.0,1.0)

    p = [exp(logv1),exp(logv2),exp(logv4),
    A_mean ~ Normal(7.0,1.0)
    B_mean ~ Normal(7.0,1.0)
    C_mean ~ Normal(7.0,1.0)
    u_mean = [exp(A_mean),exp(B_mean),exp(C_mean)]
    u_initial = eltype(u_mean).(zeros(3))
    tspan = (0.0,2.0)
    ita0 = u_initial
    d_phi(u,p,t) = [p[1]-p[3]*ita0[3]     0     -p[3]*ita0[1];
                  p[2]+p[3]*ita0[3]   -p[5]      p[3]*ita0[1];
                           0         p[4]        -p[6] ]*u
    phi_0 = zeros(3,3)
    phi_0[1,1] = 1e-16
    phi_0[2,2] = 1e-16
    phi_0[3,3] = 1e-16
    sol  = solve(ODEProblem(d_phi,eltype(p).(phi_0),tspan,p))
    phi_new = sol[end] 
    phi_new, _ = prox(IndPSD(), Symmetric(convert(Array{Float64,2},phi_new)))
    phi_new = phi_new + 1e-12*Matrix(I, 3, 3)
    y[:] ~ MvNormal(ita0,phi_new)

iterations = 20
chain = sample(bayes_lna(u_account), NUTS(0.65),iterations)
TypeError: in typeassert, expected Float64, got a value of type ForwardDiff.Dual{Nothing,Float64,9}

It looks like you can’t differentiate through ProximalOperators. I’d open an issue there about trying to do that.