Generalize logp of min to neg(max)

Author: Dhruvanshu-Joshi

pymc/logprob/order.py

@@ -129,11 +129,11 @@ def max_logprob(op, values, base_rv, **kwargs):
     return logprob
 
 
-class MeasurableMin(Max):
+class MeasurableMaxNeg(Max):
     """A placeholder used to specify a log-likelihood for a min sub-graph."""
 
 
-MeasurableVariable.register(MeasurableMin)
+MeasurableVariable.register(MeasurableMaxNeg)
 
 
 @node_rewriter(tracks=[Max])
@@ -143,7 +143,7 @@ def find_measurable_min(fgraph: FunctionGraph, node: Node) -> Optional[List[Tens
     if rv_map_feature is None:
         return None  # pragma: no cover
 
-    if isinstance(node.op, MeasurableMin):
+    if isinstance(node.op, MeasurableMaxNeg):
         return None  # pragma: no cover
 
     base_var = node.inputs[0]
@@ -183,7 +183,7 @@ def find_measurable_min(fgraph: FunctionGraph, node: Node) -> Optional[List[Tens
     if axis != base_var_dims:
         return None
 
-    measurable_min = MeasurableMin(list(axis))
+    measurable_min = MeasurableMaxNeg(list(axis))
     min_rv_node = measurable_min.make_node(base_var)
     min_rv = min_rv_node.outputs
 
@@ -198,7 +198,7 @@ def find_measurable_min(fgraph: FunctionGraph, node: Node) -> Optional[List[Tens
 )
 
 
-@_logprob.register(MeasurableMin)
+@_logprob.register(MeasurableMaxNeg)
 def min_logprob(op, values, base_var, **kwargs):
     r"""Compute the log-likelihood graph for the `Max` operation.
     The formula that we use here is :
@@ -208,8 +208,8 @@ def min_logprob(op, values, base_var, **kwargs):
     (value,) = values
     base_rv = base_var.owner.inputs[0]
 
-    logprob = _logprob_helper(base_rv, value)
-    logcdf = _logcdf_helper(base_rv, value)
+    logprob = _logprob_helper(base_rv, -value)
+    logcdf = _logcdf_helper(base_rv, -value)
 
     [n] = constant_fold([base_rv.size])
     logprob = (n - 1) * pt.math.log(1 - pt.math.exp(logcdf)) + logprob + pt.math.log(n)

tests/logprob/test_order.py

@@ -72,8 +72,7 @@ def test_non_iid_fails(if_max):
         x_m = pt.max(x, axis=-1)
         x_m_value = pt.vector("x_max_value")
     else:
-        x_min = pt.min(x, axis=-1)
-        x_m = x_min.owner.inputs[0]
+        x_m = pt.min(x, axis=-1)
         x_m_value = pt.vector("x_min_value")
     with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")):
         x_max_logprob = logp(x_m, x_m_value)
@@ -91,8 +90,7 @@ def test_non_rv_fails(if_max):
         x_m = pt.max(x, axis=-1)
         x_m_value = pt.vector("x_max_value")
     else:
-        x_min = pt.min(x, axis=-1)
-        x_m = x_min.owner.inputs[0]
+        x_m = pt.min(x, axis=-1)
         x_m_value = pt.vector("x_min_value")
     with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")):
         x_max_logprob = logp(x_m, x_m_value)
@@ -114,8 +112,7 @@ def test_multivariate_rv_fails(if_max):
         x_m = pt.max(x, axis=-1)
         x_m_value = pt.vector("x_max_value")
     else:
-        x_min = pt.min(x, axis=-1)
-        x_m = x_min.owner.inputs[0]
+        x_m = pt.min(x, axis=-1)
         x_m_value = pt.vector("x_min_value")
     with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")):
         x_max_logprob = logp(x_m, x_m_value)
@@ -136,8 +133,7 @@ def test_categorical(if_max):
         x_m = pt.max(x, axis=-1)
         x_m_value = pt.vector("x_max_value")
     else:
-        x_min = pt.min(x, axis=-1)
-        x_m = x_min.owner.inputs[0]
+        x_m = pt.min(x, axis=-1)
         x_m_value = pt.vector("x_min_value")
     with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")):
         x_max_logprob = logp(x_m, x_m_value)
@@ -158,8 +154,7 @@ def test_non_supp_axis(if_max):
         x_m = pt.max(x, axis=-1)
         x_m_value = pt.vector("x_max_value")
     else:
-        x_min = pt.min(x, axis=-1)
-        x_m = x_min.owner.inputs[0]
+        x_m = pt.min(x, axis=-1)
         x_m_value = pt.vector("x_min_value")
     with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")):
         x_max_logprob = logp(x_m, x_m_value)
@@ -225,9 +220,8 @@ def test_min_logprob(shape, value, axis):
     x = pt.random.uniform(0, 1, size=shape)
     x.name = "x"
     x_min = pt.min(x, axis=axis)
-    x_min_rv = x_min.owner.inputs[0]
     x_min_value = pt.scalar("x_min_value")
-    x_min_logprob = logp(x_min_rv, x_min_value)
+    x_min_logprob = logp(x_min, x_min_value)
 
     assert_no_rvs(x_min_logprob)
 
@@ -250,7 +244,6 @@ def test_min_non_mul_elemwise_fails():
     x = pt.log(pt.random.beta(0, 1, size=(3,)))
     x.name = "x"
     x_min = pt.min(x, axis=-1)
-    x_min_rv = x_min.owner.inputs[0]
     x_min_value = pt.vector("x_min_value")
     with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")):
-        x_min_logprob = logp(x_min_rv, x_min_value)
+        x_min_logprob = logp(x_min, x_min_value)