Turing.jl + Enyzme Examples

Specific Domains Probabilistic programming
1

Hi there, are there any working examples of using Turing.jl with the Enzyme.jl AD backend?

I’ve seen several Github issues like this one but I think they may be outdated at this point. Also the following Discourse topic where it looks like Enzyme is ready for everyday use with Turing: Enzyme ready for everyday use? (2024). However, I didn’t see an example referenced in that post.

Any info would be appreciated.

Thanks!
Josh

2 Likes
3

Hi @joshualeond! The Turing.jl PR you linked is still being worked on, it just goes in bursts of activity. Many, many Turing models work with Enzyme, but it’s not yet at a point where I would have great confidence that any given model you throw at it will be fine. Definitely worth a try though, to see how your use case fares. We are in the process of running more of our test suite with Enzyme, reporting any issues that come up.

Usage is as simple as

import ADTypes
import Enzyme
using Turing

Enzyme.API.runtimeActivity!(true)

@model function m()
    x ~ Normal()
    y ~ Normal(x, 1)
end

model = m()
sampler = NUTS(; adtype=ADTypes.AutoEnzyme(; mode=Enzyme.Forward))
sample(model, sampler, 1000)
1 Like
4

I’ll note that the above code uses an old version of Enzyme (I think Turing needs to fix some compat’s).

2 Likes