Merge pull request #1478 from AustinRochford/bugfix-multinomial-vector

Bugfix multinomial vector

Author: fonnesbeck

pymc3/distributions/multivariate.py

@@ -301,15 +301,15 @@ def logp(self, x):
 
         if x.ndim==2:
             x_sum = x.sum(axis=0)
-            n_sum = n * x.shape[0]
+            k = x.shape[0]
         else:
             x_sum = x
-            n_sum = n
+            k = 1
         return bound(
-            factln(n_sum) + tt.sum(x_sum * tt.log(p) - factln(x_sum)),
+            k * factln(n) - tt.sum(factln(x)) + tt.sum(x_sum * tt.log(p)),
             tt.all(x >= 0),
             tt.all(x <= n),
-            tt.eq(tt.sum(x_sum), n_sum),
+            tt.eq(tt.sum(x_sum), k * n),
             tt.all(p <= 1),
             tt.eq(p.sum(), 1),
             n >= 0)

pymc3/tests/test_distributions.py

@@ -568,13 +568,16 @@ def test_multinomial_vec(self):
         vals = np.array([[2,4,4], [3,3,4]])
         p = np.array([0.2, 0.3, 0.5])
         n = 10
-        with Model() as model:
-            Multinomial('m', n=10, p=p, shape=vals.shape)
-        pt = {'m': vals}
-        with Model() as model_sum:
-            Multinomial('m_sum', n=2*n, p=p, shape=len(p))
-        pt_sum = {'m_sum': vals.sum(0)}
-        assert_almost_equal(model.fastlogp(pt), model_sum.fastlogp(pt_sum), decimal=4)
+
+        with Model() as model_single:
+            Multinomial('m', n=n, p=p, shape=len(p))
+
+        with Model() as model_many:
+            Multinomial('m', n=n, p=p, shape=vals.shape)
+
+        assert_almost_equal(sum([model_single.fastlogp({'m': val}) for val in vals]),
+                            model_many.fastlogp({'m': vals}),
+                            decimal=4)
 
     def test_categorical(self):
         for n in [2, 3, 4]: