Add a hack to statically infer Dirichlet argument shapes

Author: brandonwillard

pymc3/distributions/multivariate.py

@@ -23,6 +23,7 @@
 
 from scipy import stats, linalg
 
+from theano.gof.op import get_test_value
 from theano.tensor.nlinalg import det, matrix_inverse, trace, eigh
 from theano.tensor.slinalg import Cholesky
 import pymc3 as pm
@@ -486,6 +487,21 @@ class Dirichlet(Continuous):
 
     def __init__(self, a, transform=transforms.stick_breaking,
                  *args, **kwargs):
+
+        if kwargs.get('shape') is None:
+            warnings.warn(
+                (
+                    "Shape not explicitly set. "
+                    "Please, set the value using the `shape` keyword argument. "
+                    "Using the test value to infer the shape."
+                ),
+                DeprecationWarning
+            )
+            try:
+                kwargs['shape'] = get_test_value(tt.shape(a))
+            except AttributeError:
+                pass
+
         super().__init__(transform=transform, *args, **kwargs)
 
         self.size_prefix = tuple(self.shape[:-1])

pymc3/tests/test_distributions.py

@@ -1328,6 +1328,15 @@ def test_dirichlet(self, n):
             Dirichlet, Simplex(n), {"a": Vector(Rplus, n)}, dirichlet_logpdf
         )
 
+    def test_dirichlet_shape(self):
+        a = tt.as_tensor_variable(np.r_[1, 2])
+        with pytest.warns(DeprecationWarning):
+            dir_rv = Dirichlet.dist(a)
+            assert dir_rv.shape == (2,)
+
+        with pytest.warns(DeprecationWarning), theano.change_flags(compute_test_value="ignore"):
+            dir_rv = Dirichlet.dist(tt.vector())
+
     def test_dirichlet_2D(self):
         self.pymc3_matches_scipy(
             Dirichlet,