pytensor-54: Removed yet another redundant check

Author: tamastokes

pytensor/tensor/rewriting/math.py

@@ -431,38 +431,37 @@ def local_mul_exp_to_exp_add(fgraph, node):
     This rewrite detects e^x * e^y and converts it to e^(x+y).
     Similarly, e^x / e^y becomes e^(x-y).
     """
-    if isinstance(node.op.scalar_op, (aes.Mul, aes.TrueDiv)):
-        exps = [
-            n.owner.inputs[0]
+    exps = [
+        n.owner.inputs[0]
+        for n in node.inputs
+        if n.owner
+        and hasattr(n.owner.op, "scalar_op")
+        and isinstance(n.owner.op.scalar_op, aes.Exp)
+    ]
+    # Can only do any rewrite if there are at least two exp-s
+    if len(exps) >= 2:
+        # Mul -> add; TrueDiv -> sub
+        orig_op, new_op = mul, add
+        if isinstance(node.op.scalar_op, aes.TrueDiv):
+            orig_op, new_op = true_div, sub
+        new_out = exp(new_op(*exps))
+        if new_out.dtype != node.outputs[0].dtype:
+            new_out = cast(new_out, dtype=node.outputs[0].dtype)
+        # The original Mul may have more than two factors, some of which may not be exp nodes.
+        # If so, we keep multiplying them with the new exp(sum) node.
+        # E.g.: e^x * y * e^z * w --> e^(x+z) * y * w
+        rest = [
+            n
             for n in node.inputs
-            if n.owner
-            and hasattr(n.owner.op, "scalar_op")
-            and isinstance(n.owner.op.scalar_op, aes.Exp)
+            if not n.owner
+            or not hasattr(n.owner.op, "scalar_op")
+            or not isinstance(n.owner.op.scalar_op, aes.Exp)
         ]
-        # Can only do any rewrite if there are at least two exp-s
-        if len(exps) >= 2:
-            # Mul -> add; TrueDiv -> sub
-            orig_op, new_op = mul, add
-            if isinstance(node.op.scalar_op, aes.TrueDiv):
-                orig_op, new_op = true_div, sub
-            new_out = exp(new_op(*exps))
+        if len(rest) > 0:
+            new_out = orig_op(new_out, *rest)
             if new_out.dtype != node.outputs[0].dtype:
                 new_out = cast(new_out, dtype=node.outputs[0].dtype)
-            # The original Mul may have more than two factors, some of which may not be exp nodes.
-            # If so, we keep multiplying them with the new exp(sum) node.
-            # E.g.: e^x * y * e^z * w --> e^(x+z) * y * w
-            rest = [
-                n
-                for n in node.inputs
-                if not n.owner
-                or not hasattr(n.owner.op, "scalar_op")
-                or not isinstance(n.owner.op.scalar_op, aes.Exp)
-            ]
-            if len(rest) > 0:
-                new_out = orig_op(new_out, *rest)
-                if new_out.dtype != node.outputs[0].dtype:
-                    new_out = cast(new_out, dtype=node.outputs[0].dtype)
-            return [new_out]
+        return [new_out]
 
 
 @register_specialize
@@ -472,52 +471,51 @@ def local_mul_pow_to_pow_add(fgraph, node):
     This rewrite detects a^x * a^y and converts it to a^(x+y).
     Similarly, a^x / a^y becomes a^(x-y).
     """
-    if isinstance(node.op.scalar_op, (aes.Mul, aes.TrueDiv)):
-        # search for pow-s and group them by their bases
-        pow_nodes = defaultdict(list)
-        rest = []
-        for n in node.inputs:
-            if (
-                n.owner
-                and hasattr(n.owner.op, "scalar_op")
-                and isinstance(n.owner.op.scalar_op, aes.Pow)
-            ):
-                base_node = n.owner.inputs[0]
-                # exponent is at n.owner.inputs[1], but we need to store the full node
-                # in case this particular power node remains alone and can't be rewritten
-                pow_nodes[base_node].append(n)
-            else:
-                rest.append(n)
-
-        # Can only do any rewrite if there are at least two pow-s with the same base
-        can_rewrite = [k for k, v in pow_nodes.items() if len(v) >= 2]
-        if len(can_rewrite) >= 1:
-            # Mul -> add; TrueDiv -> sub
-            orig_op, new_op = mul, add
-            if isinstance(node.op.scalar_op, aes.TrueDiv):
-                orig_op, new_op = true_div, sub
-            pow_factors = []
-            # Rewrite pow-s having the same base for each different base
-            # E.g.: a^x * a^y --> a^(x+y)
-            for base in can_rewrite:
-                exponents = [n.owner.inputs[1] for n in pow_nodes[base]]
-                new_node = base ** new_op(*exponents)
-                if new_node.dtype != node.outputs[0].dtype:
-                    new_node = cast(new_node, dtype=node.outputs[0].dtype)
-                pow_factors.append(new_node)
-            # Don't forget about those sole pow-s that couldn't be rewriten
-            sole_pows = [v[0] for k, v in pow_nodes.items() if k not in can_rewrite]
-            # Combine the rewritten pow-s and other, non-pow factors of the original Mul
-            # E.g.: a^x * y * b^z * a^w * v * b^t --> a^(x+z) * b^(z+t) * y * v
-            if len(pow_factors) > 1 or len(sole_pows) > 0 or len(rest) > 0:
-                new_out = orig_op(*pow_factors, *sole_pows, *rest)
-                if new_out.dtype != node.outputs[0].dtype:
-                    new_out = cast(new_out, dtype=node.outputs[0].dtype)
-            else:
-                # if all factors of the original mul were pows-s with the same base,
-                # we can get rid of the mul completely.
-                new_out = pow_factors[0]
-            return [new_out]
+    # search for pow-s and group them by their bases
+    pow_nodes = defaultdict(list)
+    rest = []
+    for n in node.inputs:
+        if (
+            n.owner
+            and hasattr(n.owner.op, "scalar_op")
+            and isinstance(n.owner.op.scalar_op, aes.Pow)
+        ):
+            base_node = n.owner.inputs[0]
+            # exponent is at n.owner.inputs[1], but we need to store the full node
+            # in case this particular power node remains alone and can't be rewritten
+            pow_nodes[base_node].append(n)
+        else:
+            rest.append(n)
+
+    # Can only do any rewrite if there are at least two pow-s with the same base
+    can_rewrite = [k for k, v in pow_nodes.items() if len(v) >= 2]
+    if len(can_rewrite) >= 1:
+        # Mul -> add; TrueDiv -> sub
+        orig_op, new_op = mul, add
+        if isinstance(node.op.scalar_op, aes.TrueDiv):
+            orig_op, new_op = true_div, sub
+        pow_factors = []
+        # Rewrite pow-s having the same base for each different base
+        # E.g.: a^x * a^y --> a^(x+y)
+        for base in can_rewrite:
+            exponents = [n.owner.inputs[1] for n in pow_nodes[base]]
+            new_node = base ** new_op(*exponents)
+            if new_node.dtype != node.outputs[0].dtype:
+                new_node = cast(new_node, dtype=node.outputs[0].dtype)
+            pow_factors.append(new_node)
+        # Don't forget about those sole pow-s that couldn't be rewriten
+        sole_pows = [v[0] for k, v in pow_nodes.items() if k not in can_rewrite]
+        # Combine the rewritten pow-s and other, non-pow factors of the original Mul
+        # E.g.: a^x * y * b^z * a^w * v * b^t --> a^(x+z) * b^(z+t) * y * v
+        if len(pow_factors) > 1 or len(sole_pows) > 0 or len(rest) > 0:
+            new_out = orig_op(*pow_factors, *sole_pows, *rest)
+            if new_out.dtype != node.outputs[0].dtype:
+                new_out = cast(new_out, dtype=node.outputs[0].dtype)
+        else:
+            # if all factors of the original mul were pows-s with the same base,
+            # we can get rid of the mul completely.
+            new_out = pow_factors[0]
+        return [new_out]
 
 
 @register_stabilize