Pass rv to _default_transform dispatch function

Author: ricardoV94

pymc/distributions/continuous.py

@@ -141,17 +141,17 @@ class CircularContinuous(Continuous):
 
 
 @_default_transform.register(PositiveContinuous)
-def pos_cont_transform(op):
+def pos_cont_transform(op, rv):
     return transforms.log
 
 
 @_default_transform.register(UnitContinuous)
-def unit_cont_transform(op):
+def unit_cont_transform(op, rv):
     return transforms.logodds
 
 
 @_default_transform.register(CircularContinuous)
-def circ_cont_transform(op):
+def circ_cont_transform(op, rv):
     return transforms.circular
 
 

pymc/distributions/transforms.py

@@ -25,6 +25,7 @@
     Simplex,
 )
 from aesara.graph import Op
+from aesara.tensor import TensorVariable
 
 __all__ = [
     "RVTransform",
@@ -42,7 +43,7 @@
 
 
 @singledispatch
-def _default_transform(op: Op):
+def _default_transform(op: Op, rv: TensorVariable):
     """Return default transform for a given Distribution `Op`"""
     return None
 

pymc/model.py

@@ -1422,7 +1422,7 @@ def create_value_var(
         # Make the value variable a transformed value variable,
         # if there's an applicable transform
         if transform is UNSET and rv_var.owner:
-            transform = _default_transform(rv_var.owner.op)
+            transform = _default_transform(rv_var.owner.op, rv_var)
 
         if transform is not None and transform is not UNSET:
             value_var.tag.transform = transform