Info about compatible AD backends and parallel chains

[itsdfish]

Hello,

A collegue pointed me to your package and it seems like it might be a useful method for speeding up parameter estimation for some of our models. Thank you for putting this together.

I have had some problems with version capatibility (see my previous comments) and AD requirements. For example, I received the following error with the example in the docs:

ERROR: Turing doesn't support using the Zygote backend when Threads.nthreads()>1. Use a different backend or a single-thread.

I was able to run the example after running an instance of Julia with 1 thread. Is this a strict requirement?

To circumvent this problem, I tried using ReverseDiff but encountered the following error:

Unsupposed backend from Turing: Base.RefValue{Symbol}(:reversediff)

It would be nice to use ReverseDiff because its performance with NUTS is similar to the performance of zygote with MuseInference.

Can you please provide some clarification on the AD requirements and parallel capabilities?

Thanks!

[marius311]

The Turing/Zygote/1-thread thing is unrelated to MuseInference. Last I checked, that didn't work there either.

On ReverseDiff, could you provide a MWE? It would indeed be good to have ReverseDiff working.

[itsdfish]

Thank you for clarifying about Zygote.

Thanks for looking into support for ReverseDIff.

I modified your example from the documentation for ReverseDiff. Here is the code:

using AbstractDifferentiation
using LinearAlgebra
using MuseInference
using Random
using ReverseDiff
using Turing

Turing.setadbackend(:reversediff)

@model function funnel()
    θ ~ Normal(0, 3)
    z ~ MvNormal(zeros(2048), exp(θ)*I)
    x ~ MvNormal(z, I)
end

Random.seed!(1)
(;x) = rand(funnel() | (θ=0,))
model = funnel() | (;x)

Random.seed!(2)
chain = @time sample(model, NUTS(100, 0.65, init_ϵ=0.5), 500);

nsims = round(Int, ess(chain)[:θ,:ess])

Random.seed!(3)
muse_result = @time muse(model, 0; nsims, get_covariance=true)

Here is the error message:

ERROR: Unsupposed backend from Turing: Base.RefValue{Symbol}(:reversediff)
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] TuringMuseProblem(model::DynamicPPL.Model{…}; params::Tuple{…}, autodiff::Nothing)
   @ MuseInference ~/.julia/packages/MuseInference/jjAab/src/turing.jl:125
 [3] TuringMuseProblem
   @ ~/.julia/packages/MuseInference/jjAab/src/turing.jl:112 [inlined]
 [4] #muse!#132
   @ ~/.julia/packages/MuseInference/jjAab/src/turing.jl:249 [inlined]
 [5] muse(::DynamicPPL.Model{…}, ::Vararg{…}; kwargs::@Kwargs{…})
   @ MuseInference ~/.julia/packages/MuseInference/jjAab/src/muse.jl:107
 [6] top-level scope
   @ ./timing.jl:279 [inlined]
 [7] top-level scope
   @ ~/.julia/dev/sandbox/muse_inference/basic_example.jl:0
Some type information was truncated. Use `show(err)` to see complete types.

Version Info:

Julia 1.10.4

⌅ [c29ec348] AbstractDifferentiation v0.5.3
  [b964fa9f] LaTeXStrings v1.3.1
  [43b88160] MuseInference v0.2.4
  [37e2e3b7] ReverseDiff v1.15.3
⌅ [fce5fe82] Turing v0.28.3
  [e88e6eb3] Zygote v0.6.70
  [37e2e46d] LinearAlgebra
  [9a3f8284] Random