Models with persistent memory usage

I am using Turing a lot but now I have a model which would keep some persistent memory to be computational efficient. A minimal example would be:

using Turing, Optim, ForwardDiff

mutable struct Acomputer{T <: Number}
    a::Vector{T} # <- in real I would need more 
    para::T
end

function set_para!(ac::Acomputer, p)
    ac.para = p
end

function get_a!(ac::Acomputer)
    fill!(ac.a, ac.para)  # <- in real this is more involved
    ac.a
end

the model would be defined like this:

@model function m(ac, x)
    para ~ Uniform(0,10)

    set_para!(ac, para)
    a = get_a!(ac)
    x .~ Poisson.(a)
end

I can now optimise it w/o the use of gradients:

ac = Acomputer(zeros(2), 3.0)
optimize(m(ac, [4,6]), MLE(), NelderMead())

giving the expected result:

ModeResult with maximized lp of -3.66
1-element Named Vector{Float64}
A     │ 
──────┼────────
:para │ 5.00009

However - if I want to make use AD, i.e. ForwardDiff, I cannot do this anymore:

ac = Acomputer{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 1}}(zeros(2), 3.0)
optimize(m(ac, [4,6]), MLE())

I get

MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 1})
Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat at rounding.jl:200
  (::Type{T})(::T) where T<:Number at boot.jl:772
  (::Type{T})(::AbstractChar) where T<:Union{AbstractChar, Number} at char.jl:50
  ...

… and don’t know how to resolve this issue. Is there any help?

First off, be careful here. I don’t know what you’re doing, but if you start sampling with this, you’re likely breaking the Markov assumption which is essential to Markov Chain Monte Carlo (MCMC).

But with that being said, the following seems to be the problem:

ac = Acomputer{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 1}}(zeros(2), 3.0)

You’re telling Acomputer that it should expect Dual numbers, but then you give it zeros(2)::Vector{Float64} and 3.0 (which is a Float64). This is causing the issue.

Not really - if I just run

ac = Acomputer{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 1}}(zeros(2), 3.0)

then it ac will be of the correct type, s.t. all the internal parameters are of the correct type, i.e. Dual numbers. As far as I understand this is prerequisite to make AD possible. The above error only occurs when I want to optimize the model.

But probably there is also another way to achieve what I want to do.

In principle: My model is computationally expansive and allocates a lot of memory. I want to circumvent that subsequent model calls always (re)allocate this memory.

Okay so then I think I need the full traceback :confused: