Fail graciously in local_pow_to_nested_squaring when static type shape is updated

pytensor/tensor/rewriting/math.py

@@ -2081,7 +2081,6 @@ def local_pow_to_nested_squaring(fgraph, node):
     Note: This sounds like the kind of thing any half-decent compiler can do by itself?
     """
 
-    if node.op == at_pow:
     # the idea here is that we have pow(x, y)
     odtype = node.outputs[0].dtype
     xsym = node.inputs[0]
@@ -2127,16 +2126,19 @@ def local_pow_to_nested_squaring(fgraph, node):
             if abs(y) > 2:
                 # We fuse all the pow together here to make
                 # compilation faster
-                    rval1 = Elemwise(
-                        aes.Composite([pow2_scal[0]], [rval1_scal])
-                    ).make_node(xsym)
+                rval1 = Elemwise(aes.Composite([pow2_scal[0]], [rval1_scal])).make_node(
+                    xsym
+                )
             if y < 0:
                 rval = [reciprocal(rval1)]
             else:
                 rval = [rval1]
         if rval:
             rval[0] = cast(rval[0], odtype)
-                assert rval[0].type == node.outputs[0].type, (rval, node.outputs)
+            # TODO: We can add a specify_broadcastable and/or unbroadcast to make the
+            #  output types compatible. Or work on #408 and let TensorType.filter_variable do it.
+            if rval[0].type.broadcastable != node.outputs[0].type.broadcastable:
+                return None
             return rval
 
 

tests/tensor/rewriting/test_math.py

@@ -29,8 +29,9 @@ from pytensor.graph.rewriting.db import RewriteDatabaseQuery
 from pytensor.graph.rewriting.utils import is_same_graph, rewrite_graph
 from pytensor.misc.safe_asarray import _asarray
 from pytensor.printing import debugprint
+from pytensor.scalar import Pow
 from pytensor.tensor import inplace
-from pytensor.tensor.basic import Alloc, join, second, switch
+from pytensor.tensor.basic import Alloc, constant, join, second, switch
 from pytensor.tensor.blas import Dot22, Gemv
 from pytensor.tensor.blas_c import CGemv
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
@@ -69,7 +70,7 @@ from pytensor.tensor.math import max as at_max
 from pytensor.tensor.math import maximum
 from pytensor.tensor.math import min as at_min
 from pytensor.tensor.math import minimum, mul, neg, neq
-from pytensor.tensor.math import pow as at_pow
+from pytensor.tensor.math import pow as pt_pow
 from pytensor.tensor.math import (
     prod,
     rad2deg,
@@ -1746,6 +1747,29 @@ def test_local_pow_to_nested_squaring():
     utt.assert_allclose(f(val_no0), val_no0 ** (-16))
 
 
+def test_local_pow_to_nested_squaring_fails_gracefully():
+    # Reported in #456
+
+    x = vector("x", shape=(1,))
+    # Create an Apply that does not have precise output shape
+    node = Apply(
+        op=pt_pow,
+        inputs=[x, constant([2.0])],
+        outputs=[tensor(shape=(None,))],
+    )
+    y = node.default_output()
+
+    fn = function([x], y)
+
+    # Check rewrite is not applied (this could change in the future)
+    assert any(
+        (isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, Pow))
+        for node in fn.maker.fgraph.apply_nodes
+    )
+
+    np.testing.assert_allclose(fn([2.0]), np.array([4.0]))
+
+
 class TestFuncInverse:
     def setup_method(self):
         mode = get_default_mode()
@@ -2449,7 +2473,7 @@ class TestLocalMergeSwitchSameCond:
             le,
             eq,
             neq,
-            at_pow,
+            pt_pow,
         ):
             g = rewrite(FunctionGraph(mats, [op(s1, s2)]))
             assert debugprint(g, file="str").count("Switch") == 1