Logprob derivation for Minima

pymc/logprob/order.py

@@ -36,7 +36,6 @@
 
 from typing import List, Optional
 
-import pytensor
 import pytensor.tensor as pt
 
 from pytensor.graph.basic import Node
@@ -46,6 +45,7 @@
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.math import Max
 from pytensor.tensor.random.op import RandomVariable
+from pytensor.tensor.var import TensorVariable
 
 from pymc.logprob.abstract import (
     MeasurableVariable,
@@ -57,14 +57,14 @@
 
 
 class MeasurableMax(Max):
-    """A placeholder used to specify a log-likelihood for a cmax sub-graph."""
+    """A placeholder used to specify a log-likelihood for a max sub-graph."""
 
 
 MeasurableVariable.register(MeasurableMax)
 
 
 @node_rewriter([Max])
-def find_measurable_max(fgraph: FunctionGraph, node: Node) -> Optional[List[MeasurableMax]]:
+def find_measurable_max(fgraph: FunctionGraph, node: Node) -> Optional[List[TensorVariable]]:
     rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
     if rv_map_feature is None:
         return None  # pragma: no cover
@@ -73,25 +73,27 @@ def find_measurable_max(fgraph: FunctionGraph, node: Node) -> Optional[List[Meas
         return None  # pragma: no cover
 
     base_var = node.inputs[0]
-    pytensor.dprint(base_var)
 
     if base_var.owner is None:
         return None
 
-    # NonRVS must be rejected
-    if not isinstance(base_var.owner.op, RandomVariable):
+    if not rv_map_feature.request_measurable(node.inputs):
         return None
 
-    # univariate iid test which also rules out other distributions
-    if isinstance(base_var.owner.op, RandomVariable):
-        for params in base_var.owner.inputs[3:]:
-            if params.type.ndim != 0:
-                return None
+    # 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
 
-    if not rv_map_feature.request_measurable(node.inputs):
+    # TODO: We are currently only supporting continuous rvs
+    if isinstance(base_var.owner.op, RandomVariable) and base_var.owner.op.dtype.startswith("int"):
         return None
 
-    # Check whether axis is supported or not
+    # univariate i.i.d. test which also rules out other distributions
+    for params in base_var.owner.inputs[3:]:
+        if params.type.ndim != 0:
+            return None
+
+    # Check whether axis covers all dimensions
     axis = set(node.op.axis)
     base_var_dims = set(range(base_var.ndim))
     if axis != base_var_dims:
@@ -114,12 +116,7 @@ def find_measurable_max(fgraph: FunctionGraph, node: Node) -> Optional[List[Meas
 
 @_logprob.register(MeasurableMax)
 def max_logprob(op, values, base_rv, **kwargs):
-    r"""Compute the log-likelihood graph for the `Max` operation.
-
-    The formula that we use here is :
-        \ln(f_{(n)}(x)) = \ln(n) + (n-1) \ln(F(x)) + \ln(f(x))
-    where f(x) represents the p.d.f and F(x) represents the c.d.f of the distribution respectively.
-    """
+    r"""Compute the log-likelihood graph for the `Max` operation."""
     (value,) = values
 
     logprob = _logprob_helper(base_rv, value)
@@ -133,47 +130,58 @@ def max_logprob(op, values, base_rv, **kwargs):
 
 
 class MeasurableMin(Max):
-    """A placeholder used to specify a log-likelihood for a min sub-graph."""
+    """A placeholder used to specify a log-likelihood for a cmax sub-graph."""
 
 
 MeasurableVariable.register(MeasurableMin)
 
 
 @node_rewriter(tracks=[Max])
-def find_measurable_min(fgraph: FunctionGraph, node: Node) -> Optional[List[MeasurableMin]]:
+def find_measurable_min(fgraph: FunctionGraph, node: Node) -> Optional[List[MeasurableMax]]:
     rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
     if rv_map_feature is None:
         return None  # pragma: no cover
 
-    if isinstance(node.op, MeasurableVariable):
+    if isinstance(node.op, MeasurableMin):
         return None  # pragma: no cover
 
     base_var = node.inputs[0]
 
     if base_var.owner is None:
         return None
 
-    if not isinstance(base_var.owner.op, Elemwise):
+    if not rv_map_feature.request_measurable(node.inputs):
+        return None
+
+    # Non-univariate distributions must be rejected.
+    if not (
+        isinstance(base_var.owner.op, Elemwise) and base_var.owner.inputs[0].owner.op.ndim_supp == 0
+    ):
+        return None
+
+    if isinstance(base_var.owner.op, Elemwise) and base_var.owner.inputs[
+        0
+    ].owner.op.dtype.startswith("int"):
         return None
 
-    if isinstance(base_var.owner.op, Elemwise):
-        # check if min is -1 * rv
-        if len(base_var.owner.inputs) < 2:
+    for params in base_var.owner.inputs[3:]:
+        if params.type.ndim != 0:
             return None
-        if isinstance(base_var.owner.op.scalar_op, Mul):
-            if not isinstance(base_var.owner.inputs[1].owner.op, RandomVariable) or (
-                base_var.owner.inputs[0].value != -1
-            ):
-                return None
+
+    if not isinstance(base_var.owner.op.scalar_op, Mul):
+        return None
+
+    if not rv_map_feature.request_measurable(node.inputs):
+        return None
 
     # Check whether axis is supported or not
     axis = set(node.op.axis)
     base_var_dims = set(range(base_var.ndim))
     if axis != base_var_dims:
         return None
 
-    measurable_max = MeasurableMin(list(axis))
-    min_rv_node = measurable_max.make_node(base_var)
+    measurable_min = MeasurableMin(list(axis))
+    min_rv_node = measurable_min.make_node(base_var)
     min_rv = min_rv_node.outputs
 
     return min_rv

tests/logprob/test_order.py

@@ -47,18 +47,6 @@
 from pymc.testing import assert_no_rvs
 
 
-def test_max():
-    """Test whether the logprob for ```pt.max``` is implemented"""
-    x = pt.random.normal(0, 1, size=(3,))
-    x.name = "x"
-    x_max = pt.max(x, axis=-1)
-    # pytensor.dprint(x_max)
-    x_max_value = pt.vector("x_max_value")
-    x_max_logprob = logp(x_max, x_max_value)
-
-    assert_no_rvs(x_max_logprob)
-
-
 def test_min():
     """Test whether the logprob for ```pt.min``` is implemented"""
     x = pt.random.normal(0, 1, size=(3,))
@@ -70,17 +58,6 @@ def test_min():
     assert_no_rvs(x_min_logprob)
 
 
-def test_axis_max():
-    """Test whether the rewrite takes into account ```None``` axis"""
-    x = pt.random.normal(0, 1)
-    x.name = "x"
-    x_max = pt.max(x, axis=None)
-    x_max_value = pt.vector("x_max_value")
-    x_max_logprob = logp(x_max, x_max_value)
-
-    assert_no_rvs(x_max_logprob)
-
-
 def test_argmax():
     """Test whether the logprob for ```pt.argmax``` is rejected correctly"""
     x = pt.random.normal(0, 1, size=(3,))
@@ -112,6 +89,27 @@ def test_max_non_rv_fails():
         x_max_logprob = logp(x_max, x_max_value)
 
 
+def test_max_non_mul_elemwise_fails():
+    """Test whether the logprob for ```pt.max``` for non RVs is rejected correctly"""
+    x = pt.log(pt.random.beta(0, 1, size=(3,)))
+    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_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 rejected correctly"""
     x = pm.Categorical.dist([1, 1, 1, 1], shape=(5,))
@@ -137,7 +135,18 @@ def test_non_supp_axis_max():
     assert_no_rvs(x_max_logprob)
 
 
-def test_max_logprob():
+@pytest.mark.parametrize(
+    "shape, value, axis",
+    [
+        (3, 0.85, -1),
+        (3, 0.01, 0),
+        (2, 0.2, None),
+        (4, 0.5, 0),
+        ((3, 4), 0.9, None),
+        ((3, 4), 0.75, (1, 0)),
+    ],
+)
+def test_max_logprob(shape, value, axis):
     """Test whether the logprob for ```pt.max``` produces the corrected
 
     The fact that order statistics of i.i.d. uniform RVs ~ Beta is used here:
@@ -161,8 +170,19 @@ def test_max_logprob():
     )
 
 
-def test_min_logprob():
-    """Test whether the logprob for ```pt.max``` produces the corrected
+@pytest.mark.parametrize(
+    "shape, value, axis",
+    [
+        (3, 0.85, -1),
+        (3, 0.01, 0),
+        (2, 0.2, None),
+        (4, 0.5, 0),
+        ((3, 4), 0.9, None),
+        ((3, 4), 0.75, (1, 0)),
+    ],
+)
+def test_min_logprob(shape, value, axis):
+    """Test whether the logprob for ```pt.mix``` produces the corrected
 
     The fact that order statistics of i.i.d. uniform RVs ~ Beta is used here:
         U_1, \\dots, U_n \\stackrel{\text{i.i.d.}}{\\sim} \text{Uniform}(0, 1) \\Rightarrow U_{(k)} \\sim \text{Beta}(k, n + 1- k)