Add dims to *Ordered probs variables when appropriate

Author: AustinRochford

pymc/distributions/discrete.py

@@ -1745,7 +1745,8 @@ class OrderedLogistic:
     def __new__(cls, name, *args, compute_p=True, **kwargs):
         out_rv = _OrderedLogistic(name, *args, **kwargs)
         if compute_p:
-            pm.Deterministic(f"{name}_probs", out_rv.owner.inputs[3])
+            pm.Deterministic(f"{name}_probs", out_rv.owner.inputs[3],
+                             dims=kwargs.get('dims'))
         return out_rv
 
     @classmethod
@@ -1856,7 +1857,8 @@ class OrderedProbit:
     def __new__(cls, name, *args, compute_p=True, **kwargs):
         out_rv = _OrderedProbit(name, *args, **kwargs)
         if compute_p:
-            pm.Deterministic(f"{name}_probs", out_rv.owner.inputs[3])
+            pm.Deterministic(f"{name}_probs", out_rv.owner.inputs[3],
+                             dims=kwargs.get('dims'))
         return out_rv
 
     @classmethod

pymc/distributions/multivariate.py

@@ -752,7 +752,8 @@ class OrderedMultinomial:
     def __new__(cls, name, *args, compute_p=True, **kwargs):
         out_rv = _OrderedMultinomial(name, *args, **kwargs)
         if compute_p:
-            pm.Deterministic(f"{name}_probs", out_rv.owner.inputs[4])
+            pm.Deterministic(f"{name}_probs", out_rv.owner.inputs[4],
+                             dims=kwargs.get('dims'))
         return out_rv
 
     @classmethod