Add helper to undo ignore_logprob

Author: ricardoV94

pymc/distributions/logprob.py

@@ -11,8 +11,7 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-
-
+from copy import copy
 from typing import Dict, List, Sequence, Union
 
 import numpy as np
@@ -210,3 +209,23 @@ def ignore_logprob(rv: TensorVariable) -> TensorVariable:
         return rv
     new_node = assign_custom_measurable_outputs(node, type_prefix=prefix)
     return new_node.outputs[node.outputs.index(rv)]
+
+
+def reconsider_logprob(rv: TensorVariable) -> TensorVariable:
+    """Return a duplicated variable that is considered when creating logprob graphs
+
+    This undoes the effect of `ignore_logprob`.
+
+    If a variable was not ignored, it is returned directly.
+    """
+    prefix = "Unmeasurable"
+    node = rv.owner
+    op_type = type(node.op)
+    if not op_type.__name__.startswith(prefix):
+        return rv
+
+    new_node = node.clone()
+    original_op_type = new_node.op.original_op_type
+    new_node.op = copy(new_node.op)
+    new_node.op.__class__ = original_op_type
+    return new_node.outputs[node.outputs.index(rv)]

pymc/logprob/abstract.py

@@ -220,6 +220,7 @@ def assign_custom_measurable_outputs(
 
     new_op_dict = op_type.__dict__.copy()
     new_op_dict["id_obj"] = (new_node.op, measurable_outputs_fn)
+    new_op_dict.setdefault("original_op_type", op_type)
 
     new_op_type = type(
         f"{type_prefix}{op_type.__name__}", (op_type, UnmeasurableVariable), new_op_dict

pymc/tests/distributions/test_logprob.py

@@ -47,6 +47,7 @@
     ignore_logprob,
     logcdf,
     logp,
+    reconsider_logprob,
 )
 from pymc.logprob.abstract import get_measurable_outputs
 from pymc.model import Model, Potential
@@ -315,7 +316,7 @@ def test_unexpected_rvs():
         model.logp()
 
 
-def test_ignore_logprob_basic():
+def test_ignore_reconsider_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]
@@ -328,18 +329,34 @@ def test_ignore_logprob_basic():
     assert get_measurable_outputs(new_x.owner.op, new_x.owner) == []
 
     # 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
+    assert ignore_logprob(new_x) is new_x
+
+    orig_x = reconsider_logprob(new_x)
+    assert orig_x is not new_x
+    assert isinstance(orig_x.owner.op, Normal)
+    assert type(orig_x.owner.op).__name__ == "NormalRV"
+    # Confirm that it has measurable outputs again
+    assert get_measurable_outputs(orig_x.owner.op, orig_x.owner) == [orig_x.owner.outputs[1]]
+
+    # Test that will not clone a variable that is already measurable
+    assert reconsider_logprob(x) is x
+    assert reconsider_logprob(orig_x) is orig_x
+
+
+def test_ignore_reconsider_logprob_model():
+    def custom_logp(value, x):
+        # custom_logp is just the logp of x at value
+        x = reconsider_logprob(x)
+        return _joint_logp(
+            [x],
+            rvs_to_values={x: value},
+            rvs_to_transforms={},
+            rvs_to_total_sizes={},
+        )
 
     with Model() as m:
         x = Normal.dist()
-        y = CustomDist("y", x, logp=logp)
+        y = CustomDist("y", x, logp=custom_logp)
     with pytest.warns(
         UserWarning,
         match="Found a random variable that was neither among the observations "
@@ -355,7 +372,7 @@ def logp(value, x):
     # The above warning should go away with ignore_logprob.
     with Model() as m:
         x = ignore_logprob(Normal.dist())
-        y = CustomDist("y", x, logp=logp)
+        y = CustomDist("y", x, logp=custom_logp)
     with warnings.catch_warnings():
         warnings.simplefilter("error")
         assert _joint_logp(