add test_shape_inputs for _OrderedLogistic

Author: danhphan

pymc/tests/test_distributions_random.py

@@ -1688,6 +1688,26 @@ 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)),
+            ([[1.0, -1.0, 0.0], [-1.0, 3.0, 5.0]], [-2.0, 0, 1.0], (2, 3, 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