# Enzyme.jl cannot differentiate through vector of distributions?

**URL:** <https://discourse.julialang.org/t/enzyme-jl-cannot-differentiate-through-vector-of-distributions/114338>\
**Category:** General Usage\
**Tags:** question\
**Created:** [May 16, 2024, 4:50am UTC](https://discourse.julialang.org/t/enzyme-jl-cannot-differentiate-through-vector-of-distributions/114338 "2024-05-16T04:50:09Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![biona001](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/biona001/32/16497_2.png) [@biona001](https://discourse.julialang.org/u/biona001)\
**Post date:** [May 16, 2024, 4:50am UTC](https://discourse.julialang.org/t/enzyme-jl-cannot-differentiate-through-vector-of-distributions/114338/1 "2024-05-16T04:50:09Z")

</div>

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:

```julia
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

```julia
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

---

<div class="post-metadata">

**Author:** ![biona001](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/biona001/32/16497_2.png) [@biona001](https://discourse.julialang.org/u/biona001)\
**Post date:** [May 16, 2024, 4:33pm UTC](https://discourse.julialang.org/t/enzyme-jl-cannot-differentiate-through-vector-of-distributions/114338/2 "2024-05-16T16:33:13Z")

</div>

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

```julia
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

```

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [May 17, 2024, 6:55am UTC](https://discourse.julialang.org/t/enzyme-jl-cannot-differentiate-through-vector-of-distributions/114338/3 "2024-05-17T06:55:45Z")

</div>

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

```julia
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](https://github.com/EnzymeAD/Enzyme.jl/pull/1449)). Of course for perf/etc it is better to be type stable, but for the future at least it’s good to support.