Add dims to *Ordered probs variables when appropriate (#5084)

* Add dims to *Ordered probs variables when appropriate * Make black happy * For real about black this time
pymc-devs · Oct 17, 2021 · 598dd9d · 598dd9d
1 parent e03f5bf
commit 598dd9d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py
@@ -1745,7 +1745,7 @@ 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 +1856,7 @@ 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

diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py
@@ -752,7 +752,7 @@ 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