Distribution shape controlled by discrete random variable

I’m trying to create a model in which the shape of a distribution relies on the posteriori of a discrete random variable. I have done this before in a statistics class using the law of iterated expectation, but I’m wondering if this is going to work in PyMC.
Here is a somewhat simplified example of what I’m trying to achieve:

obs1 = 3
obs2 = 10

with pm.Model() as mod:
    K = pm.Geometric("K", 1/10) # Variable specifing the length of M
    like = pm.Binomial("like", p=0.5, n=K, observed=obs1)

    M = pm.Normal("M", mu=0, sigma=1, size=(K,))
    N = pm.Normal("N", mu=at.sum(M), sigma=1, observed=obs2)

    idata = pm.sample()

But this yields a Value Error:

ValueError: Bad shape in start point:
Expected shape (8,) for var 'M', got: (10,)

I have the impression that PyMC initializes K to the expected value from the Geometric distribution, then sets this as the length of M. Then, when the sampling starts this yields the given error.

I am wondering if this is possible, and if not, if anyone can think of a workaround.
Any help is appreciated!

1 Like

We don’t have any samplers for dynamic size variables in PyMC. That’s an area that still requires investigation.

Since you only need the sum of M, you can perhaps define the variable that corresponds to that, as the sum of normals follows a closed form normal. The mean would still be zero, and sigma something like at.sqrt(K)?

1 Like

Ah, OK.
Unfortunately, in my actual application the expression is a little more complex, so I won’t be able to make a more analytic approach like this. I suppose I will have to handle K outside of the model.
Thanks for the help!