pytensor-54: Rewrite a^x * a^y to a^(x+y)

Author: tamastokes

pytensor/tensor/rewriting/math.py

@@ -466,6 +466,65 @@ def local_mulexp2expadd(fgraph, node):
             return [new_out]
 
 
+@register_specialize
+@node_rewriter([mul, true_div])
+def local_mulpow2powadd(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, Elemwise) and isinstance(
+        node.op.scalar_op, (aes.Mul, aes.TrueDiv)
+    ):
+        from collections import defaultdict
+
+        # 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
 @register_specialize
 @register_canonicalize

tests/tensor/rewriting/test_math.py

@@ -4105,6 +4105,64 @@ def test_local_mulexp2expadd():
     assert isinstance(graph[0].inputs[0], TensorConstant)
 
 
+def test_local_mulpow2powadd():
+    x = scalar("x")
+    y = scalar("y")
+    z = scalar("z")
+    w = scalar("w")
+    v = scalar("v")
+    u = scalar("u")
+    t = scalar("t")
+    s = scalar("s")
+    a = scalar("a")
+    b = scalar("b")
+    c = scalar("c")
+
+    # 2^x * 2^y * 2^z * 2^w = 2^(x+y+z+w)
+    op = 2**x * 2**y * 2**z * 2**w
+    f = function([x, y, z, w], op)
+    utt.assert_allclose(f(3, 4, 5, 6), 2 ** (3 + 4 + 5 + 6))
+    graph = f.maker.fgraph.toposort()
+    assert isinstance(graph[0].op, Elemwise)
+    inner_graph = graph[0].op.scalar_op.fgraph.toposort()
+    assert any(isinstance(n.op, aes.Add) for n in inner_graph)
+    assert not any(isinstance(n.op, aes.Mul) for n in inner_graph)
+
+    # 2^x * a^y * 2^z * b^w * c^v * a^u * s * b^t = 2^(x+z) * a^(y+u) * b^(w+t) * c^v * s
+    op = 2**x * a**y * 2**z * b**w * c**v * a**u * s * b**t
+    f = function([x, y, z, w, v, u, t, s, a, b, c], op)
+    utt.assert_allclose(
+        f(4, 5, 6, 7, 8, 9, 10, 11, 2.5, 3, 3.5),
+        2 ** (4 + 6) * 2.5 ** (5 + 9) * 3 ** (7 + 10) * 3.5**8 * 11,
+    )
+    graph = f.maker.fgraph.toposort()
+    assert isinstance(graph[0].op, Elemwise)
+    inner_graph = graph[0].op.scalar_op.fgraph.toposort()
+    assert len([True for n in inner_graph if isinstance(n.op, aes.Add)]) == 3
+    assert len([True for n in inner_graph if isinstance(n.op, aes.Pow)]) == 4
+    assert any(isinstance(n.op, aes.Mul) for n in inner_graph)
+
+    # (2^x / 2^y) * (a^z / a^w) = 2^(x-y) * a^(z-w)
+    op = 2**x / 2**y * (a**z / a**w)
+    f = function([x, y, z, w, a], op)
+    utt.assert_allclose(f(3, 5, 6, 4, 7), 2 ** (3 - 5) * 7 ** (6 - 4))
+    graph = f.maker.fgraph.toposort()
+    assert isinstance(graph[0].op, Elemwise)
+    inner_graph = graph[0].op.scalar_op.fgraph.toposort()
+    assert len([True for n in inner_graph if isinstance(n.op, aes.Sub)]) == 2
+    assert any(isinstance(n.op, aes.Mul) for n in inner_graph)
+
+    # a^x * a^y * exp(z) * exp(w) = a^(x+y) * exp(z+w)
+    op = a**x * a**y * exp(z) * exp(w)
+    f = function([x, y, z, w, a], op)
+    utt.assert_allclose(f(3, 4, 5, 6, 2), 2 ** (3 + 4) * np.exp(5 + 6))
+    graph = f.maker.fgraph.toposort()
+    assert isinstance(graph[0].op, Elemwise)
+    inner_graph = graph[0].op.scalar_op.fgraph.toposort()
+    assert len([True for n in inner_graph if isinstance(n.op, aes.Add)]) == 2
+    assert any(isinstance(n.op, aes.Mul) for n in inner_graph)
+
+
 def test_local_expm1():
     x = matrix("x")
     u = scalar("u")