Add rewrite for 1 ** x = 1

pytensor/tensor/rewriting/math.py

@@ -1905,12 +1905,27 @@ def local_reciprocal_canon(fgraph, node):
 @register_canonicalize
 @node_rewriter([pt_pow])
 def local_pow_canonicalize(fgraph, node):
-    cst = get_underlying_scalar_constant_value(
+    """
+    Rewrites for exponential functions with straight-forward simplifications:
+    1. x ** 0 -> 1
+    2. x ** 1 -> x
+    3. 1 ** x -> 1
+
+    In all cases, the shape of the output is the result of broadcasting the shapes of the inputs.
+    """
+
+    cst_base = get_underlying_scalar_constant_value(
+        node.inputs[0], only_process_constants=True, raise_not_constant=False
+    )
+    if cst_base == 1:
+        return [alloc_like(1, node.outputs[0], fgraph)]
+
+    cst_exponent = get_underlying_scalar_constant_value(
         node.inputs[1], only_process_constants=True, raise_not_constant=False
     )
-    if cst == 0:
+    if cst_exponent == 0:
         return [alloc_like(1, node.outputs[0], fgraph)]
-    if cst == 1:
+    if cst_exponent == 1:
         return [alloc_like(node.inputs[0], node.outputs[0], fgraph)]
 
 

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)