Skip to content

Runtime activity required for invoke #407

@ForceBru

Description

@ForceBru

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
end

Output

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions