Enzyme execution failed on a basic example?

[biona001]

Hi community,

I am trying to include Enzyme.jl to this benchmark which compares the speed of ForwardDiff.jl, ReverseDiff.jl, Symbolics.jl, Zygote.jl to that of JAX for computing the gradient of a relatively simple loglikelihood function. My hope is that Enzyme.jl would be as fast as JAX or at least faster than all other AD packages for this problem. The actual problem I'm trying to differentiate is a much more complicated loglikelihood function. I ran the code below on

• Julia v1.9.1
• Enzyme.jl v0.11.17

The issue is I am getting a 2 different errors depending on how I write the objective (none of them work). MWE:

using Random
using Enzyme
using LinearAlgebra
using BenchmarkTools

# ========== Benchmark setup ==========
SEED = 42
N_SAMPLES = 500
N_COMPONENTS = 4

rnd = Random.MersenneTwister(SEED)
data = randn(rnd, N_SAMPLES)
params0 = [rand(rnd, N_COMPONENTS); randn(rnd, N_COMPONENTS); 2rand(rnd, N_COMPONENTS)]

# ========== Objective function ==========
normal_pdf(x::Real, mean::Real, var::Real) =
    exp(-(x - mean)^2 / (2var)) / sqrt(2π * var)

normal_pdf(x, mean, var) =
    exp(-(x - mean)^2 / (2var)) / sqrt(2π * var)

# original objective (doesn't work)
function mixture_loglikelihood1(params::AbstractVector{<:Real}, data::AbstractVector{<:Real})::Real
    K = length(params) ÷ 3
    weights, means, stds = @views params[1:K], params[K+1:2K], params[2K+1:end]
    mat = normal_pdf.(data, means', stds' .^2) # (N, K)
    sum(mat .* weights', dims=2) .|> log |> sum
end

# another form of original objective (doesn't work)
function mixture_loglikelihood2(params::AbstractVector{<:Real}, data::AbstractVector{<:Real})::Real
    K = length(params) ÷ 3
    weights, means, stds = @views params[1:K], params[K+1:2K], params[2K+1:end]
    mat = normal_pdf.(data, means', stds' .^2) # (N, K)
    obj_true = sum(
        sum(
            weight * normal_pdf(x, mean, std^2)
            for (weight, mean, std) in zip(weights, means, stds)
        ) |> log
        for x in data
    )
end

# objective re-written by me
function mixture_loglikelihood3(params::AbstractVector{<:Real}, data::AbstractVector{<:Real})::Real
    K = length(params) ÷ 3
    weights, means, stds = @views params[1:K], params[K+1:2K], params[2K+1:end]
    mat = normal_pdf.(data, means', stds' .^2) # (N, K)

    obj = zero(eltype(mat))
    for x in data
        obj_i = zero(eltype(mat))
        for (weight, mean, std) in zip(weights, means, stds)
            obj_i += weight * normal_pdf(x, mean, std^2)
        end
        obj += log(obj_i)
    end
    return obj
end

objective1 = params -> mixture_loglikelihood1(params, data)
objective2 = params -> mixture_loglikelihood2(params, data)
objective3 = params -> mixture_loglikelihood3(params, data)

@show objective1(params0) # -443.4039737200718
@show objective2(params0) # -443.4039737200718
@show objective3(params0) # -443.4039737200718

Enzyme.jl on objective1:

julia> bparams0 = zeros(length(params0))
julia> Enzyme.autodiff(Reverse, objective1, Active, Duplicated(params0, bparams0))

BoundsError: attempt to access Core.SimpleVector at index [1]

julia> Enzyme.gradient(Reverse, objective1, params0)

BoundsError: attempt to access Core.SimpleVector at index [1]

Enzyme.jl on objective2:

julia> bparams0 = zeros(length(params0))
julia> Enzyme.autodiff(Reverse, objective2, Active, Duplicated(params0, bparams0))

Enzyme execution failed.
Mismatched activity for:   store {} addrspace(10)* %1, {} addrspace(10)* addrspace(10)* %.repack170, align 8, !dbg !415, !tbaa !174, !alias.scope !178, !noalias !179 const val: {} addrspace(10)* %1
 value=Unknown object of type Vector{Float64}
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] sum
   @ ./reduce.jl:559
 [2] mixture_loglikelihood2
   @ ./In[47]:16

Stacktrace:
 [1] throwerr(cstr::Cstring)
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/qRjan/src/compiler.jl:1288

julia> Enzyme.gradient(Reverse, objective2, params0)

Enzyme execution failed.
Mismatched activity for:   store {} addrspace(10)* %1, {} addrspace(10)* addrspace(10)* %.repack170, align 8, !dbg !415, !tbaa !174, !alias.scope !178, !noalias !179 const val: {} addrspace(10)* %1
 value=Unknown object of type Vector{Float64}
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] sum
   @ ./reduce.jl:559
 [2] mixture_loglikelihood2
   @ ./In[47]:16

Stacktrace:
 [1] throwerr(cstr::Cstring)
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/qRjan/src/compiler.jl:1288

Including Enzyme.API.runtimeActivity!(true) does not change anything.

Enzyme.jl on objective3:

julia> bparams0 = zeros(length(params0))
julia> Enzyme.autodiff(Reverse, objective3, Active, Duplicated(params0, bparams0))

BoundsError: attempt to access Core.SimpleVector at index [1]

julia> Enzyme.gradient(Reverse, objective3, params0)

BoundsError: attempt to access Core.SimpleVector at index [1]

Any tips/suggestions would be highly appreciated.

[biona001]

Well, I took out the type annotations on the objective functions, then both objective1, and objective3 worked. objective2 still throws the same error, but I guess it's not that important now. Closing this

function mixture_loglikelihood1(params::AbstractVector, data::AbstractVector)
    K = length(params) ÷ 3
    weights, means, stds = @views params[1:K], params[K+1:2K], params[2K+1:end]
    mat = normal_pdf.(data, means', stds' .^2) # (N, K)

    sum(mat .* weights', dims=2) .|> log |> sum
end

function mixture_loglikelihood2(params::AbstractVector, data::AbstractVector)
    K = length(params) ÷ 3
    weights, means, stds = @views params[1:K], params[K+1:2K], params[2K+1:end]
    mat = normal_pdf.(data, means', stds' .^2) # (N, K)

    obj_true = sum(
        sum(
            weight * normal_pdf(x, mean, std^2)
            for (weight, mean, std) in zip(weights, means, stds)
        ) |> log
        for x in data
    )
end

function mixture_loglikelihood3(params::AbstractVector, data::AbstractVector)
    K = length(params) ÷ 3
    weights, means, stds = @views params[1:K], params[K+1:2K], params[2K+1:end]
    mat = normal_pdf.(data, means', stds' .^2) # (N, K)

    # objective re-written by me
    obj = zero(eltype(mat))
    for x in data
        obj_i = zero(eltype(mat))
        for (weight, mean, std) in zip(weights, means, stds)
            obj_i += weight * normal_pdf(x, mean, std^2)
        end
        obj += log(obj_i)
    end
    return obj
end

objective1 = params -> mixture_loglikelihood1(params, data)
objective2 = params -> mixture_loglikelihood2(params, data)
objective3 = params -> mixture_loglikelihood3(params, data)

[wsmoses]

Reopening this since the bounds error should never occur for sure.

[biona001]

By the way, with N_SAMPLES = 10000 and N_COMPONENTS = 5, I'm seeing the following time

• Symbolics.jl: 24.025 ms
• ForwardDiff.jl: 11.848 ms
• ReverseDiff.jl: 96.394 ms
• Zygote.jl (reverse): 4.797 ms
• Enzyme.jl (reverse mode): 3.245 ms
• JAX (python): 2.515 ms

So it seems enzymes.jl is the fastest autodiff package in Julia as far as I can tell, but still ~30% slower than JAX (python).

[wsmoses]

Bounds error now fixed, closing