TArray and multinomial

[itsdfish]

Hi all-

I would like to ask a few related questions about the use of TArray in a model that uses particle Giggs sampling. Here is a simplified version of the model that I am developing:

using Turing 

@model model_function(y) = begin
   N ~ Poisson(5)
   θ ~ Normal(0,1)
   M = length(y)
   s = [TArray(Int64,10) for _ in 1:M]
   #s = [Vector{Int64}(undef,10) for _ in 1:M] #this "works" or at least runs
   for i in 1:M
      s[i] ~ Multinomial(N,10)
   end
   println(s)
   for i in 1:M
      y[i] ~ somedistribution(s[i],θ)
   end
end
y = [1,2,3]
chain = sample(model_function(y), PG(10,10))

The latent parameter s is a vector of vectors, such that each vector s[i] follows a multinomial distribution. somedistribution basically transforms s[i] into a log probability using a softmax rule (not defined here for simplicity). I receive the following error:

ERROR: MethodError: Cannot convert an object of type Array{Int64,1} to an object of type TArray{Int64,1}

First, I want to verify that the inner vector of s should be a TArray. If it is not necessary, I could use [Vector{Int64}(undef,10) for _ in 1:M]. However, assuming that I need to use a TArray, how can I allow each s[i] to follow a multinomial distribution? My first attempt was to create a new distribution and collect the elements from the multinomial sample in a TArray:

import Distributions: logpdf,rand

struct TMultinomial{T1,T2} <: ContinuousUnivariateDistribution
    n::T1
    k::T2
end

Broadcast.broadcastable(x::TMultinomial) = Ref(x)

function rand(d::TMultinomial)
   m = Multinomial(d.n,d.k)
   v = rand(m)
   tv = TArray(Int64,length(v))
   for (i,val) ∈ enumerate(v)
      tv[i] = val
   end 
   return tv
end 

Replacing Multinomial with TMultinomial results in the following error:

MethodError: no method matching iterate(::TMultinomial{Int64,Int64})

This happens even though I overloaded broadcastable. My second attempted solution is the following:

@model model_function(y) = begin
   N ~ Poisson(5)
   M = length(y)
   s = [TArray(Int64,10) for _ in 1:M]
   for i in 1:M
      v ~ Multinomial(N,10)
      for j in 1:10
         s[i][j] = v[j]
      end
   end
   println(s)
   for i in 1:M
      y[i] ~ somedistribution(s[i],θ)
   end
end
y = [1,2,3]
chain = sample(model_function(y), PG(10,10))

This seems to work in the sense that it does not crash and reaches the print statement, but it seems inelegant and I am concerned there might be some unintended consequences.

Any help would be greatly appreciated.

[trappmartin]

Your instantiation of a TArray is incorrect. It should be TArray{Int}(10). The TArray is necessary for particle Gibbs and SMC.

It migth also be a good idea to change the code and use a Matrix instead of a Vector of Vectors as you know the size anyways, i.e.

@model model_function(y) = begin
   N ~ Poisson(5)
   θ ~ Normal(0,1)
   M = length(y)
   s = TArray{Int64}(10,M)
   for i in 1:M
      s[:,i] ~ Multinomial(N,10)
      y[i] ~ somedistribution(s[:,i],θ)
   end
end

Alternatively, you can also write the following:

@model model_function(y, ::Type{Ts}=Matrix{Int}) = begin
   N ~ Poisson(5)
   θ ~ Normal(0,1)
   M = length(y)
   s = Ts(10,M)
   for i in 1:M
      s[:,i] ~ Multinomial(N,10)
      y[i] ~ somedistribution(s[:,i],θ)
   end
end

which automatically transforms Ts to a TArray if a particle based inference algorithm is used and uses the more efficient standard Julia Array for other methods.

[itsdfish]

Thank you so much!

I think I found the source of my confusion. It appears that TArray(Int64,2) is an alias for TArray{Int64}(2)

TArray(Int64,2)
┌ Warning: display(::TArray) prints the originating task's storage, not the current task's storage. Please use show(::TArray) to display the current task's v
ersion of a TArray.
└ @ Libtask ~/.julia/packages/Libtask/RjRkK/src/tarray.jl:105
2-element Array{Int64,1}:
 140319208111840
 140319208094336

but it does not extend to multidimensional arrays, which prompted me to go down various dead ends.