Add FutureWarning when accessing RV mappings via tag

Author: ricardoV94

pymc/distributions/distribution.py

@@ -49,7 +49,7 @@
     shape_from_dims,
 )
 from pymc.printing import str_for_dist
-from pymc.util import UNSET
+from pymc.util import UNSET, _add_future_warning_tag
 from pymc.vartypes import string_types
 
 __all__ = [
@@ -371,6 +371,7 @@ def dist(
         rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
         rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)", "pm.logcdf(rv, x)")
         rv_out.random = _make_nice_attr_error("rv.random()", "pm.draw(rv)")
+        _add_future_warning_tag(rv_out)
         return rv_out
 
 

pymc/distributions/shape_utils.py

@@ -50,6 +50,7 @@
 
 from pymc.aesaraf import PotentialShapeType
 from pymc.exceptions import ShapeError
+from pymc.util import _add_future_warning_tag
 
 
 def to_tuple(shape):
@@ -600,6 +601,7 @@ def change_dist_size(
         new_size = tuple(new_size)  # type: ignore
 
     new_dist = _change_dist_size(dist.owner.op, dist, new_size=new_size, expand=expand)
+    _add_future_warning_tag(new_dist)
 
     new_dist.name = dist.name
     for k, v in dist.tag.__dict__.items():

pymc/model.py

@@ -67,6 +67,7 @@
 from pymc.util import (
     UNSET,
     WithMemoization,
+    _add_future_warning_tag,
     get_transformed_name,
     get_value_vars_from_user_vars,
     get_var_name,
@@ -1314,6 +1315,7 @@ def register_rv(
         """
         name = self.name_for(name)
         rv_var.name = name
+        _add_future_warning_tag(rv_var)
         rv_var.tag.total_size = total_size
         self.rvs_to_total_sizes[rv_var] = total_size
 
@@ -1495,6 +1497,7 @@ def create_value_var(
         if aesara.config.compute_test_value != "off":
             value_var.tag.test_value = rv_var.tag.test_value
 
+        _add_future_warning_tag(value_var)
         rv_var.tag.value_var = value_var
 
         # Make the value variable a transformed value variable,

pymc/tests/distributions/test_distribution.py

@@ -28,8 +28,9 @@
 
 from pymc.distributions import DiracDelta, Flat, MvNormal, MvStudentT, logp
 from pymc.distributions.distribution import SymbolicRandomVariable, _moment, moment
-from pymc.distributions.shape_utils import to_tuple
+from pymc.distributions.shape_utils import change_dist_size, to_tuple
 from pymc.tests.distributions.util import assert_moment_is_expected
+from pymc.util import _FutureWarningValidatingScratchpad
 
 
 class TestBugfixes:
@@ -358,3 +359,40 @@ class TestSymbolicRV(SymbolicRandomVariable):
         dirac_delta_2_ = DiracDelta.dist(10)
         node = TestSymbolicRV([], [dirac_delta_1_, dirac_delta_2_], ndim_supp=0)().owner
         assert get_measurable_outputs(node.op, node) == [node.outputs[default_output_idx]]
+
+
+def test_tag_future_warning_dist():
+    # Test no unexpected warnings
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+
+        x = pm.Normal.dist()
+        assert isinstance(x.tag, _FutureWarningValidatingScratchpad)
+
+        x.tag.banana = "banana"
+        assert x.tag.banana == "banana"
+
+        # Check we didn't break test_value filtering
+        x.tag.test_value = np.array(1)
+        assert x.tag.test_value == 1
+        with pytest.raises(TypeError, match="Wrong number of dimensions"):
+            x.tag.test_value = np.array([1, 1])
+        assert x.tag.test_value == 1
+
+        # No warning if deprecated attribute is not present
+        with pytest.raises(AttributeError):
+            x.tag.value_var
+
+        # Warning if present
+        x.tag.value_var = "1"
+        with pytest.warns(FutureWarning, match="Use model.rvs_to_values"):
+            value_var = x.tag.value_var
+        assert value_var == "1"
+
+        # Check that PyMC method that copies tag contents does not erase special tag
+        new_x = change_dist_size(x, new_size=5)
+        assert new_x.tag is not x.tag
+        assert isinstance(new_x.tag, _FutureWarningValidatingScratchpad)
+        with pytest.warns(FutureWarning, match="Use model.rvs_to_values"):
+            value_var = new_x.tag.value_var
+        assert value_var == "1"

pymc/tests/test_model.py

@@ -41,10 +41,12 @@
 from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.distributions import Normal, transforms
 from pymc.distributions.logprob import _joint_logp
+from pymc.distributions.transforms import log
 from pymc.exceptions import ImputationWarning, ShapeError, ShapeWarning
 from pymc.model import Point, ValueGradFunction, modelcontext
 from pymc.tests.helpers import SeededTest
 from pymc.tests.models import simple_model
+from pymc.util import _FutureWarningValidatingScratchpad
 
 
 class NewModel(pm.Model):
@@ -1406,3 +1408,59 @@ def test_deterministic(self):
             assert np.all(
                 np.isclose(model.compile_logp(sum=False)({}), st.norm().logpdf(data_values))
             )
+
+
+def test_tag_future_warning_model():
+    # Test no unexpected warnings
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+
+        model = pm.Model()
+
+        x = at.random.normal()
+        x.tag.something_else = "5"
+        x.tag.test_value = 0
+        assert not isinstance(x.tag, _FutureWarningValidatingScratchpad)
+
+        # Test that model changes the tag type, but copies exsiting contents
+        x = model.register_rv(x, name="x", transform=log)
+        assert isinstance(x.tag, _FutureWarningValidatingScratchpad)
+        assert x.tag.something_else == "5"
+        assert x.tag.test_value == 0
+
+        # Test expected warnings
+        with pytest.warns(FutureWarning, match="model.rvs_to_values"):
+            x_value = x.tag.value_var
+
+        assert isinstance(x_value.tag, _FutureWarningValidatingScratchpad)
+        with pytest.warns(FutureWarning, match="model.rvs_to_transforms"):
+            transform = x_value.tag.transform
+        assert transform is log
+
+        with pytest.raises(AttributeError):
+            x.tag.observations
+
+        with pytest.warns(FutureWarning, match="model.rvs_to_total_sizes"):
+            total_size = x.tag.total_size
+        assert total_size is None
+
+        # Cloning a node will keep the same tag type and contents
+        y = x.owner.clone().default_output()
+        assert y is not x
+        assert y.tag is not x.tag
+        assert isinstance(y.tag, _FutureWarningValidatingScratchpad)
+        y = model.register_rv(y, name="y", data=5)
+        assert isinstance(y.tag, _FutureWarningValidatingScratchpad)
+
+        # Test expected warnings
+        with pytest.warns(FutureWarning, match="model.rvs_to_values"):
+            y_value = y.tag.value_var
+        with pytest.warns(FutureWarning, match="model.rvs_to_values"):
+            y_obs = y.tag.observations
+        assert y_value is y_obs
+        assert y_value.eval() == 5
+
+        assert isinstance(y_value.tag, _FutureWarningValidatingScratchpad)
+        with pytest.warns(FutureWarning, match="model.rvs_to_total_sizes"):
+            total_size = y.tag.total_size
+        assert total_size is None

pymc/util.py

@@ -13,6 +13,7 @@
 #   limitations under the License.
 
 import functools
+import warnings
 
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 
@@ -23,6 +24,7 @@
 
 from aesara import Variable
 from aesara.compile import SharedVariable
+from aesara.graph.utils import ValidatingScratchpad
 from cachetools import LRUCache, cachedmethod
 
 
@@ -481,3 +483,32 @@ def get_value_vars_from_user_vars(
         )
 
     return value_vars
+
+
+class _FutureWarningValidatingScratchpad(ValidatingScratchpad):
+    def __getattribute__(self, name):
+        for deprecated_names, alternative in (
+            (("value_var", "observations"), "model.rvs_to_values[rv]"),
+            (("transform",), "model.rvs_to_transforms[rv]"),
+            (("total_size",), "model.rvs_to_total_sizes[rv]"),
+        ):
+            if name in deprecated_names:
+                try:
+                    super().__getattribute__(name)
+                except AttributeError:
+                    pass
+                else:
+                    warnings.warn(
+                        f"The tag attribute {name} is deprecated. Use {alternative} instead",
+                        FutureWarning,
+                    )
+        return super().__getattribute__(name)
+
+
+def _add_future_warning_tag(var) -> None:
+    old_tag = var.tag
+    if not isinstance(old_tag, _FutureWarningValidatingScratchpad):
+        new_tag = _FutureWarningValidatingScratchpad("test_value", var.type.filter)
+        for k, v in old_tag.__dict__.items():
+            new_tag.__dict__.setdefault(k, v)
+        var.tag = new_tag