Rewrite local_merge_alloc to remove a non-strict zip

pytensor/tensor/rewriting/basic.py

@@ -1207,25 +1207,23 @@ def local_merge_alloc(fgraph, node):
     inputs_inner = node.inputs[0].owner.inputs
     dims_outer = inputs_outer[1:]
     dims_inner = inputs_inner[1:]
-    dims_outer_rev = dims_outer[::-1]
-    dims_inner_rev = dims_inner[::-1]
+    assert len(dims_inner) <= len(dims_outer)
     # check if the pattern of broadcasting is matched, in the reversed ordering.
     # The reverse ordering is needed when an Alloc add an implicit new
     # broadcasted dimensions to its inputs[0]. Eg:
     # Alloc(Alloc(m, y, 1, 1), x, y, z, w) -> Alloc(m, x, y, z, w)
-    i = 0
-    for dim_inner, dim_outer in zip(dims_inner_rev, dims_outer_rev, strict=False):
-        if dim_inner != dim_outer:
+    for i, dim_inner in enumerate(reversed(dims_inner)):
+        dim_outer = dims_outer[-1 - i]
+        if dim_inner == dim_outer:
+            continue
         if isinstance(dim_inner, Constant) and dim_inner.data == 1:
-                pass
-            else:
+            continue
         dims_outer[-1 - i] = Assert(
             "You have a shape error in your graph. To see a better"
             " error message and a stack trace of where in your code"
             " the error is created, use the PyTensor flags"
             " optimizer=None or optimizer=fast_compile."
         )(dim_outer, eq(dim_outer, dim_inner))
-        i += 1
     return [alloc(inputs_inner[0], *dims_outer)]