Issue with sampling from batches of Dirichlet distributions

[nhartland]

Hi pymc3 devs,

This is an issue initially reported on discourse. I'm trying to sample a batch of Dirichlet distributions of batch size 1 (for broadcasting against other distributions). This is giving me the below error. The same example is working fine in pymc 3.8.

Minimal example

import numpy as np
import pymc3 as pm

with pm.Model() as model:
    α = pm.Dirichlet(r'α', a=np.ones((1, 2)))
    pm.sample()

Traceback

Traceback (most recent call last):
  File "./bugtest.py", line 6, in <module>
    pm.sample()
  File "/usr/local/Caskroom/miniconda/base/envs/bugtest/lib/python3.8/site-packages/pymc3/sampling.py", line 427, in sample
    check_start_vals(model.test_point, model)
  File "/usr/local/Caskroom/miniconda/base/envs/bugtest/lib/python3.8/site-packages/pymc3/util.py", line 235, in check_start_vals
    raise SamplingError(
pymc3.exceptions.SamplingError: Initial evaluation of model at starting point failed!
Starting values:
{'α_stickbreaking__': array([[0.]])}

Initial evaluation results:
α_stickbreaking__   -inf
Name: Log-probability of test_point, dtype: float64

Versions and main components

• PyMC3 Version: 3.11.0
• Theano Version: theano-pymc-1.1.0
• Python Version: 3.8.6
• Operating system: macos
• How did you install PyMC3: conda

[twiecki]

CC our shape experts @rpgoldman @lucianopaz @Sayam753.

[Sayam753]

Here is the commit f07c273 that introduced this bug.

[nhartland]

Thanks for such a fast response everyone!

[brandonwillard]

There might be a little more to this, because even before those changes—and after the reversion in #4454—it produces nans for the degenerate case of a = 1. In other words, whatever that original logic was/is doing, it might not be sound.