support for multi-dimensional Dirichlet is not working

[hstrey]

I just tested whether the Dirichlet distribution takes a two-dimensional variable
This was supposed to have been fixed by #844

 import pymc3 as pm
 import numpy as np
 with pm.Model():
     P = pm.Dirichlet('P', np.ones((3, 3)))

results in:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-2-d05340b35be4> in <module>()
      2 import numpy as np
      3 with pm.Model():
----> 4     P = pm.Dirichlet('P', np.ones((3, 3)))

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/pymc3/distributions/distribution.py in __new__(cls, name, *args, **kwargs)
     29             data = kwargs.pop('observed', None)
     30             dist = cls.dist(*args, **kwargs)
---> 31             return model.Var(name, dist, data)
     32         else:
     33             raise TypeError("Name needs to be a string but got: %s" % name)

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/pymc3/model.py in Var(self, name, dist, data)
    282             else:
    283                 var = TransformedRV(name=name, distribution=dist, model=self,
--> 284                                     transform=dist.transform)
    285                 pm._log.debug('Applied {transform}-transform to {name}'
    286                               ' and added transformed {orig_name} to model.'.format(

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/pymc3/model.py in __init__(self, type, owner, index, name, distribution, model, transform)
    684             transformed_name = "{}_{}_".format(name, transform.name)
    685             self.transformed = model.Var(
--> 686                 transformed_name, transform.apply(distribution))
    687 
    688             normalRV = transform.backward(self.transformed)

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/pymc3/distributions/transforms.py in apply(self, dist)
     30 
     31     def apply(self, dist):
---> 32         return TransformedDistribution.dist(dist, self)
     33 
     34     def __str__(self):

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/pymc3/distributions/distribution.py in dist(cls, *args, **kwargs)
     39     def dist(cls, *args, **kwargs):
     40         dist = object.__new__(cls)
---> 41         dist.__init__(*args, **kwargs)
     42         return dist
     43 

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/pymc3/distributions/transforms.py in __init__(self, dist, transform, *args, **kwargs)
     59         self.dist = dist
     60         self.transform_used = transform
---> 61         v = forward(FreeRV(name='v', distribution=dist))
     62         self.type = v.type
     63 

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/pymc3/model.py in __init__(self, type, owner, index, name, distribution, model)
    506             self.tag.test_value = np.ones(
    507                 distribution.shape, distribution.dtype) * distribution.default()
--> 508             self.logp_elemwiset = distribution.logp(self)
    509             self.model = model
    510 

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/pymc3/distributions/multivariate.py in logp(self, value)
    243 
    244         # only defined for sum(value) == 1
--> 245         return bound(tt.sum(logpow(value, a - 1) - gammaln(a), axis=-1)
    246                      + gammaln(tt.sum(a, axis=-1)),
    247                      tt.all(value >= 0), tt.all(value <= 1),

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/pymc3/distributions/dist_math.py in logpow(x, m)
     56     """
     57     # return m * log(x)
---> 58     return tt.switch(tt.eq(x, 0), -np.inf, m * tt.log(x))
     59 
     60 

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/theano/gof/op.py in __call__(self, *inputs, **kwargs)
    619             for i, ins in enumerate(node.inputs):
    620                 try:
--> 621                     storage_map[ins] = [self._get_test_value(ins)]
    622                     compute_map[ins] = [True]
    623                 except AttributeError:

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/theano/gof/op.py in _get_test_value(cls, v)
    547             # ensure that the test value is correct
    548             try:
--> 549                 ret = v.type.filter(v.tag.test_value)
    550             except Exception as e:
    551                 # Better error message.

/Users/hstrey/anaconda/envs/python3/lib/python3.5/site-packages/theano/tensor/type.py in filter(self, data, strict, allow_downcast)
    175             raise TypeError("Wrong number of dimensions: expected %s,"
    176                             " got %s with shape %s." % (self.ndim, data.ndim,
--> 177                                                         data.shape))
    178         if not data.flags.aligned:
    179             try:

TypeError: For compute_test_value, one input test value does not have the requested type.

The error when converting the test value to that variable type:
Wrong number of dimensions: expected 1, got 2 with shape (3, 3).

[twiecki]

Can add you add shape=(3, 3)?

[hstrey]

cool. Now it works. Why is the Dirichlet distribution not calling shape on a? Wouldn't this be a more straightforward way of handling multi-dimensional parameters. The Normal distribution does not require a shape argument.

Thanks for your help

Helmut

[twiecki]

Yeah, the shape thing is a bit odd. There is a PR to implement what you suggest: #1125 it's not straight-forward though.

[twiecki]

Closing this.