Sampling from Turing model combining discrete and continuous variables fails on second loop through the model

Hi, I’m trying to write a model that simultaneously estimates the mean of two underlying components, while also inferring component membership, with by-subject random intercepts. The data I have are synthetic, 120 subjects that each give numerical ratings for each of 40 items (verbs, but the fact that it’s about language is immaterial). In the data-generating process I wrote, each verb belongs to one of two components (with mean either -10 or 10), and each subject has some random noise by which their numerical ratings are offset (normal(0,1)). The python script I used to generate the fake data is here:


import random
import numpy as np

np.random.seed(3)


out = open("fake_dat_120subjects.csv","w",encoding="utf8")
out.write("SubjectID,VerbID,VerbMean,Observation\n")

num_subjects = 120
num_verbs = 20


subject_adjustment_dict = {}
verb_assigment_dict = {}

for i in range(num_verbs):
    if np.random.binomial(1, .5) == 0:
        category_offset = -10
    else:
        category_offset = 10

    verb_assigment_dict[i] = category_offset


for j in range(num_subjects):
    subject_offset = np.random.normal(0,1)
    subject_adjustment_dict[j] = subject_offset



for verb in range(num_verbs):
    for subject in subject_adjustment_dict:
        subject_contribution = np.random.normal(subject_adjustment_dict[subject],1)
        verb_contribution = verb_assigment_dict[verb]
        obs = subject_contribution+verb_contribution
        out.write(
            str(subject) + ',' +
            str(verb) + ',' +
            str(verb_contribution) + ',' +
            str(obs) + '\n'
                  )

I transform it into square data in R, like this:

fake_dat <- read_csv("fake_dat_120subjects.csv") %>% 
  mutate(Observation = as.numeric(Observation))


f <- fake_dat %>% 
  select(SubjectID,VerbID,Observation) %>% 
  pivot_wider(names_from = VerbID, values_from = Observation) %>% 
  select(-SubjectID) %>% 
  as.matrix()
  
write.csv(f,"square_dat.csv", row.names = FALSE)

The Julia code I wrote trying to mimic the generative process is here:

using Turing
using DataFrames
using CSV
using ArviZ
using MCMCChains
using StatsPlots

indat = CSV.read("square_dat.csv", DataFrame)
obs_in = Matrix(indat)



@model function mixture_classifier(obs,num_verbs,num_subjects)
    π1 ~ Beta(2,2)
    π2 = 1-π1
    
    mu_class_b ~ Normal(0,5)
    mu_class_a ~ Normal(0,5)
    println("mu class a type is ", mu_class_a)

    println("mu class b type is ", mu_class_b)
    subject_specific_offsets = [Normal(0,1) for _ in 1:num_subjects]
    for v in 1:num_verbs
        println("verb is number",v)
        cluster_assigment ~ Categorical([π1, π2])
        println("cluster assignment is of type ",typeof(cluster_assigment))
        if cluster_assigment == 1
            mu_for_this_verb  = mu_class_a # TODO this should be a rand statement I think
            println("mu class a is  of type ",typeof(mu_class_a))
            println("mu for this verb is of type ",typeof(mu_for_this_verb))

        else
            mu_for_this_verb = mu_class_b
            println("mu class b is of type ",typeof(mu_class_b))
            println("mu for this verb is of type ",typeof(mu_for_this_verb))

        end
        for s in 1:num_subjects
            t = rand(subject_specific_offsets[s])
            obs[s,v] = mu_for_this_verb+t
        end
    end
end

c4 = sample(mixture_classifier(obs_in,20,120), 
    Gibbs(PG(20, :cluster_assigment), NUTS(500, 0.65,:mu_class_a, :mu_class_b)), 
    10;  # draws - for some reason it doesn't like the keyword here, but that's what this is.
    progress = true,
    thinning = 1,
    discard_initial = 0)


When this runs, for the first few loops through the verbs, everything is fine; the mu_class_a/b variables are floats, and so is the t variable, and things work fine. However, after a few iterations (I don’t really know how this works internally, sorry!) it turns out that the mu_class_a/b things are distribution objects, and things fair; I get the error below:

verb is number18
cluster assignment is of type Int64
mu class a is  of type Float64
mu for this verb is of type Float64
verb is number19
cluster assignment is of type Int64
mu class a is  of type Float64
mu for this verb is of type Float64
verb is number20
cluster assignment is of type Int64
mu class a is  of type Float64
mu for this verb is of type Float64
mu class a type is Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:π1, :mu_class_b, :mu_class_a, :cluster_assigment), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:π1, Setfield.IdentityLens}, Int64}, Vector{Beta{Float64}}, Vector{AbstractPPL.VarName{:π1, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:mu_class_b, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:mu_class_b, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:mu_class_a, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:mu_class_a, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:cluster_assigment, Setfield.IdentityLens}, Int64}, Vector{Categorical{Float64, Vector{Float64}}}, Vector{AbstractPPL.VarName{:cluster_assigment, Setfield.IdentityLens}}, Vector{Int64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(mixture_classifier), (:obs, :num_verbs, :num_subjects), (), (), Tuple{Matrix{Float64}, Int64, Int64}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (:mu_class_a, :mu_class_b), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}}(-0.7455153559237497,0.0,1.0)
mu class b type is Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:π1, :mu_class_b, :mu_class_a, :cluster_assigment), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:π1, Setfield.IdentityLens}, Int64}, Vector{Beta{Float64}}, Vector{AbstractPPL.VarName{:π1, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:mu_class_b, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:mu_class_b, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:mu_class_a, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:mu_class_a, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:cluster_assigment, Setfield.IdentityLens}, Int64}, Vector{Categorical{Float64, Vector{Float64}}}, Vector{AbstractPPL.VarName{:cluster_assigment, Setfield.IdentityLens}}, Vector{Int64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(mixture_classifier), (:obs, :num_verbs, :num_subjects), (), (), Tuple{Matrix{Float64}, Int64, Int64}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (:mu_class_a, :mu_class_b), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}}(-1.0910178332169398,1.0,0.0)
verb is number1
cluster assignment is of type Int64
mu class a is  of type ForwardDiff.Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:π1, :mu_class_b, :mu_class_a, :cluster_assigment), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:π1, Setfield.IdentityLens}, Int64}, Vector{Beta{Float64}}, Vector{AbstractPPL.VarName{:π1, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:mu_class_b, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:mu_class_b, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:mu_class_a, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:mu_class_a, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:cluster_assigment, Setfield.IdentityLens}, Int64}, Vector{Categorical{Float64, Vector{Float64}}}, Vector{AbstractPPL.VarName{:cluster_assigment, Setfield.IdentityLens}}, Vector{Int64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(mixture_classifier), (:obs, :num_verbs, :num_subjects), (), (), Tuple{Matrix{Float64}, Int64, Int64}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (:mu_class_a, :mu_class_b), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}, Float64, 2}
mu for this verb is of type ForwardDiff.Dual{ForwardDiff.Tag{Turing.Essential.var"#f#4"{DynamicPPL.TypedVarInfo{NamedTuple{(:π1, :mu_class_b, :mu_class_a, :cluster_assigment), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:π1, Setfield.IdentityLens}, Int64}, Vector{Beta{Float64}}, Vector{AbstractPPL.VarName{:π1, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:mu_class_b, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:mu_class_b, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:mu_class_a, Setfield.IdentityLens}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:mu_class_a, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:cluster_assigment, Setfield.IdentityLens}, Int64}, Vector{Categorical{Float64, Vector{Float64}}}, Vector{AbstractPPL.VarName{:cluster_assigment, Setfield.IdentityLens}}, Vector{Int64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(mixture_classifier), (:obs, :num_verbs, :num_subjects), (), (), Tuple{Matrix{Float64}, Int64, Int64}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (:mu_class_a, :mu_class_b), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext}, Float64}, Float64, 2}
ERROR: TypeError: in typeassert, expected Float64, got a value of type ForwardDiff.Dual{Nothing, Float64, 2}

Any advice on how to fix this would be very much appreciated; I don’t feel like the model should be particularly difficult in principle, but I can’t figure this out. (also I apologize for the messiness of doing all of the data prep etc. in 3 different languages). Also also, the reason I’m having each verb be observed independently is b/c I’d like to eventually have each verb have some offset from the group (just an additional noise parameter that is motivated by the theory I’m working with) but decided to drop it here b/c I was having trouble in other areas.