-
Notifications
You must be signed in to change notification settings - Fork 82
Closed
Labels
Milestone
Description
Code
import Pkg; Pkg.status()
import Random
import Enzyme
const AV = AbstractVector{T} where T
# ===== Set up objective function =====
normal_pdf(x::Real, mean::Real, var::Real) =
exp(-(x - mean)^2 / (2var)) / sqrt(2π * var)
function mixture_loglikelihood(params::AV{<:Real}, data::AV{<:Real})::Real
K = length(params) ÷ 3
weights, means, stds = @views params[1:K], params[(K+1):2K], params[(2K+1):end]
sum(
sum(
weight * normal_pdf(x, mean, std)
for (weight, mean, std) in zip(weights, means, stds)
) |> log
for x in data
)
end
rnd = Random.MersenneTwister(42)
data = randn(rnd, 500)
params0 = [[0.2, 0.8]; [-1, 1]; [.3, .6]]
objective = params -> mixture_loglikelihood(params, data)
# ===== Compute gradient =====
# 1. From docs (https://docs.juliahub.com/Enzyme/G1p5n/0.10.4/api/#Enzyme.gradient!-Tuple{Enzyme.ReverseMode,%20Any,%20Any,%20Any})
let
f(x) = x[1]*x[2]
dx = [0.0, 0.0]
Enzyme.gradient!(Enzyme.Reverse, dx, f, [2.0, 3.0])
@show dx
end
# 2. My function
let
grad_storage = similar(params0)
Enzyme.gradient!(Enzyme.Reverse, grad_storage, objective, params0)
@show grad_storage
endOutput
Status `~/test/autodiff_bench/Project.toml`
[6e4b80f9] BenchmarkTools v1.3.1
[7da242da] Enzyme v0.10.4
[f6369f11] ForwardDiff v0.10.32
[37e2e3b7] ReverseDiff v1.14.1
[e88e6eb3] Zygote v0.6.43
dx = [3.0, 2.0]
ERROR: LoadError: UndefRefError: access to undefined reference
Stacktrace:
[1] unsafe_convert
@ ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4277 [inlined]
[2] macro expansion
@ ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4520 [inlined]
[3] enzyme_call(::Ptr{Nothing}, ::Type{Enzyme.Compiler.AugmentedForwardThunk}, ::Type{Val{0x0000000000000001}}, ::Val{true}, ::Type{Tuple{Enzyme.Duplicated{Vector{Float64}}, Enzyme.Const{Vector{Float64}}}}, ::Type{Enzyme.Const{Real}}, ::typeof(mixture_loglikelihood), ::Nothing, ::Enzyme.Duplicated{Vector{Float64}}, ::Enzyme.Const{Vector{Float64}})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4316
[4] (::Enzyme.Compiler.AugmentedForwardThunk{typeof(mixture_loglikelihood), Enzyme.Const{Real}, Tuple{Enzyme.Duplicated{Vector{Float64}}, Enzyme.Const{Vector{Float64}}}, Val{0x0000000000000001}, Nothing, Val{true}()})(::Enzyme.Duplicated{Vector{Float64}}, ::Vararg{Any})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4307
in expression starting at /Users/forcebru/test/autodiff_bench/Enzyme_bug.jl:39
Other autodiff packages (I tried ForwardDiff, ReverseDiff and Zygote) don't have any issues differentiating objective.
Versions
- Enzyme v0.10.4
julia> versioninfo()
Julia Version 1.8.0-rc3
Commit 33f19bcbd25 (2022-07-13 19:10 UTC)
Platform Info:
OS: macOS (x86_64-apple-darwin21.4.0)
CPU: 4 × Intel(R) Core(TM) i5-3330S CPU @ 2.70GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-13.0.1 (ORCJIT, ivybridge)
Threads: 1 on 4 virtual cores