Rename RVTransform to Transform

Author: ricardoV94

pymc/distributions/transforms.py

@@ -31,12 +31,12 @@
     IntervalTransform,
     LogOddsTransform,
     LogTransform,
-    RVTransform,
     SimplexTransform,
+    Transform,
 )
 
 __all__ = [
-    "RVTransform",
+    "Transform",
     "simplex",
     "logodds",
     "Interval",
@@ -60,6 +60,10 @@ def __getattr__(name):
         warnings.warn(f"{name} has been deprecated, use sum_to_1 instead.", FutureWarning)
         return sum_to_1
 
+    if name == "RVTransform":
+        warnings.warn("RVTransform has been renamed to Transform", FutureWarning)
+        return Transform
+
     raise AttributeError(f"module {__name__} has no attribute {name}")
 
 
@@ -69,7 +73,7 @@ def _default_transform(op: Op, rv: TensorVariable):
     return None
 
 
-class LogExpM1(RVTransform):
+class LogExpM1(Transform):
     name = "log_exp_m1"
 
     def backward(self, value, *inputs):
@@ -87,7 +91,7 @@ def log_jac_det(self, value, *inputs):
         return -pt.softplus(-value)
 
 
-class Ordered(RVTransform):
+class Ordered(Transform):
     name = "ordered"
 
     def __init__(self, ndim_supp=None):
@@ -110,7 +114,7 @@ def log_jac_det(self, value, *inputs):
         return pt.sum(value[..., 1:], axis=-1)
 
 
-class SumTo1(RVTransform):
+class SumTo1(Transform):
     """
     Transforms K - 1 dimensional simplex space (k values in [0,1] and that sum to 1) to a K - 1 vector of values in [0,1]
     This Transformation operates on the last dimension of the input tensor.
@@ -134,7 +138,7 @@ def log_jac_det(self, value, *inputs):
         return pt.sum(y, axis=-1)
 
 
-class CholeskyCovPacked(RVTransform):
+class CholeskyCovPacked(Transform):
     """
     Transforms the diagonal elements of the LKJCholeskyCov distribution to be on the
     log scale
@@ -162,7 +166,7 @@ def log_jac_det(self, value, *inputs):
         return pt.sum(value[..., self.diag_idxs], axis=-1)
 
 
-class Chain(RVTransform):
+class Chain(Transform):
     __slots__ = ("param_extract_fn", "transform_list", "name")
 
     def __init__(self, transform_list):
@@ -297,7 +301,7 @@ def bounds_fn(*rv_inputs):
         super().__init__(args_fn=bounds_fn)
 
 
-class ZeroSumTransform(RVTransform):
+class ZeroSumTransform(Transform):
     """
     Constrains any random samples to sum to zero along the user-provided ``zerosum_axes``.
 

pymc/initial_point.py

@@ -24,7 +24,7 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.tensor.variable import TensorVariable
 
-from pymc.logprob.transforms import RVTransform
+from pymc.logprob.transforms import Transform
 from pymc.pytensorf import compile_pymc, find_rng_nodes, replace_rng_nodes, reseed_rngs
 from pymc.util import get_transformed_name, get_untransformed_name, is_transformed_name
 
@@ -177,7 +177,7 @@ def inner(seed, *args, **kwargs):
 def make_initial_point_expression(
     *,
     free_rvs: Sequence[TensorVariable],
-    rvs_to_transforms: Dict[TensorVariable, RVTransform],
+    rvs_to_transforms: Dict[TensorVariable, Transform],
     initval_strategies: Dict[TensorVariable, Optional[Union[np.ndarray, Variable, str]]],
     jitter_rvs: Set[TensorVariable] = None,
     default_strategy: str = "moment",

pymc/logprob/basic.py

@@ -64,7 +64,7 @@
 )
 from pymc.logprob.rewriting import cleanup_ir, construct_ir_fgraph
 from pymc.logprob.transform_value import TransformValuesRewrite
-from pymc.logprob.transforms import RVTransform
+from pymc.logprob.transforms import Transform
 from pymc.logprob.utils import find_rvs_in_graph, rvs_to_value_vars
 
 TensorLike: TypeAlias = Union[Variable, float, np.ndarray]
@@ -589,7 +589,7 @@ def transformed_conditional_logp(
     rvs: Sequence[TensorVariable],
     *,
     rvs_to_values: Dict[TensorVariable, TensorVariable],
-    rvs_to_transforms: Dict[TensorVariable, RVTransform],
+    rvs_to_transforms: Dict[TensorVariable, Transform],
     jacobian: bool = True,
     **kwargs,
 ) -> List[TensorVariable]:

pymc/logprob/transform_value.py

@@ -27,7 +27,7 @@
 
 from pymc.logprob.abstract import MeasurableVariable, _logprob
 from pymc.logprob.rewriting import PreserveRVMappings, cleanup_ir_rewrites_db
-from pymc.logprob.transforms import RVTransform
+from pymc.logprob.transforms import Transform
 
 
 class TransformedValue(Op):
@@ -67,7 +67,7 @@ class TransformedValueRV(Op):
 
     __props__ = ("transforms",)
 
-    def __init__(self, transforms: Sequence[RVTransform]):
+    def __init__(self, transforms: Sequence[Transform]):
         self.transforms = tuple(transforms)
         super().__init__()
 
@@ -320,7 +320,7 @@ class TransformValuesRewrite(GraphRewriter):
 
     def __init__(
         self,
-        values_to_transforms: Dict[TensorVariable, Union[RVTransform, None]],
+        values_to_transforms: Dict[TensorVariable, Union[Transform, None]],
     ):
         """
         Parameters

pymc/logprob/transforms.py

@@ -120,7 +120,7 @@
 from pymc.logprob.utils import CheckParameterValue, check_potential_measurability
 
 
-class RVTransform(abc.ABC):
+class Transform(abc.ABC):
     ndim_supp = None
 
     @abc.abstractmethod
@@ -174,10 +174,10 @@ class MeasurableTransform(MeasurableElemwise):
 
     # Cannot use `transform` as name because it would clash with the property added by
     # the `TransformValuesRewrite`
-    transform_elemwise: RVTransform
+    transform_elemwise: Transform
     measurable_input_idx: int
 
-    def __init__(self, *args, transform: RVTransform, measurable_input_idx: int, **kwargs):
+    def __init__(self, *args, transform: Transform, measurable_input_idx: int, **kwargs):
         self.transform_elemwise = transform
         self.measurable_input_idx = measurable_input_idx
         super().__init__(*args, **kwargs)
@@ -444,7 +444,7 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[Li
     scalar_op = node.op.scalar_op
     measurable_input_idx = 0
     transform_inputs: Tuple[TensorVariable, ...] = (measurable_input,)
-    transform: RVTransform
+    transform: Transform
 
     transform_dict = {
         Exp: ExpTransform(),
@@ -559,7 +559,7 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[Li
 )
 
 
-class SinhTransform(RVTransform):
+class SinhTransform(Transform):
     name = "sinh"
     ndim_supp = 0
 
@@ -570,7 +570,7 @@ def backward(self, value, *inputs):
         return pt.arcsinh(value)
 
 
-class CoshTransform(RVTransform):
+class CoshTransform(Transform):
     name = "cosh"
     ndim_supp = 0
 
@@ -589,7 +589,7 @@ def log_jac_det(self, value, *inputs):
         )
 
 
-class TanhTransform(RVTransform):
+class TanhTransform(Transform):
     name = "tanh"
     ndim_supp = 0
 
@@ -600,7 +600,7 @@ def backward(self, value, *inputs):
         return pt.arctanh(value)
 
 
-class ArcsinhTransform(RVTransform):
+class ArcsinhTransform(Transform):
     name = "arcsinh"
     ndim_supp = 0
 
@@ -611,7 +611,7 @@ def backward(self, value, *inputs):
         return pt.sinh(value)
 
 
-class ArccoshTransform(RVTransform):
+class ArccoshTransform(Transform):
     name = "arccosh"
     ndim_supp = 0
 
@@ -622,7 +622,7 @@ def backward(self, value, *inputs):
         return pt.cosh(value)
 
 
-class ArctanhTransform(RVTransform):
+class ArctanhTransform(Transform):
     name = "arctanh"
     ndim_supp = 0
 
@@ -633,7 +633,7 @@ def backward(self, value, *inputs):
         return pt.tanh(value)
 
 
-class ErfTransform(RVTransform):
+class ErfTransform(Transform):
     name = "erf"
     ndim_supp = 0
 
@@ -644,7 +644,7 @@ def backward(self, value, *inputs):
         return pt.erfinv(value)
 
 
-class ErfcTransform(RVTransform):
+class ErfcTransform(Transform):
     name = "erfc"
     ndim_supp = 0
 
@@ -655,7 +655,7 @@ def backward(self, value, *inputs):
         return pt.erfcinv(value)
 
 
-class ErfcxTransform(RVTransform):
+class ErfcxTransform(Transform):
     name = "erfcx"
     ndim_supp = 0
 
@@ -681,7 +681,7 @@ def calc_delta_x(value, prior_result):
         return result[-1]
 
 
-class LocTransform(RVTransform):
+class LocTransform(Transform):
     name = "loc"
 
     def __init__(self, transform_args_fn):
@@ -699,7 +699,7 @@ def log_jac_det(self, value, *inputs):
         return pt.zeros_like(value)
 
 
-class ScaleTransform(RVTransform):
+class ScaleTransform(Transform):
     name = "scale"
 
     def __init__(self, transform_args_fn):
@@ -718,7 +718,7 @@ def log_jac_det(self, value, *inputs):
         return -pt.log(pt.abs(pt.broadcast_to(scale, value.shape)))
 
 
-class LogTransform(RVTransform):
+class LogTransform(Transform):
     name = "log"
 
     def forward(self, value, *inputs):
@@ -731,7 +731,7 @@ def log_jac_det(self, value, *inputs):
         return value
 
 
-class ExpTransform(RVTransform):
+class ExpTransform(Transform):
     name = "exp"
 
     def forward(self, value, *inputs):
@@ -744,7 +744,7 @@ def log_jac_det(self, value, *inputs):
         return -pt.log(value)
 
 
-class AbsTransform(RVTransform):
+class AbsTransform(Transform):
     name = "abs"
 
     def forward(self, value, *inputs):
@@ -758,7 +758,7 @@ def log_jac_det(self, value, *inputs):
         return pt.switch(value >= 0, 0, np.nan)
 
 
-class PowerTransform(RVTransform):
+class PowerTransform(Transform):
     name = "power"
 
     def __init__(self, power=None):
@@ -801,7 +801,7 @@ def log_jac_det(self, value, *inputs):
         return res
 
 
-class IntervalTransform(RVTransform):
+class IntervalTransform(Transform):
     name = "interval"
 
     def __init__(self, args_fn: Callable[..., Tuple[Optional[Variable], Optional[Variable]]]):
@@ -909,7 +909,7 @@ def log_jac_det(self, value, *inputs):
             return pt.zeros_like(value)
 
 
-class LogOddsTransform(RVTransform):
+class LogOddsTransform(Transform):
     name = "logodds"
 
     def backward(self, value, *inputs):
@@ -923,7 +923,7 @@ def log_jac_det(self, value, *inputs):
         return pt.log(sigmoid_value) + pt.log1p(-sigmoid_value)
 
 
-class SimplexTransform(RVTransform):
+class SimplexTransform(Transform):
     name = "simplex"
 
     def forward(self, value, *inputs):
@@ -950,7 +950,7 @@ def log_jac_det(self, value, *inputs):
         return pt.sum(res, -1)
 
 
-class CircularTransform(RVTransform):
+class CircularTransform(Transform):
     name = "circular"
 
     def backward(self, value, *inputs):
@@ -963,7 +963,7 @@ def log_jac_det(self, value, *inputs):
         return pt.zeros(value.shape)
 
 
-class ChainedTransform(RVTransform):
+class ChainedTransform(Transform):
     name = "chain"
 
     def __init__(self, transform_list, base_op):

pymc/model/fgraph.py

@@ -24,7 +24,7 @@
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.sharedvar import ScalarSharedVariable
 
-from pymc.logprob.transforms import RVTransform
+from pymc.logprob.transforms import Transform
 from pymc.model.core import Model
 from pymc.pytensorf import StringType, find_rng_nodes, toposort_replace
 
@@ -59,8 +59,8 @@ def perform(self, *args, **kwargs):
 class ModelValuedVar(ModelVar):
     __props__ = ("transform",)
 
-    def __init__(self, transform: Optional[RVTransform] = None):
-        if transform is not None and not isinstance(transform, RVTransform):
+    def __init__(self, transform: Optional[Transform] = None):
+        if transform is not None and not isinstance(transform, Transform):
             raise TypeError(f"transform must be None or RVTransform type, got {type(transform)}")
         self.transform = transform
         super().__init__()

pymc/model/transform/conditioning.py

@@ -23,7 +23,7 @@
 from pytensor.tensor.random.op import RandomVariable
 
 from pymc import Model
-from pymc.logprob.transforms import RVTransform
+from pymc.logprob.transforms import Transform
 from pymc.model.fgraph import (
     ModelDeterministic,
     ModelFreeRV,
@@ -263,7 +263,7 @@ def do(
 
 def change_value_transforms(
     model: Model,
-    vars_to_transforms: Mapping[ModelVariable, Union[RVTransform, None]],
+    vars_to_transforms: Mapping[ModelVariable, Union[Transform, None]],
 ) -> Model:
     """Change the value variables transforms in the model