Replace deprecated tag.ignore_logprob

Author: ricardoV94

pymc/distributions/bound.py

@@ -23,7 +23,7 @@
 from pymc.distributions.continuous import BoundedContinuous, bounded_cont_transform
 from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import Continuous, Discrete
-from pymc.distributions.logprob import logp
+from pymc.distributions.logprob import ignore_logprob, logp
 from pymc.distributions.shape_utils import to_tuple
 from pymc.distributions.transforms import _default_transform
 from pymc.model import modelcontext
@@ -193,7 +193,7 @@ def __new__(
                 raise ValueError("Given dims do not exist in model coordinates.")
 
         lower, upper, initval = cls._set_values(lower, upper, size, shape, initval)
-        dist.tag.ignore_logprob = True
+        dist = ignore_logprob(dist)
 
         if isinstance(dist.owner.op, Continuous):
             res = _ContinuousBounded(
@@ -228,7 +228,7 @@ def dist(
 
         cls._argument_checks(dist, **kwargs)
         lower, upper, initval = cls._set_values(lower, upper, size, shape, initval=None)
-        dist.tag.ignore_logprob = True
+        dist = ignore_logprob(dist)
         if isinstance(dist.owner.op, Continuous):
             res = _ContinuousBounded.dist(
                 [dist, lower, upper],

pymc/distributions/logprob.py

@@ -20,6 +20,7 @@
 import numpy as np
 
 from aeppl import factorized_joint_logprob
+from aeppl.abstract import assign_custom_measurable_outputs
 from aeppl.logprob import logcdf as logcdf_aeppl
 from aeppl.logprob import logprob as logp_aeppl
 from aeppl.transforms import TransformValuesOpt
@@ -221,7 +222,11 @@ def joint_logpt(
 
     transform_opt = TransformValuesOpt(transform_map)
     temp_logp_var_dict = factorized_joint_logprob(
-        tmp_rvs_to_values, extra_rewrites=transform_opt, use_jacobian=jacobian, **kwargs
+        tmp_rvs_to_values,
+        extra_rewrites=transform_opt,
+        use_jacobian=jacobian,
+        warn_missing_rvs=False,
+        **kwargs,
     )
 
     # Raise if there are unexpected RandomVariables in the logp graph
@@ -276,3 +281,20 @@ def logcdf(rv, value):
 
     value = at.as_tensor_variable(value, dtype=rv.dtype)
     return logcdf_aeppl(rv, value)
+
+
+def ignore_logprob(rv):
+    """Return a duplicated variable that is ignored when creating Aeppl logprob graphs
+
+    This is used in SymbolicDistributions that use other RVs as inputs but account
+    for their logp terms explicitly.
+
+    If the variable is already ignored, it is returned directly.
+    """
+    prefix = "Unmeasurable"
+    node = rv.owner
+    op_type = type(node.op)
+    if op_type.__name__.startswith(prefix):
+        return rv
+    new_node = assign_custom_measurable_outputs(node, type_prefix=prefix)
+    return new_node.outputs[node.outputs.index(rv)]

pymc/distributions/mixture.py

@@ -30,7 +30,7 @@
 from pymc.distributions.continuous import Normal, get_tau_sigma
 from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import SymbolicDistribution, _moment, moment
-from pymc.distributions.logprob import logcdf, logp
+from pymc.distributions.logprob import ignore_logprob, logcdf, logp
 from pymc.distributions.shape_utils import to_tuple
 from pymc.distributions.transforms import _default_transform
 from pymc.util import check_dist_not_registered
@@ -252,6 +252,10 @@ def rv_op(cls, weights, *components, size=None, rngs=None):
 
         assert weights_ndim_batch == 0
 
+        # Component RVs terms are accounted by the Mixture logprob, so they can be
+        # safely ignored by Aeppl
+        components = [ignore_logprob(component) for component in components]
+
         # Create a OpFromGraph that encapsulates the random generating process
         # Create dummy input variables with the same type as the ones provided
         weights_ = weights.type()
@@ -299,11 +303,6 @@ def rv_op(cls, weights, *components, size=None, rngs=None):
         mix_out.tag.components = components
         mix_out.tag.choices_rng = mix_indexes_rng
 
-        # Component RVs terms are accounted by the Mixture logprob, so they can be
-        # safely ignore by Aeppl (this tag prevents UserWarning)
-        for component in components:
-            component.tag.ignore_logprob = True
-
         return mix_out
 
     @classmethod

pymc/distributions/multivariate.py

@@ -57,6 +57,7 @@
     multigammaln,
 )
 from pymc.distributions.distribution import Continuous, Discrete, moment
+from pymc.distributions.logprob import ignore_logprob
 from pymc.distributions.shape_utils import (
     broadcast_dist_samples_to,
     rv_size_is_none,
@@ -1182,11 +1183,9 @@ def dist(cls, eta, n, sd_dist, **kwargs):
 
         # sd_dist is part of the generative graph, but should be completely ignored
         # by the logp graph, since the LKJ logp explicitly includes these terms.
-        # Setting sd_dist.tag.ignore_logprob to True, will prevent Aeppl warning about
-        # an unnacounted RandomVariable in the graph
         # TODO: Things could be simplified a bit if we managed to extract the
         #  sd_dist prior components from the logp expression.
-        sd_dist.tag.ignore_logprob = True
+        sd_dist = ignore_logprob(sd_dist)
 
         return super().dist([n, eta, sd_dist], **kwargs)
 

pymc/distributions/timeseries.py

@@ -26,6 +26,7 @@
 from pymc.distributions import distribution, logprob, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
 from pymc.distributions.dist_math import check_parameters
+from pymc.distributions.logprob import ignore_logprob
 from pymc.distributions.shape_utils import rv_size_is_none, to_tuple
 from pymc.util import check_dist_not_registered
 
@@ -206,7 +207,7 @@ def dist(
                 raise TypeError("init must be a univariate distribution variable")
 
         # Ignores logprob of init var because that's accounted for in the logp method
-        init.tag.ignore_logprob = True
+        init = ignore_logprob(init)
 
         return super().dist([mu, sigma, init, steps], size=size, **kwargs)
 

pymc/tests/test_logprob.py

@@ -17,6 +17,7 @@
 import pytest
 import scipy.stats.distributions as sp
 
+from aeppl.abstract import get_measurable_outputs
 from aesara.graph.basic import ancestors
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.subtensor import (
@@ -32,7 +33,7 @@
 from pymc.aesaraf import floatX, walk_model
 from pymc.distributions.continuous import HalfFlat, Normal, TruncatedNormal, Uniform
 from pymc.distributions.discrete import Bernoulli
-from pymc.distributions.logprob import joint_logpt, logcdf, logp
+from pymc.distributions.logprob import ignore_logprob, joint_logpt, logcdf, logp
 from pymc.model import Model, Potential
 from pymc.tests.helpers import select_by_precision
 
@@ -227,3 +228,38 @@ def test_unexpected_rvs():
 
     with pytest.raises(ValueError, match="^Random variables detected in the logp graph"):
         model.logpt()
+
+
+def test_ignore_logprob_basic():
+    x = Normal.dist()
+    (measurable_x_out,) = get_measurable_outputs(x.owner.op, x.owner)
+    assert measurable_x_out is x.owner.outputs[1]
+
+    new_x = ignore_logprob(x)
+    assert new_x is not x
+    assert isinstance(new_x.owner.op, Normal)
+    assert type(new_x.owner.op).__name__ == "UnmeasurableNormalRV"
+    # Confirm that it does not have measurable output
+    assert get_measurable_outputs(new_x.owner.op, new_x.owner) is None
+
+    # Test that it will not clone a variable that is already unmeasurable
+    new_new_x = ignore_logprob(new_x)
+    assert new_new_x is new_x
+
+
+def test_ignore_logprob_model():
+    # logp that does not depend on input
+    def logp(value, x):
+        return value
+
+    with Model() as m:
+        x = Normal.dist()
+        y = DensityDist("y", x, logp=logp)
+    # Aeppl raises a KeyError when it finds an unexpected RV
+    with pytest.raises(KeyError):
+        joint_logpt([y], {y: y.type()})
+
+    with Model() as m:
+        x = ignore_logprob(Normal.dist())
+        y = DensityDist("y", x, logp=logp)
+    assert joint_logpt([y], {y: y.type()})