Allow OrderedProbit distribution to take vector inputs

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

@@ -21,6 +21,7 @@
 import numpy as np
 import numpy.random as nr
 import numpy.testing as npt
+import pandas as pd
 import pytest
 import scipy.stats as st
 
@@ -368,7 +369,7 @@ def check_pymc_params_match_rv_op(self):
             assert_almost_equal(expected_value, actual_variable.eval(), decimal=self.decimal)
 
     def check_rv_size(self):
-        # test sizes
+        # test sizes    
         sizes_to_check = self.sizes_to_check or [None, (), 1, (1,), 5, (4, 5), (2, 4, 2)]
         sizes_expected = self.sizes_expected or [(), (), (1,), (1,), (5,), (4, 5), (2, 4, 2)]
         for size, expected in zip(sizes_to_check, sizes_expected):
@@ -1698,6 +1699,35 @@ class TestOrderedProbit(BaseTestDistributionRandom):
         "check_rv_size",
     ]
 
+    def test_vector_inputs(self):
+        """ 
+        This test checks when providing vector inputs for `eta` and `sigma` parameters using advanced indexing.
+        """
+        df = pd.DataFrame({    
+            'X' : ['A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B', 'B', 'B'],
+            'Y' : [1, 1, 1, 2, 2, 3, 4, 5, 1, 1, 1, 1, 2, 2, 3, 3]
+        })
+
+        df.Y = df.Y.astype(int)
+        grp_idx = pd.Categorical(df.X).codes
+        K = df.Y.nunique()
+
+        with pm.Model() as opb:    
+            cutpoints = pm.Normal("cutpoints", 0.0, 1.5, shape=K-1,
+                        transform=pm.distributions.transforms.ordered,
+                        initval=np.arange(K-1))
+
+            mu = pm.Normal("mu", mu=K/2, sd=K, shape=2)
+            sigma = pm.HalfNormal("sigma", 1, shape=2)    
+
+            y_obs = pm.OrderedProbit("y_obs", 
+                            cutpoints=cutpoints, 
+                            eta=mu[grp_idx], 
+                            sigma=sigma[grp_idx],                      
+                            observed=df.Y-1)
+
+        assert df.Y.shape == y_obs.eval().shape
+
 
 class TestOrderedMultinomial(BaseTestDistributionRandom):
     pymc_dist = _OrderedMultinomial
@@ -1825,6 +1855,7 @@ def check_errors(self):
                     shape=15,
                 )
 
+
     def check_random_variable_prior(self):
         """
         This test checks for shape correctness when using MatrixNormal distribution