Add rewrite for 1 ** x = 1

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -9,7 +9,8 @@
 
 import pytensor.scalar.basic as ps
 import pytensor.scalar.math as ps_math
-from pytensor.graph.basic import Constant, Variable
+from pytensor.graph import FunctionGraph
+from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.rewriting.basic import (
     NodeRewriter,
     PatternNodeRewriter,
@@ -1914,6 +1915,33 @@ def local_pow_canonicalize(fgraph, node):
         return [alloc_like(node.inputs[0], node.outputs[0], fgraph)]
 
 
+@register_canonicalize
+@node_rewriter([pt_pow])
+def local_pow_canonicalize_base_1(
+    fgraph: FunctionGraph, node: Apply
+) -> list[TensorVariable] | None:
+    """
+    Replace `1 ** x` with 1, broadcast to the shape of the output.
+
+    Parameters
+    ----------
+    fgraph: FunctionGraph
+        Full function graph being rewritten
+    node: Apply
+        Specific node being rewritten
+
+    Returns
+    -------
+    rewritten_output: list[TensorVariable] | None
+        Rewritten output of node, or None if no rewrite is possible
+    """
+    cst = get_underlying_scalar_constant_value(
+        node.inputs[0], only_process_constants=True, raise_not_constant=False
+    )
+    if cst == 1:
+        return [alloc_like(node.inputs[0], node.outputs[0], fgraph)]
+
+
 @register_specialize
 @node_rewriter([mul])
 def local_mul_to_sqr(fgraph, node):

tests/tensor/rewriting/test_math.py

@@ -4571,3 +4571,22 @@ def test_log_kv_stabilization():
         out.eval({x: 1000.0}, mode=mode),
         -1003.2180912984705,
     )
+
+
+@pytest.mark.parametrize("shape", [(), (4, 5, 6)], ids=["scalar", "tensor"])
+def test_pow_1_rewrite(shape):
+    x = pt.tensor("x", shape=shape)
+    z = 1**x
+
+    f1 = pytensor.function([x], z, mode=get_default_mode().excluding("canonicalize"))
+    assert debugprint(f1, file="str").count("Pow") == 1
+
+    x_val = np.random.random(shape).astype(config.floatX)
+    z_val_1 = f1(x_val)
+
+    f2 = pytensor.function([x], z)
+    assert debugprint(f2, file="str").count("Pow") == 0
+
+    z_val_2 = f2(x_val)
+
+    np.testing.assert_allclose(z_val_1, z_val_2)