Allow OrderedProbit distribution to take vector inputs (#5418)

* fix and add test_vector_inputs for OrderedProbit

* add test_shape_inputs for _OrderedLogistic

Author: danhphan

pymc/distributions/discrete.py

@@ -1964,7 +1964,9 @@ def dist(cls, eta, cutpoints, sigma=1, *args, **kwargs):
         _log_p = at.concatenate(
             [
                 at.shape_padright(normal_lccdf(0, sigma, probits[..., 0])),
-                log_diff_normal_cdf(0, sigma, probits[..., :-1], probits[..., 1:]),
+                log_diff_normal_cdf(
+                    0, at.shape_padright(sigma), probits[..., :-1], probits[..., 1:]
+                ),
                 at.shape_padright(normal_lcdf(0, sigma, probits[..., -1])),
             ],
             axis=-1,

pymc/tests/test_distributions_random.py

@@ -1688,6 +1688,28 @@ class TestOrderedLogistic(BaseTestDistributionRandom):
         "check_rv_size",
     ]
 
+    @pytest.mark.parametrize(
+        "eta, cutpoints, expected",
+        [
+            (0, [-2.0, 0, 2.0], (4,)),
+            ([-1], [-2.0, 0, 2.0], (1, 4)),
+            ([1.0, -2.0], [-1.0, 0, 1.0], (2, 4)),
+            (np.zeros((3, 2)), [-2.0, 0, 1.0], (3, 2, 4)),
+            (np.ones((5, 2)), [[-2.0, 0, 1.0], [-1.0, 0, 1.0]], (5, 2, 4)),
+            (np.ones((3, 5, 2)), [[-2.0, 0, 1.0], [-1.0, 0, 1.0]], (3, 5, 2, 4)),
+        ],
+    )
+    def test_shape_inputs(self, eta, cutpoints, expected):
+        """
+        This test checks when providing different shapes for `eta` parameters.
+        """
+        categorical = _OrderedLogistic.dist(
+            eta=eta,
+            cutpoints=cutpoints,
+        )
+        p = categorical.owner.inputs[3].eval()
+        assert p.shape == expected
+
 
 class TestOrderedProbit(BaseTestDistributionRandom):
     pymc_dist = _OrderedProbit
@@ -1698,6 +1720,30 @@ class TestOrderedProbit(BaseTestDistributionRandom):
         "check_rv_size",
     ]
 
+    @pytest.mark.parametrize(
+        "eta, cutpoints, sigma, expected",
+        [
+            (0, [-2.0, 0, 2.0], 1.0, (4,)),
+            ([-1], [-1.0, 0, 2.0], [2.0], (1, 4)),
+            ([1.0, -2.0], [-1.0, 0, 1.0], 1.0, (2, 4)),
+            ([1.0, -2.0, 3.0], [-1.0, 0, 2.0], np.ones((1, 3)), (1, 3, 4)),
+            (np.zeros((2, 3)), [-2.0, 0, 1.0], [1.0, 2.0, 5.0], (2, 3, 4)),
+            (np.ones((2, 3)), [-1.0, 0, 1.0], np.ones((2, 3)), (2, 3, 4)),
+            (np.zeros((5, 2)), [[-2, 0, 1], [-1, 0, 1]], np.ones((2, 5, 2)), (2, 5, 2, 4)),
+        ],
+    )
+    def test_shape_inputs(self, eta, cutpoints, sigma, expected):
+        """
+        This test checks when providing different shapes for `eta` and `sigma` parameters.
+        """
+        categorical = _OrderedProbit.dist(
+            eta=eta,
+            cutpoints=cutpoints,
+            sigma=sigma,
+        )
+        p = categorical.owner.inputs[3].eval()
+        assert p.shape == expected
+
 
 class TestOrderedMultinomial(BaseTestDistributionRandom):
     pymc_dist = _OrderedMultinomial