Remove default transformed RV Ops

This functionality was not being used by PyMC. In theory, it could be used to cache one type of transformed RV Op instead of recreating it for every instance of the same RV + transform, resulting in a faster logp rewrite (?). We can always revisit this idea later, but for now it unnecessarily complicated our codebase.

Author: ricardoV94

pymc/logprob/transforms.py

@@ -37,7 +37,6 @@
 import abc
 
 from copy import copy
-from functools import partial, singledispatch
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import pytensor.tensor as at
@@ -69,21 +68,6 @@
 from pymc.logprob.utils import walk_model
 
 
-@singledispatch
-def _default_transformed_rv(
-    op: Op,
-    node: Node,
-) -> Optional[Apply]:
-    """Create a node for a transformed log-probability of a `MeasurableVariable`.
-
-    This function dispatches on the type of `op`.  If you want to implement
-    new transforms for a `MeasurableVariable`, register a function on this
-    dispatcher.
-
-    """
-    return None
-
-
 class TransformedVariable(Op):
     """A no-op that identifies a transform and its un-transformed input."""
 
@@ -136,13 +120,6 @@ def log_jac_det(self, value: TensorVariable, *inputs) -> TensorVariable:
         return at.log(at.abs(at.nlinalg.det(at.atleast_2d(jacobian(phi_inv, [value])[0]))))
 
 
-class DefaultTransformSentinel:
-    pass
-
-
-DEFAULT_TRANSFORM = DefaultTransformSentinel()
-
-
 @node_rewriter(tracks=None)
 def transform_values(fgraph: FunctionGraph, node: Node) -> Optional[List[Node]]:
     """Apply transforms to value variables.
@@ -176,17 +153,12 @@ def transform_values(fgraph: FunctionGraph, node: Node) -> Optional[List[Node]]:
 
     if transform is None:
         return None
-    elif transform is DEFAULT_TRANSFORM:
-        trans_node = _default_transformed_rv(node.op, node)
-        if trans_node is None:
-            return None
-        transform = trans_node.op.transform
-    else:
-        new_op = _create_transformed_rv_op(node.op, transform)
-        # Create a new `Apply` node and outputs
-        trans_node = node.clone()
-        trans_node.op = new_op
-        trans_node.outputs[rv_var_out_idx].name = node.outputs[rv_var_out_idx].name
+
+    new_op = _create_transformed_rv_op(node.op, transform)
+    # Create a new `Apply` node and outputs
+    trans_node = node.clone()
+    trans_node.op = new_op
+    trans_node.outputs[rv_var_out_idx].name = node.outputs[rv_var_out_idx].name
 
     # We now assume that the old value variable represents the *transformed space*.
     # This means that we need to replace all instance of the old value variable
@@ -216,24 +188,22 @@ def on_attach(self, fgraph):
 
 
 class TransformValuesRewrite(GraphRewriter):
-    r"""Transforms value variables according to a map and/or per-`RandomVariable` defaults."""
+    r"""Transforms value variables according to a map."""
 
-    default_transform_rewrite = in2out(transform_values, ignore_newtrees=True)
+    transform_rewrite = in2out(transform_values, ignore_newtrees=True)
 
     def __init__(
         self,
-        values_to_transforms: Dict[
-            TensorVariable, Union[RVTransform, DefaultTransformSentinel, None]
-        ],
+        values_to_transforms: Dict[TensorVariable, Union[RVTransform, None]],
     ):
         """
         Parameters
         ==========
         values_to_transforms
             Mapping between value variables and their transformations.  Each
-            value variable can be assigned one of `RVTransform`,
-            ``DEFAULT_TRANSFORM``, or ``None``. If a transform is not specified
-            for a specific value variable it will not be transformed.
+            value variable can be assigned one of `RVTransform`, or ``None``.
+            If a transform is not specified for a specific value variable it will
+            not be transformed.
 
         """
 
@@ -244,7 +214,7 @@ def add_requirements(self, fgraph):
         fgraph.attach_feature(values_transforms_feature)
 
     def apply(self, fgraph: FunctionGraph):
-        return self.default_transform_rewrite.rewrite(fgraph)
+        return self.transform_rewrite.rewrite(fgraph)
 
 
 class MeasurableTransform(MeasurableElemwise):
@@ -583,7 +553,6 @@ def _create_transformed_rv_op(
     rv_op: Op,
     transform: RVTransform,
     *,
-    default: bool = False,
     cls_dict_extra: Optional[Dict] = None,
 ) -> Op:
     """Create a new transformed variable instance given a base `RandomVariable` `Op`.
@@ -600,8 +569,6 @@ def _create_transformed_rv_op(
         The `RandomVariable` for which we want to construct a `TransformedRV`.
     transform
         The `RVTransform` for `rv_op`.
-    default
-        If ``False`` do not make `transform` the default transform for `rv_op`.
     cls_dict_extra
         Additional class members to add to the constructed `TransformedRV`.
 
@@ -642,85 +609,7 @@ def transformed_logprob(op, values, *inputs, use_jacobian=True, **kwargs):
 
         return logprob
 
-    transform_op = rv_op_type if default else new_op_type
-
-    @_default_transformed_rv.register(transform_op)
-    def class_transformed_rv(op, node):
-        new_op = new_op_type()
-        res = new_op.make_node(*node.inputs)
-        res.outputs[1].name = node.outputs[1].name
-        return res
-
     new_op = copy(rv_op)
     new_op.__class__ = new_op_type
 
     return new_op
-
-
-create_default_transformed_rv_op = partial(_create_transformed_rv_op, default=True)
-
-
-TransformedUniformRV = create_default_transformed_rv_op(
-    at.random.uniform,
-    # inputs[3] = lower; inputs[4] = upper
-    IntervalTransform(lambda *inputs: (inputs[3], inputs[4])),
-)
-TransformedParetoRV = create_default_transformed_rv_op(
-    at.random.pareto,
-    # inputs[3] = alpha
-    IntervalTransform(lambda *inputs: (inputs[3], None)),
-)
-TransformedTriangularRV = create_default_transformed_rv_op(
-    at.random.triangular,
-    # inputs[3] = lower; inputs[5] = upper
-    IntervalTransform(lambda *inputs: (inputs[3], inputs[5])),
-)
-TransformedHalfNormalRV = create_default_transformed_rv_op(
-    at.random.halfnormal,
-    # inputs[3] = loc
-    IntervalTransform(lambda *inputs: (inputs[3], None)),
-)
-TransformedWaldRV = create_default_transformed_rv_op(
-    at.random.wald,
-    LogTransform(),
-)
-TransformedExponentialRV = create_default_transformed_rv_op(
-    at.random.exponential,
-    LogTransform(),
-)
-TransformedLognormalRV = create_default_transformed_rv_op(
-    at.random.lognormal,
-    LogTransform(),
-)
-TransformedHalfCauchyRV = create_default_transformed_rv_op(
-    at.random.halfcauchy,
-    LogTransform(),
-)
-TransformedGammaRV = create_default_transformed_rv_op(
-    at.random.gamma,
-    LogTransform(),
-)
-TransformedInvGammaRV = create_default_transformed_rv_op(
-    at.random.invgamma,
-    LogTransform(),
-)
-TransformedChiSquareRV = create_default_transformed_rv_op(
-    at.random.chisquare,
-    LogTransform(),
-)
-TransformedWeibullRV = create_default_transformed_rv_op(
-    at.random.weibull,
-    LogTransform(),
-)
-TransformedBetaRV = create_default_transformed_rv_op(
-    at.random.beta,
-    LogOddsTransform(),
-)
-TransformedVonMisesRV = create_default_transformed_rv_op(
-    at.random.vonmises,
-    CircularTransform(),
-)
-TransformedDirichletRV = create_default_transformed_rv_op(
-    at.random.dirichlet,
-    SimplexTransform(),
-)

pymc/tests/logprob/test_transforms.py

@@ -45,9 +45,9 @@
 from pytensor.graph.basic import equal_computations
 from pytensor.graph.fg import FunctionGraph
 
+from pymc.distributions.transforms import _default_transform, log, logodds
 from pymc.logprob.joint_logprob import factorized_joint_logprob, joint_logprob
 from pymc.logprob.transforms import (
-    DEFAULT_TRANSFORM,
     ChainedTransform,
     ExpTransform,
     IntervalTransform,
@@ -58,7 +58,6 @@
     ScaleTransform,
     TransformValuesMapping,
     TransformValuesRewrite,
-    _default_transformed_rv,
     transformed_variable,
 )
 from pymc.tests.helpers import assert_no_rvs
@@ -216,6 +215,8 @@ def test_transformed_logprob(at_dist, dist_params, sp_dist, size):
     elsewhere in the graph (i.e. in ``b``), by comparing the graph for the
     transformed log-probability with the SciPy-derived log-probability--using a
     numeric approximation to the Jacobian term.
+
+    TODO: This test is rather redundant with those in tess/distributions/test_transform.py
     """
 
     a = at_dist(*dist_params, size=size)
@@ -228,16 +229,14 @@ def test_transformed_logprob(at_dist, dist_params, sp_dist, size):
     b_value_var = b.clone()
     b_value_var.name = "b_value"
 
-    transform_rewrite = TransformValuesRewrite({a_value_var: DEFAULT_TRANSFORM})
+    transform = _default_transform(a.owner.op, a)
+    transform_rewrite = TransformValuesRewrite({a_value_var: transform})
     res = joint_logprob({a: a_value_var, b: b_value_var}, extra_rewrites=transform_rewrite)
 
     test_val_rng = np.random.RandomState(3238)
 
     logp_vals_fn = pytensor.function([a_value_var, b_value_var], res)
 
-    a_trans_op = _default_transformed_rv(a.owner.op, a.owner).op
-    transform = a_trans_op.transform
-
     a_forward_fn = pytensor.function([a_value_var], transform.forward(a_value_var, *a.owner.inputs))
     a_backward_fn = pytensor.function(
         [a_value_var], transform.backward(a_value_var, *a.owner.inputs)
@@ -305,7 +304,7 @@ def test_simple_transformed_logprob_nojac(use_jacobian):
     x_vv = X_rv.clone()
     x_vv.name = "x"
 
-    transform_rewrite = TransformValuesRewrite({x_vv: DEFAULT_TRANSFORM})
+    transform_rewrite = TransformValuesRewrite({x_vv: log})
     tr_logp = joint_logprob(
         {X_rv: x_vv}, extra_rewrites=transform_rewrite, use_jacobian=use_jacobian
     )
@@ -359,9 +358,9 @@ def test_hierarchical_uniform_transform():
 
     transform_rewrite = TransformValuesRewrite(
         {
-            lower: DEFAULT_TRANSFORM,
-            upper: DEFAULT_TRANSFORM,
-            x: DEFAULT_TRANSFORM,
+            lower: _default_transform(lower_rv.owner.op, lower_rv),
+            upper: _default_transform(upper_rv.owner.op, upper_rv),
+            x: _default_transform(x_rv.owner.op, x_rv),
         }
     )
     logp = joint_logprob(
@@ -425,28 +424,7 @@ def test_default_transform_multiout():
     x_rv = at.random.normal(0, sd, name="x")
     x = x_rv.clone()
 
-    transform_rewrite = TransformValuesRewrite({x: DEFAULT_TRANSFORM})
-
-    logp = joint_logprob(
-        {x_rv: x},
-        extra_rewrites=transform_rewrite,
-    )
-
-    assert np.isclose(
-        logp.eval({x: 1}),
-        sp.stats.norm(0, 1).logpdf(1),
-    )
-
-
-def test_nonexistent_default_transform():
-    """
-    Test that setting `DEFAULT_TRANSFORM` to a variable that has no default
-    transform does not fail
-    """
-    x_rv = at.random.normal(name="x")
-    x = x_rv.clone()
-
-    transform_rewrite = TransformValuesRewrite({x: DEFAULT_TRANSFORM})
+    transform_rewrite = TransformValuesRewrite({x: None})
 
     logp = joint_logprob(
         {x_rv: x},
@@ -480,7 +458,7 @@ def test_original_values_output_dict():
     p_rv = at.random.beta(1, 1, name="p")
     p_vv = p_rv.clone()
 
-    tr = TransformValuesRewrite({p_vv: DEFAULT_TRANSFORM})
+    tr = TransformValuesRewrite({p_vv: logodds})
     logp_dict = factorized_joint_logprob({p_rv: p_vv}, extra_rewrites=tr)
 
     assert p_vv in logp_dict