Closed
Description
Repro:
module MWE
using DynamicPPL
using Distributions
using Mooncake
using ADTypes
using LogDensityProblems
@model function lda(K, d, w)
V = length(unique(w))
D = length(unique(d))
N = length(d)
@assert length(w) == N
ϕ = Vector{Vector{Real}}(undef, K)
for i in 1:K
ϕ[i] ~ Dirichlet(ones(V) / V)
end
θ = Vector{Vector{Real}}(undef, D)
for i in 1:D
θ[i] ~ Dirichlet(ones(K) / K)
end
z = zeros(Int, N)
for i in 1:N
z[i] ~ Categorical(θ[d[i]])
w[i] ~ Categorical(ϕ[d[i]])
end
return (; ϕ=ϕ, θ=θ, z=z)
end
model = begin
w = [1, 2, 3, 2, 1, 1]
d = [1, 1, 1, 2, 2, 2]
lda(2, d, w)
end
vi = DynamicPPL.link(VarInfo(model), model)
x = randn(length(vi[:]))
ldf_mooncake = LogDensityFunction(
model, vi, DefaultContext(); adtype=AutoMooncake(; config=nothing)
)
LogDensityProblems.logdensity_and_gradient(ldf_mooncake, x)
endERROR: LoadError: UndefRefError: access to undefined reference
Stacktrace:
[1] getindex
@ ./essentials.jl:13 [inlined]
[2] getindex
@ ./subarray.jl:323 [inlined]
[3] _broadcast_getindex
@ ./broadcast.jl:675 [inlined]
[4] _getindex
@ ./broadcast.jl:705 [inlined]
[5] _broadcast_getindex
@ ./broadcast.jl:681 [inlined]
[6] getindex
@ ./broadcast.jl:636 [inlined]
[7] macro expansion
@ ./broadcast.jl:1004 [inlined]
[8] macro expansion
@ ./simdloop.jl:77 [inlined]
[9] copyto!
@ ./broadcast.jl:1003 [inlined]
[10] copyto!
@ ./broadcast.jl:956 [inlined]
[11] materialize!
@ ./broadcast.jl:914 [inlined]
[12] materialize!
@ ./broadcast.jl:911 [inlined]
[13] (::Mooncake.var"#unsafe_copyto_pb!!#638"{…})(::Mooncake.NoRData)
@ Mooncake ~/.julia/packages/Mooncake/hkc4e/src/rrules/array_legacy.jl:265
[14] _copyto_impl!
@ ./array.jl:372 [inlined]
[15] (::Tuple{…})(none::Mooncake.NoRData)
@ Base.Experimental ./<missing>:0
[16] Pullback
@ ~/.julia/packages/Mooncake/hkc4e/src/interpreter/s2s_reverse_mode_ad.jl:916 [inlined]
[17] _setindex
@ ~/.julia/packages/BangBang/VnHxZ/src/NoBang/base.jl:130 [inlined]
[18] (::Tuple{…})(none::Mooncake.NoRData)
@ Base.Experimental ./<missing>:0
[19] Pullback
@ ~/.julia/packages/Mooncake/hkc4e/src/interpreter/s2s_reverse_mode_ad.jl:916 [inlined]
[20] lda
@ /tmp/mwe2.jl:9 [inlined]
[21] (::Tuple{…})(none::Union{…})
@ Base.Experimental ./<missing>:0
[22] Pullback
@ ~/.julia/packages/Mooncake/hkc4e/src/interpreter/s2s_reverse_mode_ad.jl:916 [inlined]
[23] (::Mooncake.RRuleWrapperPb{Mooncake.Pullback{…}, Mooncake.LazyZeroRData{…}})(dy::Mooncake.NoRData)
@ Mooncake ~/.julia/packages/Mooncake/hkc4e/src/interpreter/s2s_reverse_mode_ad.jl:299
[24] logdensity_at
@ ~/.julia/packages/DynamicPPL/3OIqq/src/logdensityfunction.jl:183 [inlined]
[25] (::Tuple{…})(none::Float64)
@ Base.Experimental ./<missing>:0
[26] Pullback
@ ~/.julia/packages/Mooncake/hkc4e/src/interpreter/s2s_reverse_mode_ad.jl:916 [inlined]
[27] __value_and_gradient!!(::Mooncake.DerivedRule{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
@ Mooncake ~/.julia/packages/Mooncake/hkc4e/src/interface.jl:85
[28] value_and_gradient!!
@ ~/.julia/packages/Mooncake/hkc4e/src/interface.jl:300 [inlined]
[29] value_and_gradient(::typeof(DynamicPPL.logdensity_at), ::DifferentiationInterfaceMooncakeExt.MooncakeGradientPrep{…}, ::AutoMooncake{…}, ::Vector{…}, ::Constant{…}, ::Constant{…}, ::Constant{…})
@ DifferentiationInterfaceMooncakeExt ~/.julia/packages/DifferentiationInterface/qrWdQ/ext/DifferentiationInterfaceMooncakeExt/onearg.jl:130
[30] prepare_gradient(::typeof(DynamicPPL.logdensity_at), ::AutoMooncake{…}, ::Vector{…}, ::Constant{…}, ::Constant{…}, ::Constant{…})
@ DifferentiationInterfaceMooncakeExt ~/.julia/packages/DifferentiationInterface/qrWdQ/ext/DifferentiationInterfaceMooncakeExt/onearg.jl:123
[31] _
@ ~/.julia/packages/DynamicPPL/3OIqq/src/logdensityfunction.jl:131 [inlined]
[32] top-level scope
@ /tmp/mwe2.jl:42
[33] include(fname::String)
@ Base.MainInclude ./client.jl:494
[34] top-level scope
@ REPL[34]:1
julia> versioninfo()
Julia Version 1.10.8
Commit 4c16ff44be8 (2025-01-22 10:06 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: macOS (arm64-apple-darwin24.0.0)
CPU: 10 × Apple M1 Pro
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)
Threads: 1 default, 0 interactive, 1 GC (on 8 virtual cores)
Mooncake v0.4.103