NUTS fails to initalize with multi-dim catergorical model

[bdyetton]

I've just switched to the master branch so i can take advantage of this bug fix: #3535.
However, i'm getting:

Initializing NUTS failed. Falling back to elementwise auto-assignment.

on some models.
Its seems to be to do with having a 2-dimensional shape for a categorical (and therefore makes me think it's related to #3535)

Note the two models below, where the top one is fine (besides the divergences), the bottom one fails:

data = np.random.randint(0, 3, size=(1000, 1))

    with pm.Model() as model:
        tp1 = pm.Dirichlet('tp1', a=np.array([0.25]*4), shape=(4,))
        obs = pm.Categorical('obs', p=tp1, observed=data)
        trace = pm.sample()

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [tp1]
Sampling 4 chains, 221 divergences: 100%|█| 4000/4000 [00:01<00:00, 3044.22draws/s]
There were 64 divergences after tuning. Increase `target_accept` or reparameterize.
There were 49 divergences after tuning. Increase `target_accept` or reparameterize.
There were 56 divergences after tuning. Increase `target_accept` or reparameterize.
There were 52 divergences after tuning. Increase `target_accept` or reparameterize.

data = np.random.randint(0,3,size=(1000,2))
with pm.Model() as model:
    tp1 = pm.Dirichlet('tp1', a=np.array([0.25]*4), shape=(4,4)) #difference is the shape here here
    obs = pm.Categorical('obs', p=tp1[data[:,0],:], observed=data[:,1]) #and shape here
    trace = pm.sample()

which produces:

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Initializing NUTS failed. Falling back to elementwise auto-assignment.
Multiprocess sampling (4 chains in 4 jobs)
Slice: [tp1]
Sampling 4 chains, 0 divergences: 100%|█| 4000/4000 [00:08<00:00, 454.37draws/s]

Versions and main components

• PyMC3 Version: current master (Fixed categorical logp with tt.choose #3563)
• Theano Version: 1.0.4
• Python Version: 3.7.2
• Operating system: ubuntu 16.04
• How did you install PyMC3: pip

[lucianopaz]

Thanks for reporting this. It seems that the cause of this is that tt.choose has no grad method. I'll try to find a workaround.

[rpgoldman]

Could the failing test case go in the set of tests somewhere?

[bdyetton]

There is still an issue somewhere with multi-dim categorical.
Because #3535 was closed by #3572 (and the problem still exisits), i've added this issue here.

This issue is best highlighted by this toy example, where the first model runs fast (3000it/s) but the second is very slow (3it/s) when sampling with NUTS. When using the slice sampler, the second model will not even begin sampling.

    data = np.random.randint(0, 3, size=(1000, 1))
    with pm.Model() as model:
        tp1 = pm.Dirichlet('tp1', a=np.array([0.25]*4), shape=(4,))
        obs = pm.Categorical('obs', p=tp1, observed=data)
        trace = pm.sample() #super fast!

    data_indexer = np.random.randint(0,2,size=(1000,))

    with pm.Model() as model:
        tp1 = pm.Dirichlet('tp1', a=np.array([0.25]*4), shape=(2,4))
        obs = pm.Categorical('obs', p=tp1[data_indexer, :], observed=data)
        trace = pm.sample() #takes ages!

[junpenglao]

In your second model the logp is incorrectly broadcasted - it should be:

data_indexer = np.random.randint(0,2,size=(1000,))

with pm.Model() as model1:
    tp1 = pm.Dirichlet('tp1', a=np.array([0.25]*4), shape=(2,4))
    obs = pm.Categorical('obs', p=tp1[data_indexer, :], observed=data.T)
    trace = pm.sample() #takes ages!

I did obs.distribution.logp(data).tag.test_value.shape that find out the problem.

Maybe we should have better shape check in Categorical (in general Categorical is a difficult distribution because we want it to work with the general case that data_rank = parameter_rank - 1 and a repeated measure case which data_rank = parameter_rank)

[bdyetton]

Hi @junpenglao,
Thanks for fixing :)
However, i'm pretty confused at why data requires that transpose.
The Print(attrs=['shape']))(tp1[data_indexer, :]) gives (1000,4), i.e. 1000 different catergorical p distributions, one for each data point. Why would data then need to be (1,1000)?
It would be much more intuitive if the axis for the number of data points and axis for the number of different categorical p distributions is the same.
This is how it works for other distributions, is it not?

    data = np.random.normal(0, 1, size=(1000, 2))
    data_indexer = np.random.randint(0, 2, size=(1000,))
    with pm.Model() as model1:
        mean_params = pm.Normal('m', 0, 1, shape=(2, 2))
        error = pm.HalfNormal('error', 5)
        obs = pm.Normal('obs', mu=mean_params[data_indexer, :], sd=error, observed=data)
        trace = pm.sample()

[bdyetton]

Ok, so I just tried this on the latest master, and it samples fine without the transpose :)
Something must have changed in the categorical dim for that push.
I'll close this :)

[junpenglao]

Glad that this is fixed. @lucianopaz do you know if we have a test for this?

[lucianopaz]

@junpenglao, it's not an explicit test for Categorical.logp. #3572 added tests to verify that take_along_axis has a grad, and that it is correct. Given that Categorical.logp basically only uses take_along_axis, I thought that adding another test was reduntant.

[junpenglao]

Yep sg :-)