Dirichlet and Categorical

I want to estimate a latent group (ii) that varies the mean (a) of a simple regression.

I can run the sampling steps with non convergence but I cannot run the prior steps. I might have made some shape errors.

alpha = [1, 1]

with pm.Model() as model:
  a = pm.Normal('a', shape=2)
  b = pm.Normal('b', shape=1)
  sigma = pm.Gamma('sigma', 1, 1, shape=1)
  theta = pm.Dirichlet('theta', alpha)
  ii = pm.Categorical('ii', theta, shape=N)
  pred = []
  for i in range(N):
  	pred.append(pm.Normal('pred' + str(i), a[ii[i]] + b * x[i], sigma, observed=y[i], shape=1))

  prior = pm.sample_prior_predictive() # doesnt work
  posterior = pm.sample(200, tune=100, chains=2) # works
  posterior_pred = pm.sample_posterior_predictive(posterior)

error

File "/usr/lib/python3/dist-packages/numpy/lib/stride_tricks.py", line 191, in _broadcast_shape
    b = np.broadcast(*args[:32])
ValueError: shape mismatch: objects cannot be broadcast to a single shape

Hi Dirk, welcome :wink:
Can confirm the error, which seems to be a PyMC3 shape bug, meanwhile, you can try the workaround below:

import numpy as np
import pymc3 as pm

N = 100
x = np.linspace(0, 1, N)
y = np.random.randn(N)
alpha = np.ones(2)

with pm.Model() as model:
    a = pm.Normal('a', shape=2)
    b = pm.Normal('b')
    sigma = pm.Gamma('sigma', 1, 1)
    theta = pm.Dirichlet('theta', alpha)
    ii = pm.Categorical('ii', theta[None, ...], shape=N)
    pred = pm.Normal('pred', a[ii] + b * x, sigma, observed=y)
    prior_sample = pm.sample_prior_predictive()

cc @lucianopaz

Thank you, that works now.

My theta posterior is not as expected (not even), so I need to improve the model.