I'm pretty sure your trouble are from trying to get samples for the simulation study from the numpyro.sample object, rather than from the dist object.
My sense is that you should parameterize the dist object, and then my_dist.sample() to get your random samples. Then you can wrap around the model function that gives you the parameterized dist with the priors/sampling function that returns the sample object. I call these functions one_division_dist and one_division_sample (or something).
This code doesn't work because I run into a problem when doing all the sampling in one big dist.Multinomial() call. I'm not sure if plating is supposed to make this easy when inference-sampling, but it makes it hard when you're trying to just draw samples from the dist object with dist.sample(), because in fact there are multiple multinomial samples, one per "observation" (lineage + time combination). I not sure I can make the z = beta_0 + beta_1 * time thing work, because the dimensions of the different things is so crazy.
I'm pretty sure your trouble are from trying to get samples for the simulation study from the
numpyro.sample
object, rather than from thedist
object.My sense is that you should parameterize the
dist
object, and thenmy_dist.sample()
to get your random samples. Then you can wrap around the model function that gives you the parameterizeddist
with the priors/sampling function that returns thesample
object. I call these functionsone_division_dist
andone_division_sample
(or something).This code doesn't work because I run into a problem when doing all the sampling in one big
dist.Multinomial()
call. I'm not sure if plating is supposed to make this easy when inference-sampling, but it makes it hard when you're trying to just draw samples from thedist
object withdist.sample()
, because in fact there are multiple multinomial samples, one per "observation" (lineage + time combination). I not sure I can make thez = beta_0 + beta_1 * time
thing work, because the dimensions of the different things is so crazy.