Implementing softmax regression with a latent Gaussian variable

Specifying the shape should do the tricks, but you do need to be careful checking them to make sure they are correct:

with pm.Model() as model:
    alpha = pm.Normal('alpha', mu=0, sd=1, shape=omega.shape[0])
    beta = pm.Normal('beta', mu=0, sd=1, shape=omega.shape[0])
    sigma = pm.HalfNormal('sigma', sd=1, shape=omega.shape[0])
    mu = x[:, None] * beta + alpha
    latent = pm.Normal('latent', mu=mu, sd=sigma, shape=(10, 4))
    probs = pm.math.dot(pm.math.exp(latent), omega) / pm.math.exp(latent).sum(1, keepdims=True)
    y = pm.Categorical('y', p=probs, observed=samples)
    trace = pm.sample()

However, this will not work currently, as samples you defined above is float type array, and categorical accept observed being one-hot encoded int array (in this case, a length 10 array)
Also, normalizing random variables (i.e., the softmax in your model) makes your model unidentifiable - you might find this discussion here helpful: Multinomial hierarchical regression with multiple observations per group ("Bad energy issue")

1 Like