Using intermediate results in conditional Gibbs in Turing

I have a problem with an expensive intermediate computation that I have to use both within the normal model definition, and in a conditional Gibbs update - is there some way that I can reuse the computation directly in the conditional update?

Relevant subsections of code are (M is a potentially expensive operation)

        #this is inside the model definition
        μ = μ₀ .* exp.(M(l, σ, w))
        μi = S*μ
        Turing.@addlogprob! sum(log.(μi[x.==1])) - hatsum(μ)*params.M.h


    function cond_x_b_c(c, b::BranchingProcess{T}, catalog::Catalog{T}, M::MaternSPDE) where T
        μ = c.μ₀ .* exp.(M(c.l, c.σ, c.w))
        μi = S*μ
        update_weights!(b, catalog, μi, c.K, c.α, c.c, c.p̃+1)
        return Product([Categorical(b.bnodes[i].bweights) for i = 1:length(catalog.t)])
    cond_x(c) = cond_x_b_c(c, b, catalog, params.M)

where I compute the same deterministic intermediate value μi.

Would something like Memoization help your problem? The Turing docs have an example here.

1 Like