Enzyme.jl
Enzyme.jl copied to clipboard
EnzymeAD
Enzyme execution failed on a basic example?
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.
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)
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).