Dirichlet distribution low probabilities invalid despite constraints

With low values in the concentration parameter the dirichlet prior often fails:

def model(data=None, args=None):
    # Globals.
    background_words = pyro.sample(
        "background_words", dist.Dirichlet(torch.ones(args.num_words) / args.num_words),
    )

    # Locals.
    with pyro.plate("documents", args.num_docs) as ind:
        data = pyro.sample(
                'obs',
                dist.Multinomial(args.num_words_per_doc, background_words),
                obs=data
            )

    return background_words, data

def guide(data, args, batch_size=None):
    background_words_posterior = pyro.param(
        "background_words_posterior",
        lambda: torch.ones(1, args.num_words),
        constraint=constraints.positive,
    )
    pyro.sample("background_words", dist.Dirichlet(background_words_posterior))

Using this minimal model/guide and these args:

class Namespace:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        
    def __hash__(self):
        return hash(frozenset(self.__dict__.items()))

args = {}

args["num_words"] = 1024
args["num_docs"] = 1000
args["num_words_per_doc"] = 20
args["num_steps"] = 1000
args["learning_rate"] = 1e-1
args["batch_size"] = 32

args = Namespace(**args)

I get this failure:

ValueError: The parameter concentration has invalid values
             Trace Shapes:       
              Param Sites:       
background_words_posterior 1 1024
             Sample Sites:       

This is a minimal example but I also get these errors when using Dirichlet distributions in pyro to parameterize simple topic models.

Hi @njwfish, I believe you need to constrain by constrains.simplex rather than constraints.positive.