Fixed categorical logp with tt.choose

RELEASE-NOTES.md

@@ -11,6 +11,8 @@
 
 ### Maintenance
 - Moved math operations out of `Rice`, `TruncatedNormal`, `Triangular` and `ZeroInflatedNegativeBinomial` `random` methods. Math operations on values returned by `draw_values` might not broadcast well, and all the `size` aware broadcasting is left to `generate_samples`. Fixes [#3481](https://github.com/pymc-devs/pymc3/issues/3481) and [#3508](https://github.com/pymc-devs/pymc3/issues/3508)
+- Fixed a bug in `Categorical.logp`. In the case of multidimensional `p`'s, the indexing was done wrong leading to incorrectly shaped tensors that consumed `O(n**2)` memory instead of `O(n)`. This fixes issue [#3535](https://github.com/pymc-devs/pymc3/issues/3535)
+- Fixed a defect in `OrderedLogistic.__init__` that unnecessarily increased the dimensionality of the underlying `p`. Related to issue issue [#3535](https://github.com/pymc-devs/pymc3/issues/3535) but was not the true cause of it.
 
 ## PyMC3 3.7 (May 29 2019)
 

pymc3/distributions/discrete.py

@@ -998,7 +998,12 @@ def logp(self, value):
 
         if p.ndim > 1:
             pattern = (p.ndim - 1,) + tuple(range(p.ndim - 1))
-            a = tt.log(p.dimshuffle(pattern)[value_clip])
+            a = tt.log(
+                tt.choose(
+                    value_clip,
+                    p.dimshuffle(pattern),
+                )
+            )
         else:
             a = tt.log(p[value_clip])
 
@@ -1570,13 +1575,13 @@ def __init__(self, eta, cutpoints, *args, **kwargs):
         self.eta = tt.as_tensor_variable(floatX(eta))
         self.cutpoints = tt.as_tensor_variable(cutpoints)
 
-        pa = sigmoid(tt.shape_padleft(self.cutpoints) - tt.shape_padright(self.eta))
+        pa = sigmoid(self.cutpoints - tt.shape_padright(self.eta))
         p_cum = tt.concatenate([
-            tt.zeros_like(tt.shape_padright(pa[:, 0])),
+            tt.zeros_like(tt.shape_padright(pa[..., 0])),
             pa,
-            tt.ones_like(tt.shape_padright(pa[:, 0]))
+            tt.ones_like(tt.shape_padright(pa[..., 0]))
         ], axis=-1)
-        p = p_cum[:, 1:] - p_cum[:, :-1]
+        p = p_cum[..., 1:] - p_cum[..., :-1]
 
         super().__init__(p=p, *args, **kwargs)
 

pymc3/tests/test_distributions.py

@@ -1335,3 +1335,35 @@ def test_discrete_trafo():
         with pytest.raises(ValueError) as err:
             Binomial('a', n=5, p=0.5, transform='log')
         err.match('Transformations for discrete distributions')
+
+
+@pytest.mark.parametrize("shape", [tuple(), (1,), (3, 1), (3, 2)], ids=str)
+def test_orderedlogistic_dimensions(shape):
+    # Test for issue #3535
+    loge = np.log10(np.exp(1))
+    size = 7
+    p = np.ones(shape + (10,)) / 10
+    cutpoints = np.tile(logit(np.linspace(0, 1, 11)[1:-1]), shape + (1,))
+    obs = np.random.randint(0, 1, size=(size,) + shape)
+    with Model():
+        ol = OrderedLogistic(
+            "ol",
+            eta=np.zeros(shape),
+            cutpoints=cutpoints,
+            shape=shape,
+            observed=obs
+        )
+        c = Categorical(
+            "c",
+            p=p,
+            shape=shape,
+            observed=obs
+        )
+    ologp = ol.logp({"ol": 1}) * loge
+    clogp = c.logp({"c": 1}) * loge
+    expected = -np.prod((size,) + shape)
+
+    assert c.distribution.p.ndim == (len(shape) + 1)
+    assert np.allclose(clogp, expected)
+    assert ol.distribution.p.ndim == (len(shape) + 1)
+    assert np.allclose(ologp, expected)