Remove useless switch on log transformed parameters

Author: ricardoV94

pymc/logprob/utils.py

@@ -43,14 +43,14 @@
 
 from pytensor import tensor as pt
 from pytensor.graph import Apply, Op, node_rewriter
-from pytensor.graph.basic import Constant, Variable, clone_get_equiv, graph_inputs, walk
+from pytensor.graph.basic import Constant, Variable, ancestors, clone_get_equiv, graph_inputs, walk
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import HasInnerGraph
 from pytensor.link.c.type import CType
 from pytensor.raise_op import CheckAndRaise
-from pytensor.scalar.basic import Mul
+from pytensor.scalar.basic import GE, LE, Exp, Mul
 from pytensor.tensor.basic import get_underlying_scalar_constant_value
-from pytensor.tensor.elemwise import Elemwise
+from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.variable import TensorVariable
@@ -228,6 +228,52 @@ def local_remove_check_parameter(fgraph, node):
         return [node.inputs[0]]
 
 
+@node_rewriter(tracks=[pt.switch])
+def local_remove_useless_bound_checks(fgraph, node):
+    """Remove bound checks ensured by the transformations.
+
+    switch(exp(x) >= 0, cond1, cond2) -> cond1 if exp(x) in cond1
+    The reason we don't set it to simply True is that x could be `nan`.
+    If we see exp(x) exists in cond1 we assume `nan` will be propagated anyway.
+
+    This isn't guaranteed to be True, for instance if exp(x) is inside another switch statement.
+    """
+    cond, true_branch, false_branch = node.inputs
+    if not (cond.owner is not None and isinstance(cond.owner.op, Elemwise)):
+        return
+    scalar_op = cond.owner.op.scalar_op
+    if isinstance(scalar_op, LE):
+        maybe_zero, var = cond.owner.inputs
+    elif isinstance(scalar_op, GE):
+        var, maybe_zero = cond.owner.inputs
+    else:
+        return None
+
+    if not (isinstance(maybe_zero, Constant) and maybe_zero.unique_value == 0):
+        return None
+
+    # Check if var is exp(x), x is a root variable and x is present in the true branch
+    if (
+        var.owner is not None
+        and isinstance(var.owner.op, Elemwise)
+        and isinstance(var.owner.op.scalar_op, Exp)
+    ) or (
+        (
+            isinstance(var.owner.op, DimShuffle)
+            and (
+                var.owner.inputs[0].owner is not None
+                and isinstance(var.owner.inputs[0].owner.op, Elemwise)
+                and isinstance(var.owner.inputs[0].owner.op.scalar_op, Exp)
+            )
+        )
+        # and var.owner.inputs[0].owner is None
+        and var in ancestors([true_branch])
+    ):
+        return [true_branch]
+
+    return None
+
+
 @node_rewriter(tracks=[CheckParameterValue])
 def local_check_parameter_to_ninf_switch(fgraph, node):
     if not node.op.can_be_replaced_by_ninf:
@@ -248,17 +294,23 @@ def local_check_parameter_to_ninf_switch(fgraph, node):
 
 
 pytensor.compile.optdb["canonicalize"].register(
-    "local_remove_check_parameter",
+    local_remove_check_parameter.__name__,
     local_remove_check_parameter,
     use_db_name_as_tag=False,
 )
 
 pytensor.compile.optdb["canonicalize"].register(
-    "local_check_parameter_to_ninf_switch",
+    local_check_parameter_to_ninf_switch.__name__,
     local_check_parameter_to_ninf_switch,
     use_db_name_as_tag=False,
 )
 
+pytensor.compile.optdb["canonicalize"].register(
+    local_remove_useless_bound_checks.__name__,
+    local_remove_useless_bound_checks,
+    use_db_name_as_tag=False,
+)
+
 
 class DiracDelta(MeasurableOp, Op):
     """An `Op` that represents a Dirac-delta distribution."""

pymc/pytensorf.py

@@ -937,12 +937,16 @@ def compile(
         check_bounds = model.check_bounds
     except TypeError:
         check_bounds = True
-    check_parameter_opt = (
-        "local_check_parameter_to_ninf_switch" if check_bounds else "local_remove_check_parameter"
-    )
+    if check_bounds:
+        check_parameter_opt = ("local_check_parameter_to_ninf_switch",)
+    else:
+        check_parameter_opt = (
+            "local_remove_check_parameter",
+            "local_remove_useless_bound_checks",
+        )
 
     mode = get_mode(mode)
-    opt_qry = mode.provided_optimizer.including("random_make_inplace", check_parameter_opt)
+    opt_qry = mode.provided_optimizer.including("random_make_inplace", *check_parameter_opt)
     mode = Mode(linker=mode.linker, optimizer=opt_qry)
     pytensor_function = pytensor.function(
         inputs,