Do not rely on tag information for rv and logp conversions

Author: ricardoV94

pymc/aesaraf.py

@@ -11,6 +11,8 @@
 #   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.
+import warnings
+
 from typing import (
     Callable,
     Dict,
@@ -31,6 +33,7 @@
 import scipy.sparse as sps
 
 from aeppl.logprob import CheckParameterValue
+from aeppl.transforms import RVTransform
 from aesara import scalar
 from aesara.compile.mode import Mode, get_mode
 from aesara.gradient import grad
@@ -205,7 +208,7 @@ def expand(var):
     yield from walk(graphs, expand, bfs=False)
 
 
-def replace_rvs_in_graphs(
+def _replace_rvs_in_graphs(
     graphs: Iterable[TensorVariable],
     replacement_fn: Callable[[TensorVariable], Dict[TensorVariable, TensorVariable]],
     **kwargs,
@@ -282,6 +285,10 @@ def rvs_to_value_vars(
     apply_transforms
         If ``True``, apply each value variable's transform.
     """
+    warnings.warn(
+        "rvs_to_value_vars is deprecated. Use model.replace_rvs_by_values instead",
+        FutureWarning,
+    )
 
     def populate_replacements(
         random_var: TensorVariable, replacements: Dict[TensorVariable, TensorVariable]
@@ -313,7 +320,7 @@ def populate_replacements(
     equiv = clone_get_equiv(inputs, graphs, False, False, {})
     graphs = [equiv[n] for n in graphs]
 
-    graphs, _ = replace_rvs_in_graphs(
+    graphs, _ = _replace_rvs_in_graphs(
         graphs,
         replacement_fn=populate_replacements,
         **kwargs,
@@ -322,6 +329,69 @@ def populate_replacements(
     return graphs
 
 
+def replace_rvs_by_values(
+    graphs: Sequence[TensorVariable],
+    *,
+    rvs_to_values: Dict[TensorVariable, TensorVariable],
+    rvs_to_transforms: Dict[TensorVariable, RVTransform],
+    **kwargs,
+) -> List[TensorVariable]:
+    """Clone and replace random variables in graphs with their value variables.
+
+    This will *not* recompute test values in the resulting graphs.
+
+    Parameters
+    ----------
+    graphs
+        The graphs in which to perform the replacements.
+    rvs_to_values
+        Mapping between the original graph RVs and respective value variables
+    rvs_to_transforms
+        Mapping between the original graph RVs and respective value transforms
+    """
+
+    # Clone original graphs so that we don't modify variables in place
+    inputs = [i for i in graph_inputs(graphs) if not isinstance(i, Constant)]
+    equiv = clone_get_equiv(inputs, graphs, False, False, {})
+    graphs = [equiv[n] for n in graphs]
+
+    # Get needed mappings for equivalent cloned variables
+    equiv_rvs_to_values = {}
+    equiv_rvs_to_transforms = {}
+    for rv, value in rvs_to_values.items():
+        equiv_rv = equiv.get(rv, rv)
+        equiv_rvs_to_values[equiv_rv] = equiv.get(value, value)
+        equiv_rvs_to_transforms[equiv_rv] = rvs_to_transforms[rv]
+
+    def poulate_replacements(rv, replacements):
+        # Populate replacements dict with {rv: value} pairs indicating which graph
+        # RVs should be replaced by what value variables.
+
+        # No value variable to replace RV with
+        value = equiv_rvs_to_values.get(rv, None)
+        if value is None:
+            return []
+
+        transform = equiv_rvs_to_transforms.get(rv, None)
+        if transform is not None:
+            # We want to replace uses of the RV by the back-transformation of its value
+            value = transform.backward(value, *rv.owner.inputs)
+            value.name = rv.name
+
+        replacements[rv] = value
+        # Also walk the graph of the value variable to make any additional
+        # replacements if that is not a simple input variable
+        return [value]
+
+    graphs, _ = _replace_rvs_in_graphs(
+        graphs,
+        replacement_fn=poulate_replacements,
+        **kwargs,
+    )
+
+    return graphs
+
+
 def inputvars(a):
     """
     Get the inputs into Aesara variables

pymc/backends/arviz.py

@@ -47,7 +47,7 @@ def find_observations(model: "Model") -> Dict[str, Var]:
     """If there are observations available, return them as a dictionary."""
     observations = {}
     for obs in model.observed_RVs:
-        aux_obs = getattr(obs.tag, "observations", None)
+        aux_obs = model.rvs_to_values.get(obs, None)
         if aux_obs is not None:
             try:
                 obs_data = extract_obs_data(aux_obs)
@@ -261,7 +261,7 @@ def log_likelihood_vals_point(self, point, var, log_like_fun):
 
         if isinstance(var.owner.op, (AdvancedIncSubtensor, AdvancedIncSubtensor1)):
             try:
-                obs_data = extract_obs_data(var.tag.observations)
+                obs_data = extract_obs_data(self.model.rvs_to_values[var])
             except TypeError:
                 warnings.warn(f"Could not extract data from symbolic observation {var}")
 

pymc/distributions/__init__.py

@@ -16,7 +16,6 @@
     logcdf,
     logp,
     joint_logp,
-    joint_logpt,
 )
 
 from pymc.distributions.bound import Bound
@@ -199,7 +198,6 @@
     "Censored",
     "CAR",
     "PolyaGamma",
-    "joint_logpt",
     "joint_logp",
     "logp",
     "logcdf",

pymc/distributions/logprob.py

@@ -25,18 +25,18 @@
 from aeppl.logprob import logcdf as logcdf_aeppl
 from aeppl.logprob import logprob as logp_aeppl
 from aeppl.tensor import MeasurableJoin
-from aeppl.transforms import TransformValuesRewrite
+from aeppl.transforms import RVTransform, TransformValuesRewrite
 from aesara import tensor as at
 from aesara.graph.basic import graph_inputs, io_toposort
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.var import TensorVariable
 
 from pymc.aesaraf import constant_fold, floatX
 
+TOTAL_SIZE = Union[int, Sequence[int], None]
 
-def _get_scaling(
-    total_size: Optional[Union[int, Sequence[int]]], shape, ndim: int
-) -> TensorVariable:
+
+def _get_scaling(total_size: TOTAL_SIZE, shape, ndim: int) -> TensorVariable:
     """
     Gets scaling constant for logp.
 
@@ -104,12 +104,26 @@ def _get_scaling(
     return at.as_tensor(coef, dtype=aesara.config.floatX)
 
 
-def joint_logpt(*args, **kwargs):
-    warnings.warn(
-        "joint_logpt has been deprecated. Use joint_logp instead.",
-        FutureWarning,
-    )
-    return joint_logp(*args, **kwargs)
+def _check_no_rvs(logp_terms: Sequence[TensorVariable]):
+    # Raise if there are unexpected RandomVariables in the logp graph
+    # Only SimulatorRVs are allowed
+    from pymc.distributions.simulator import SimulatorRV
+
+    unexpected_rv_nodes = [
+        node
+        for node in aesara.graph.ancestors(logp_terms)
+        if (
+            node.owner
+            and isinstance(node.owner.op, RandomVariable)
+            and not isinstance(node.owner.op, SimulatorRV)
+        )
+    ]
+    if unexpected_rv_nodes:
+        raise ValueError(
+            f"Random variables detected in the logp graph: {unexpected_rv_nodes}.\n"
+            "This can happen when DensityDist logp or Interval transform functions "
+            "reference nonlocal variables."
+        )
 
 
 def joint_logp(
@@ -151,6 +165,10 @@ def joint_logp(
         Sum the log-likelihood or return each term as a separate list item.
 
     """
+    warnings.warn(
+        "joint_logp has been deprecated, use model.logp instead",
+        FutureWarning,
+    )
     # TODO: In future when we drop support for tag.value_var most of the following
     # logic can be removed and logp can just be a wrapper function that calls aeppl's
     # joint_logprob directly.
@@ -223,33 +241,15 @@ def joint_logp(
         **kwargs,
     )
 
-    # Raise if there are unexpected RandomVariables in the logp graph
-    # Only SimulatorRVs are allowed
-    from pymc.distributions.simulator import SimulatorRV
-
-    unexpected_rv_nodes = [
-        node
-        for node in aesara.graph.ancestors(list(temp_logp_var_dict.values()))
-        if (
-            node.owner
-            and isinstance(node.owner.op, RandomVariable)
-            and not isinstance(node.owner.op, SimulatorRV)
-        )
-    ]
-    if unexpected_rv_nodes:
-        raise ValueError(
-            f"Random variables detected in the logp graph: {unexpected_rv_nodes}.\n"
-            "This can happen when DensityDist logp or Interval transform functions "
-            "reference nonlocal variables."
-        )
-
     # aeppl returns the logp for every single value term we provided to it. This includes
     # the extra values we plugged in above, so we filter those we actually wanted in the
     # same order they were given in.
     logp_var_dict = {}
     for value_var in rv_values.values():
         logp_var_dict[value_var] = temp_logp_var_dict[value_var]
 
+    _check_no_rvs(list(logp_var_dict.values()))
+
     if scaling:
         for value_var in logp_var_dict.keys():
             if value_var in rv_scalings:
@@ -263,6 +263,52 @@ def joint_logp(
     return logp_var
 
 
+def _joint_logp(
+    rvs: Sequence[TensorVariable],
+    *,
+    rvs_to_values: Dict[TensorVariable, TensorVariable],
+    rvs_to_transforms: Dict[TensorVariable, RVTransform],
+    jacobian: bool = True,
+    rvs_to_total_sizes: Dict[TensorVariable, TOTAL_SIZE],
+    **kwargs,
+) -> List[TensorVariable]:
+    """Thin wrapper around aeppl.factorized_joint_logprob, extended with PyMC specific
+    concerns such as transforms, jacobian, and scaling"""
+
+    transform_rewrite = None
+    values_to_transforms = {
+        rvs_to_values[rv]: transform
+        for rv, transform in rvs_to_transforms.items()
+        if transform is not None
+    }
+    if values_to_transforms:
+        # There seems to be an incorrect type hint in TransformValuesRewrite
+        transform_rewrite = TransformValuesRewrite(values_to_transforms)  # type: ignore
+
+    temp_logp_terms = factorized_joint_logprob(
+        rvs_to_values,
+        extra_rewrites=transform_rewrite,
+        use_jacobian=jacobian,
+        **kwargs,
+    )
+
+    # aeppl returns the logp for every single value term we provided to it. This includes
+    # the extra values we plugged in above, so we filter those we actually wanted in the
+    # same order they were given in.
+    logp_terms = {}
+    for rv in rvs:
+        value_var = rvs_to_values[rv]
+        logp_term = temp_logp_terms[value_var]
+        total_size = rvs_to_total_sizes.get(rv, None)
+        if total_size is not None:
+            scaling = _get_scaling(total_size, value_var.shape, value_var.ndim)
+            logp_term *= scaling
+        logp_terms[value_var] = logp_term
+
+    _check_no_rvs(list(logp_terms.values()))
+    return list(logp_terms.values())
+
+
 def logp(rv: TensorVariable, value) -> TensorVariable:
     """Return the log-probability graph of a Random Variable"""
 

pymc/initial_point.py

@@ -20,6 +20,7 @@
 import aesara.tensor as at
 import numpy as np
 
+from aeppl.transforms import RVTransform
 from aesara.graph.basic import Variable
 from aesara.graph.fg import FunctionGraph
 from aesara.tensor.var import TensorVariable
@@ -43,9 +44,7 @@ def convert_str_to_rv_dict(
         if isinstance(key, str):
             if is_transformed_name(key):
                 rv = model[get_untransformed_name(key)]
-                initvals[rv] = model.rvs_to_values[rv].tag.transform.backward(
-                    initval, *rv.owner.inputs
-                )
+                initvals[rv] = model.rvs_to_transforms[rv].backward(initval, *rv.owner.inputs)
             else:
                 initvals[model[key]] = initval
         else:
@@ -158,7 +157,7 @@ def make_initial_point_fn(
 
     initial_values = make_initial_point_expression(
         free_rvs=model.free_RVs,
-        rvs_to_values=model.rvs_to_values,
+        rvs_to_transforms=model.rvs_to_transforms,
         initval_strategies=initval_strats,
         jitter_rvs=jitter_rvs,
         default_strategy=default_strategy,
@@ -172,7 +171,7 @@ def make_initial_point_fn(
 
     varnames = []
     for var in model.free_RVs:
-        transform = getattr(model.rvs_to_values[var].tag, "transform", None)
+        transform = model.rvs_to_transforms[var]
         if transform is not None and return_transformed:
             name = get_transformed_name(var.name, transform)
         else:
@@ -197,7 +196,7 @@ def inner(seed, *args, **kwargs):
 def make_initial_point_expression(
     *,
     free_rvs: Sequence[TensorVariable],
-    rvs_to_values: Dict[TensorVariable, TensorVariable],
+    rvs_to_transforms: Dict[TensorVariable, RVTransform],
     initval_strategies: Dict[TensorVariable, Optional[Union[np.ndarray, Variable, str]]],
     jitter_rvs: Set[TensorVariable] = None,
     default_strategy: str = "moment",
@@ -265,7 +264,7 @@ def make_initial_point_expression(
         else:
             value = at.as_tensor(strategy, dtype=variable.dtype).astype(variable.dtype)
 
-        transform = getattr(rvs_to_values[variable].tag, "transform", None)
+        transform = rvs_to_transforms.get(variable, None)
 
         if transform is not None:
             value = transform.forward(value, *variable.owner.inputs)