Fix dirichlet shape (#1241)

* BUG Reverse shape order of Dirichlet, thanks to @taku-y.

* DOC Various fixes to example notebooks. Mainly rename ElemwiseCategoricalStep to ElemwiseCategorical.

* MAINT Remove MyDirichlet.

* Fix variable shape's inconsistency.

* Fix test function for stickbreaking transform.

Author: twiecki

docs/source/notebooks/dp_mix.ipynb

@@ -80,9 +80,9 @@ class MvStudentT(Continuous):
     Multivariate Student T log-likelihood.
 
     .. math::
-    
+
     f(\mathbf{x}| \nu,\mu,\Sigma) = \frac{\Gamma\left[(\nu+p)/2\right]}{\Gamma(\nu/2)\nu^{p/2}\pi^{p/2}\left|{\Sigma}\right|^{1/2}\left[1+\frac{1}{\nu}({\mathbf x}-{\mu})^T{\Sigma}^{-1}({\mathbf x}-{\mu})\right]^{(\nu+p)/2}}
-    
+
     ========  ==========================
     Support   :math:`x \in \mathbb{R}^k`
     Mean      :math:`\mu` if :math:`\nu > 1`  else undefined
@@ -103,37 +103,35 @@ def __init__(self, nu, Sigma, mu=None, *args, **kwargs):
         self.nu = nu
         self.mu = np.zeros(Sigma.shape[0]) if mu is None else mu
         self.Sigma = Sigma
-        
+
         self.mean = self.median = self.mode = self.mu = mu
-        
+
     def random(self, point=None, size=None):
         chi2 = np.random.chisquare
         mvn = np.random.multivariate_normal
-        
+
         nu, S, mu = draw_values([self.nu, self.Sigma, self.mu], point=point)
-        
+
         return (np.sqrt(nu) * (mvn(np.zeros(len(S)), S, size).T
                                            / chi2(nu, size))).T + mu
-        
-    def logp(self, value): 
-        
+
+    def logp(self, value):
+
         S = self.Sigma
         nu = self.nu
         mu = self.mu
 
         d = S.shape[0]
         n = value.shape[0]
-        
+
         X = value - mu
-    
+
         Q = X.dot(matrix_inverse(S)).dot(X.T).sum()
         log_det = tt.log(det(S))
         log_pdf = gammaln((nu + d)/2.) - 0.5 * (d*tt.log(np.pi*nu) + log_det) - gammaln(nu/2.)
         log_pdf -= 0.5*(nu + d)*tt.log(1 + Q/nu)
-    
+
         return log_pdf
-    
-        
 
 
 class Dirichlet(Continuous):
@@ -167,7 +165,7 @@ class Dirichlet(Continuous):
     """
     def __init__(self, a, transform=transforms.stick_breaking,
                  *args, **kwargs):
-        shape = a.shape[0]
+        shape = a.shape[-1]
         kwargs.setdefault("shape", shape)
         super(Dirichlet, self).__init__(transform=transform, *args, **kwargs)
 
@@ -195,8 +193,8 @@ def logp(self, value):
         a = self.a
 
         # only defined for sum(value) == 1
-        return bound(tt.sum(logpow(value, a - 1) - gammaln(a), axis=0)
-                     + gammaln(tt.sum(a, axis=0)),
+        return bound(tt.sum(logpow(value, a - 1) - gammaln(a), axis=-1)
+                     + gammaln(tt.sum(a, axis=-1)),
                      tt.all(value >= 0), tt.all(value <= 1),
                      k > 1, tt.all(a > 0))
 

docs/source/notebooks/gaussian-mixture-model-advi.ipynb

@@ -62,6 +62,11 @@ def __init__(self, dist, transform, *args, **kwargs):
             v.shape.tag.test_value, v.dtype,
             testval, dist.defaults, *args, **kwargs)
 
+        if transform.name == 'stickbreaking':
+            b = np.hstack(((np.atleast_1d(self.shape) == 1)[:-1], False))
+            self.type = tt.TensorType(v.dtype, b) # force the last dim not broadcastable
+
+
     def logp(self, x):
         return (self.dist.logp(self.transform_used.backward(x)) +
                 self.transform_used.jacobian_det(x))
@@ -177,46 +182,46 @@ def jacobian_det(self, x):
 
 sum_to_1 = SumTo1()
 
-
 class StickBreaking(Transform):
     """Transforms K dimensional simplex space (values in [0,1] and sum to 1) to K - 1 vector of real values.
-
     Primarily borrowed from the STAN implementation.
     """
 
     name = "stickbreaking"
 
-    def forward(self, x):
+    def forward(self, x_):
+        x = x_.T
         # reverse cumsum
         x0 = x[:-1]
         s = tt.extra_ops.cumsum(x0[::-1], 0)[::-1] + x[-1]
         z = x0/s
         Km1 = x.shape[0] - 1
         k = tt.arange(Km1)[(slice(None), ) + (None, ) * (x.ndim - 1)]
-        eq_share = - tt.log(Km1 - k)   # logit(1./(Km1 + 1 - k))
+        eq_share = logit(1./(Km1 + 1 - k)) # - tt.log(Km1 - k)
         y = logit(z) - eq_share
-        return y
+        return y.T
 
-    def backward(self, y):
+    def backward(self, y_):
+        y = y_.T
         Km1 = y.shape[0]
         k = tt.arange(Km1)[(slice(None), ) + (None, ) * (y.ndim - 1)]
-        eq_share = - tt.log(Km1 - k)   # logit(1./(Km1 + 1 - k))
+        eq_share = logit(1./(Km1 + 1 - k)) #- tt.log(Km1 - k)
         z = inverse_logit(y + eq_share)
         yl = tt.concatenate([z, tt.ones(y[:1].shape)])
         yu = tt.concatenate([tt.ones(y[:1].shape), 1-z])
         S = tt.extra_ops.cumprod(yu, 0)
         x = S * yl
-        return x
+        return x.T
 
-    def jacobian_det(self, y):
+    def jacobian_det(self, y_):
+        y = y_.T
         Km1 = y.shape[0]
         k = tt.arange(Km1)[(slice(None), ) + (None, ) * (y.ndim - 1)]
-        eq_share = -tt.log(Km1 - k)  # logit(1./(Km1 + 1 - k))
+        eq_share = logit(1./(Km1 + 1 - k)) # -tt.log(Km1 - k)
         yl = y + eq_share
         yu = tt.concatenate([tt.ones(y[:1].shape), 1-inverse_logit(yl)])
         S = tt.extra_ops.cumprod(yu, 0)
-        return tt.sum(tt.log(S[:-1]) - tt.log1p(tt.exp(yl)) - tt.log1p(tt.exp(-yl)),
-                     0)
+        return tt.sum(tt.log(S[:-1]) - tt.log1p(tt.exp(yl)) - tt.log1p(tt.exp(-yl)), 0).T
 
 stick_breaking = StickBreaking()
 
@@ -229,7 +234,7 @@ class Circular(ElemwiseTransform):
     def backward(self, y):
         return tt.arctan2(tt.sin(y), tt.cos(y))
 
-    def forward(self, x): 
+    def forward(self, x):
         return x
 
 circular = Circular()

docs/source/notebooks/gaussian_mixture_model.ipynb

@@ -104,10 +104,9 @@ def __init__(self, n):
 
 class MultiSimplex(object):
     def __init__(self, n_dependent, n_independent):
-        transposed_vals = list(itertools.product(list(simplex_values(n_dependent)), repeat=n_independent))
-        self.vals = list(np.transpose(transposed_vals, (0, 2, 1)))
+        self.vals = [np.vstack(v) for v in list(itertools.product(list(simplex_values(n_dependent)), repeat=n_independent))]
 
-        self.shape = (n_dependent, n_independent)
+        self.shape = (n_independent, n_dependent)
         self.dtype = Unit.dtype
         return
 
@@ -450,13 +449,13 @@ def check_lkj(x, n, p, lp):
                         lp, decimal=6, err_msg=str(pt))
 
 def betafn(a):
-    return scipy.special.gammaln(a).sum(0) - scipy.special.gammaln(a.sum(0))
+    return scipy.special.gammaln(a).sum(-1) - scipy.special.gammaln(a.sum(-1))
 
 def logpow(v, p):
     return np.choose(v==0, [p * np.log(v), 0])
 
 def dirichlet_logpdf(value, a):
-    return (-betafn(a) + logpow(value, a-1).sum(0)).sum()
+    return (-betafn(a) + logpow(value, a-1).sum(-1)).sum()
 
 def test_dirichlet():
     for n in [2,3]:
@@ -471,7 +470,7 @@ def check_dirichlet(n):
 
 def check_dirichlet2D(ndep, nind):
         pymc3_matches_scipy(
-                Dirichlet, MultiSimplex(ndep, nind), {'a': Vector(Vector(Rplus, nind), ndep) },
+                Dirichlet, MultiSimplex(ndep, nind), {'a': Vector(Vector(Rplus, ndep), nind) },
                 dirichlet_logpdf
                 )