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,55 @@ def local_remove_check_parameter(fgraph, node):
         return [node.inputs[0]]
 
 
+@node_rewriter(tracks=[pt.switch])
+def local_remove_useless_bound_switch(fgraph, node):
+    """Remove bound checks ensured by the transformations.
+
+    switch(exp(x) >= 0, cond1, -inf) -> 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)
+        and (isinstance(false_branch, Constant) and false_branch.unique_value == -np.inf)
+    ):
+        return None
+
+    # Check if var is exp(x) 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 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 +297,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_switch.__name__,
+    local_remove_useless_bound_switch,
+    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_switch",
+        )
 
     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,

tests/test_pytensorf.py

@@ -354,6 +354,13 @@ def test_check_parameters_can_be_replaced_by_ninf(self):
         with pytest.raises(ParameterValueError, match="test"):
             fn([-1, 2, 3])
 
+    def test_useless_bound_switch(self):
+        # Without default transform the switch is never removed
+        # Even if check_bounds = False
+        with pm.Model(check_bounds=False) as m:
+            x = pm.HalfNormal("x", default_transform=None)
+        assert m.compile_logp()({"x": -1}) == -np.inf
+
     def test_compile_pymc_sets_rng_updates(self):
         rng = pytensor.shared(np.random.default_rng(0))
         x = pm.Normal.dist(rng=rng)