InexactError with loglikelihood(model, chain) on multiple threads when models have parameters of different types

[jerlich]

I am trying to fit some Mixed Multinomial models. The models sample fine, but when computing the loglikelihood, i get an inexact error.

MWE:

using Turing, FillArrays, NNlib

ij = [[1, 2], [1, 3], [1, 4], [1, 5], [1, 6], [2, 3], [2, 4], [2, 5], [2, 6], [3, 4], [3, 5], [3, 6], [4, 5], [4, 6], [5, 6], [1, 2, 3], [1, 2, 4], [1, 2, 5], [1, 2, 6], [1, 3, 4], [1, 3, 5], [1, 3, 6], [1, 4, 5], [1, 4, 6], [1, 5, 6], [2, 3, 4], [2, 3, 5], [2, 3, 6], [2, 4, 5], [2, 4, 6], [2, 5, 6], [3, 4, 5], [3, 4, 6], [3, 5, 6], [4, 5, 6]]

a = [[825, 175], [956, 44], [994, 6], [999, 1], [1000, 0], [831, 169], [969, 31], [992, 8], [999, 1], [831, 169], [970, 30], [990, 10], [827, 173], [964, 36], [858, 142], [761, 179, 60], [797, 173, 30], [816, 163, 21], [802, 186, 12], [910, 60, 30], [920, 62, 18], [916, 67, 17], [945, 29, 26], [962, 23, 15], [962, 18, 20], [745, 182, 73], [792, 178, 30], [808, 181, 11], [917, 56, 27], [936, 47, 17], [963, 17, 20], [770, 186, 44], [807, 164, 29], [933, 48, 19], [794, 162, 44]]

t = [1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000]

n = 6

"""
	mod1(ij, a, t, n)

Simplest Multinomial model. Can generate likelihood.
"""
@model function mod1(ij, a, t, n)
	n_trialtypes = length(ij)
	logu ~ filldist(Normal(0,2.0),n-1)		
	getu(i) = i == 6 ? zero(logu[1]) : logu[i]
	for i in 1:n_trialtypes
		us = exp.(getu.(ij[i]))
		ps = softmax(us)
		a[i] ~  Multinomial(t[i],ps)
	end
end

"""
	lapse2(ij, a, t, n)

Adding a softmax noise. This samples, but can't compute likelihood even though the conditional is set to false. 
"""
@model function lapse2(ij, a, t, n)
	n_trialtypes = length(ij)
	logu ~ filldist(Normal(0,2.0),n-1)
	lapse ~ Beta(1,5)
	getu(i) = i == 6 ? zero(logu[1]) : logu[i]
	
	for i in 1:n_trialtypes
		nopts = length(ij[i])
		us = exp.(getu.(ij[i]))
		if false 
			ps = softmax(us * lapse)
		else
			ps = softmax(us)
		end
		a[i] ~ Multinomial(t[i],ps)
	end
end

m1 = mod1(ij,a,t,n)
m2 = lapse2(ij,a,t,n)
c1 = sample(m1, NUTS(), MCMCThreads(),1000, 4);
c2 = sample(m2, NUTS(), MCMCThreads(),1000, 4);

ll1 = loglikelihood(m1, c1) # This works
ll2 = loglikelihood(m2, c2) # This fails if NThreads() > 1

I have tried to simplify the example, but in my simpler models i don't get the error.

When i run with only 1 thread i don't get the error.

I have put debug statements in the the model and the error happens after evaluating the model once (e.g. the first call to Multinomial on the second model evaluation generates the error).

versioninfo():

Julia Version 1.10.5
Commit 6f3fdf7b362 (2024-08-27 14:19 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 20 × Intel(R) Xeon(R) Silver 4210R CPU @ 2.40GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, cascadelake)
Threads: 10 default, 0 interactive, 5 GC (on 20 virtual cores)
Environment:
  JULIA_REVISE_WORKER_ONLY = 1

package info:

Status `/tmp/jl_zxoO1r/Project.toml`
⌃ [13f3f980] CairoMakie v0.12.5
  [861a8166] Combinatorics v1.0.2
  [1a297f60] FillArrays v1.13.0
  [872c559c] NNlib v0.9.24
  [0ff47ea0] PlutoHooks v0.0.5
  [7f904dfe] PlutoUI v0.7.60
  [7f36be82] PosteriorStats v0.2.5
  [2913bbd2] StatsBase v0.34.3
  [fce5fe82] Turing v0.34.1
  [44cfe95a] Pkg v1.10.0
Info Packages marked with ⌃ have new versions available and may be upgradable.

[jerlich]

btw, i have confirmed on my M1 Mac that the likelihood computation for some Turing models is not threadsafe. I haven't figured out what about the model makes it not threadsafe.

[penelopeysm]

A slimmer MWE:

using Turing

@model function f(x)
    ns ~ filldist(Normal(0, 2.0), 3)
    m ~ Uniform(0, 1)
    x ~ Normal(m, 1)
end

model = f(1)
chain = sample(model, NUTS(), MCMCThreads(), 10, 4);

# all of these error with > 1 thread
loglikelihood(model, chain)
logprior(model, chain)
logjoint(model, chain)

It has something to do with the fact that ns and m have different types (in your case, it's logu and lapse), so the type of a dictionary containing these values gets inferred as an Any.

[penelopeysm]

This should be fixed in the v0.28.5 release.

Thanks for reporting @jerlich!

[penelopeysm]

Reopening because this was fixed in 0.28 but the fix hasn't been successfully ported to newer versions. See #686