Reject multivarate and nonrvs for logp of max

pymc/logprob/order.py

@@ -78,11 +78,8 @@ def find_measurable_max(fgraph: FunctionGraph, node: Node) -> Optional[List[Tens
     if not rv_map_feature.request_measurable(node.inputs):
         return None
 
-    # Non-univariate non-RVs must be rejected
-    if (
-        not isinstance(base_var.owner.op, RandomVariable)
-        and base_var.owner.inputs[0].owner.op.ndim_supp != 0
-    ):
+    # Non-univariate distributions and non-RVs must be rejected
+    if not (isinstance(base_var.owner.op, RandomVariable) and base_var.owner.op.ndim_supp == 0):
         return None
 
     # TODO: We are currently only supporting continuous rvs

tests/logprob/test_order.py

@@ -78,6 +78,17 @@ def test_max_non_rv_fails():
         x_max_logprob = logp(x_max, x_max_value)
 
 
+def test_max_multivariate_rv_fails():
+    _alpha = pt.scalar()
+    _k = pt.iscalar()
+    x = pm.StickBreakingWeights.dist(_alpha, _k)
+    x.name = "x"
+    x_max = pt.max(x, axis=-1)
+    x_max_value = pt.vector("x_max_value")
+    with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")):
+        x_max_logprob = logp(x_max, x_max_value)
+
+
 def test_max_categorical():
     """Test whether the logprob for ```pt.max``` for unsupported distributions is correctly rejected"""
     x = pm.Categorical.dist([1, 1, 1, 1], shape=(5,))