Enzyme.jl cannot differentiate through vector of distributions?

General Usage
1

I have a vector y which is a Vector{Float64}, but each y[i] is sampled from a different marginal distribution specified by vecdist[i].

I want the gradient of the loglikelihood function (scalar output), but I’m getting a strange looking Mismatched activity error. I’m not sure if I’m misusing `Enzyme.jl somehow? Any help is highly appreciated.

MWE:

using Enzyme
using Distributions
using GLM
using Random
import GLM.loglik_obs

struct MultiResponse{T, D, L}
    y::Vector{T} # d by 1 vector
    vecdist::Vector{D} # length d vector of marginal distributions, one for each y[i]
    veclink::Vector{L} # length d vector of link functions, one for each y[i]
end

function component_loglikelihood(data::MultiResponse, η::Vector)
    logl = 0.0
    for j in eachindex(data.y)
        dist = data.vecdist[j]
        link = data.veclink[j]
        μ_j = GLM.linkinv(link, η[j])
        logl += loglik_obs(dist, y[j], μ_j, 1.0, 1.0)
    end
    return logl::Float64 # type annotation prevents "Duplicate return not supported" error
end

# simulate data
d = 10
possible_distributions = [Bernoulli(), Poisson(), Normal()]
vecdist = rand(possible_distributions, d)
veclink = [canonicallink(vecdist[j]) for j in 1:d]
y = zeros(d)
for j in 1:d
    dist = vecdist[j]
    y[j] = rand(dist)
end
data = MultiResponse(y, vecdist, veclink)

# eval obj
η = randn(d)
component_loglikelihood(data, η) # -16.602324713991795

# compute grad with Enzyme.jl
grad_storage = zeros(length(η))
Enzyme.autodiff(
    Reverse, component_loglikelihood,
    Const(data),
    Duplicated(η, zero(η))
)

the last step causes the following error

Enzyme execution failed.
Mismatched activity for:   %value_phi6.i = phi {} addrspace(10)* [ addrspacecast ({}* inttoptr (i64 5306670080 to {}*) to {} addrspace(10)*), %entry ], [ %19, %L26.i.loopexit ] const val: {} addrspace(10)* addrspacecast ({}* inttoptr (i64 5306670080 to {}*) to {} addrspace(10)*)
 value=0.0 of type Float64
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#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] component_loglikelihood
   @ ./In[1]:21
 [2] component_loglikelihood
   @ ./In[1]:0


Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/srACB/src/compiler.jl:1325
  [2] component_loglikelihood
    @ ./In[1]:15 [inlined]
  [3] component_loglikelihood
    @ ./In[1]:0 [inlined]
  [4] diffejulia_component_loglikelihood_2266_inner_1wrap
    @ ./In[1]:0
  [5] macro expansion
    @ ~/.julia/packages/Enzyme/srACB/src/compiler.jl:5719 [inlined]
  [6] enzyme_call
    @ ~/.julia/packages/Enzyme/srACB/src/compiler.jl:5385 [inlined]
  [7] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/srACB/src/compiler.jl:5264 [inlined]
  [8] autodiff
    @ ~/.julia/packages/Enzyme/srACB/src/Enzyme.jl:291 [inlined]
  [9] autodiff
    @ ~/.julia/packages/Enzyme/srACB/src/Enzyme.jl:315 [inlined]
 [10] autodiff(::ReverseMode{false, FFIABI, false}, ::typeof(component_loglikelihood), ::Const{MultiResponse{Float64, UnivariateDistribution, Link}}, ::Duplicated{Vector{Float64}})
    @ Enzyme ~/.julia/packages/Enzyme/srACB/src/Enzyme.jl:300
 [11] top-level scope
    @ In[2]:3

This is on julia v1.10.3 and Enzyme v0.12.6

2

After some experimentation, I think GLM.loglik_obs is the culprit. If I annotated GLM.loglik_obs()::Float64, I will get no error.

The following now works

using Enzyme
using Distributions
using GLM
using Random
import GLM.loglik_obs

struct MultiResponse{T, D, L}
    y::Vector{T} # d by 1 vector
    vecdist::Vector{D} # length d vector of marginal distributions, one for each y[i]
    veclink::Vector{L} # length d vector of link functions, one for each y[i]
end

function component_loglikelihood(data::MultiResponse, η::Vector)
    logl = 0.0
    for j in eachindex(data.y)
        dist = data.vecdist[j]
        link = data.veclink[j]
        μ_j = GLM.linkinv(link, η[j])
        logl += loglik_obs(dist, y[j], μ_j, 1.0, 1.0)::Float64
    end
    return logl
end

# simulate data
d = 10
possible_distributions = [Bernoulli(), Poisson(), Normal()]
vecdist = rand(possible_distributions, d)
veclink = [canonicallink(vecdist[j]) for j in 1:d]
y = zeros(d)
for j in 1:d
    dist = vecdist[j]
    y[j] = rand(dist)
end
data = MultiResponse(y, vecdist, veclink)

# eval obj
η = randn(d)
component_loglikelihood(data, η) # -16.602324713991795

# compute grad with Enzyme.jl
grad_storage = zeros(length(η))
Enzyme.autodiff(
    Reverse, component_loglikelihood,
    Const(data),
    Duplicated(η, grad_storage)
)

julia> grad_storage
10-element Vector{Float64}:
  0.7817690684853906
 -2.271274822893419
  0.5252750764625924
  2.187566486142698
 -0.3231934317570125
 -0.09679141622994089
  1.3365174072378545
 -0.6137291517755701
  0.7146628592512362
  1.4827989283231573
1 Like
3

Yeah looks like the result of loglik_obs is type unstable which is why you’re getting these errors. In the first case, a type unstable return looks like an Any, so it’s assumed duplicated. Your fix making a type stable return is fine. Alternatively you could explicitly specify the activity of the return to be active like below

Enzyme.autodiff(
    Reverse, component_loglikelihood,
    Active
    Const(data),
    Duplicated(η, zero(η))
)

The second error is that your logl variable is type unstable and activity unstable (an unfortunate combo). Type unstable for the return, activity unstable since from the start of the loop it is a constant 0.0, and inside the loop it is a differentiable variable.

However, in this case we should be able to handle it and i’ve made a fix here which will resolve without you needing to mark it as type stable (Handle mixed activity of literal 0 constant by wsmoses · Pull Request #1449 · EnzymeAD/Enzyme.jl · GitHub). Of course for perf/etc it is better to be type stable, but for the future at least it’s good to support.

1 Like