Inconsistency between Dirichlet samples and Categorical argument

[shobhit6993]

There seems to be an inconsistency in the way samples are generated from a Dirichlet, and the way a Categorical interprets its p argument. When drawing multiple draws from a k-dimensional Dirichlet, the elements of one draw are along axis=1 (elements sum to 1 along axis=1), while Categorical seems to expect them to be along axis=0. I need to take the transpose of the Dirichlet stochastic to get the Categorical to work, without which it throws an error. I am guessing that the same would be true for Multinomial as well, although I haven't checked. This behavior is non-intuitive, especially considering how often Dirichlet is used as a prior for Multinomial/Categorical.

Here is a basic model to reproduce the error (uncomment/comment the appropriate Categorical line)

# artificial data
k = 2
ndata = 10
data = np.array([1, 2.2, 2.5, 2.8, 4, 6, 7.2, 7.5, 7.8, 9])

# model
alpha = 0.1 * np.ones((k, ndata))
with pm.Model() as model:
    p = pm.Dirichlet('p', alpha, shape=(k, ndata))
    mu = pm.Normal('mu', mu=5, sd=3, shape=k)
    sd = pm.Uniform('sd', lower=0.1, upper=0.5, shape=k)
    categ = pm.Categorical('categ', p=p.T, shape=(1, ndata))   # works
    # categ = pm.Categorical('categ', p=p, shape=(1, ndata))   # throws an error
    obs = pm.Normal('obs',
                    mu=mu[categ],
                    sd=sd[categ],
                    observed=data)

    step1 = pm.Metropolis(vars=[p, sigma, mu, obs])
    step2 = pm.ElemwiseCategoricalStep(vars=[categ])

    tr = pm.sample(100, step=[step1, step2])
    print tr['p'][0, :]

For reference, the trace for p looks like this:
[[ 0.75233998 0.42171357 0.7272167 0.6370857 0.62833266 0.51940564 0.74280775 0.7106556 0.58151824 0.30892646] [ 0.24766002 0.57828643 0.2727833 0.3629143 0.37166734 0.48059436 0.25719225 0.2893444 0.41848176 0.69107354]]

[twiecki]

@shobhit6993 The first dimension needs to be the dimension of the data:

import numpy as np
import pymc3 as pm

# artificial data
k = 2
data = np.array([1, 2.2, 2.5, 2.8, 4, 6, 7.2, 7.5, 7.8, 9])
ndata = len(data)

# model
alpha = 0.1 * np.ones((ndata, k))
with pm.Model() as model:
    p = pm.Dirichlet('p', alpha, shape=(ndata, k))
    mu = pm.Normal('mu', mu=5, sd=3, shape=k)
    sd = pm.Uniform('sd', lower=0.1, upper=0.5, shape=k)
    categ = pm.Categorical('categ', p=p, shape=ndata)
    obs = pm.Normal('obs',
                    mu=mu[categ],
                    sd=sd[categ],
                    observed=data)

    step1 = pm.Metropolis(vars=[p, sd, mu, obs])
    step2 = pm.ElemwiseCategorical(vars=[categ])

    tr = pm.sample(100, step=[step1, step2])
    print(tr['p'][-1, :])