Curious slow down (through Turing)

[DominiqueMakowski]

I know that Enzyme is not yet fully supported by Turing, so apologies if this issue is mistimed.

While Enzyme boosts efficiency with a simple linear model, I noticed that adding a simple variable transformation (that computes the polynomials) within the model critically slows down turning (but not the default AD).

using Distributions
using Enzyme
using Turing  
using ADTypes

Enzyme.API.runtimeActivity!(true)

function data_poly(x, degree=2)
    return hcat([x .^ d for d in 1:degree]...)
end

@model function model(y, x)

    X = data_poly(x, 2)

    # Priors
    μ_intercept ~ Normal(0, 3)
    μ_x1 ~ Normal(0, 3)
    μ_x2 ~ Normal(0, 3)
    σ ~ Normal(0, 3)

    # Likelihood
    for i in eachindex(y)
        μ = μ_intercept + μ_x1 * X[i, 1] + μ_x2 * X[i, 2]
        y[i] ~ Normal(μ, exp(σ))
    end
end

y = rand(Normal(0, 1), 1000)
x = y .+ rand(Normal(0, 1.5), 1000)

sample(model(y, x), NUTS(), 1000)  # 5.7 seconds
sample(model(y, x), NUTS(; adtype = AutoEnzyme()), 1000)  # ~6min

Any pointers to what I'm doing wrong?

[wsmoses]

If you run sample with enzyme a second time, is it fast or slow (I want to know whether that's compilation time or runtime)?

[DominiqueMakowski]

It doesn't seem like it: the first time took 595.75 seconds, and the second 647.95 seconds (as compared to 13 sec this time on default AD)

[yebai]

This might be related: https://github.com/JuliaLang/julia/issues/55638

@willtebbutt observed a significant slowdown due to the above issue on some Turing models when using Tapir.