Revert "Fixed categorical logp with tt.choose (#3563)" (#3568)

This reverts commit 8191b49.

Author: junpenglao

RELEASE-NOTES.md

@@ -11,8 +11,6 @@
 
 ### 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,12 +998,7 @@ def logp(self, value):
 
         if p.ndim > 1:
             pattern = (p.ndim - 1,) + tuple(range(p.ndim - 1))
-            a = tt.log(
-                tt.choose(
-                    value_clip,
-                    p.dimshuffle(pattern),
-                )
-            )
+            a = tt.log(p.dimshuffle(pattern)[value_clip])
         else:
             a = tt.log(p[value_clip])
 
@@ -1575,13 +1570,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(self.cutpoints - tt.shape_padright(self.eta))
+        pa = sigmoid(tt.shape_padleft(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,35 +1335,3 @@ 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)