Make non-strict zips strict in tensor/elemwise_cgen

Author: Armavica

pytensor/tensor/elemwise_cgen.py

@@ -208,7 +208,13 @@ def make_alloc(loop_orders, dtype, sub, fortran="0"):
     """
 
 
-def make_loop(loop_orders, dtypes, loop_tasks, sub, openmp=None):
+def make_loop(
+    loop_orders: list[tuple[int | str, ...]],
+    dtypes: list,
+    loop_tasks: list,
+    sub: dict[str, str],
+    openmp: bool = False,
+):
     """
     Make a nested loop over several arrays and associate specific code
     to each level of nesting.
@@ -226,7 +232,7 @@ def make_loop(loop_orders, dtypes, loop_tasks, sub, openmp=None):
         string is code to be executed before the ith loop starts, the second
         one contains code to be executed just before going to the next element
         of the ith dimension.
-        The last element if loop_tasks is a single string, containing code
+        The last element of loop_tasks is a single string, containing code
         to be executed at the very end.
     sub : dictionary
         Maps 'lv#' to a suitable variable name.
@@ -259,7 +265,7 @@ def loop_over(preloop, code, indices, i):
         }}
         """
 
-    preloops = {}
+    preloops: dict[int, str] = {}
     for i, (loop_order, dtype) in enumerate(zip(loop_orders, dtypes, strict=True)):
         for j, index in enumerate(loop_order):
             if index != "x":
@@ -276,16 +282,8 @@ def loop_over(preloop, code, indices, i):
 
     s = ""
 
-    for i, (pre_task, task), indices in reversed(
-        list(
-            zip(
-                range(len(loop_tasks) - 1),
-                loop_tasks,
-                list(zip(*loop_orders, strict=True)),
-                strict=False,
-            )
-        )
-    ):
+    tasks_indices = zip(loop_tasks[:-1], zip(*loop_orders, strict=True), strict=True)
+    for i, ((pre_task, task), indices) in reversed(list(enumerate(tasks_indices))):
         s = loop_over(preloops.get(i, "") + pre_task, s + task, indices, i)
 
     s += loop_tasks[-1]
@@ -549,16 +547,10 @@ def loop_over(preloop, code, indices, i):
         s = preloops.get(0, "")
     else:
         s = ""
-        for i, (pre_task, task), indices in reversed(
-            list(
-                zip(
-                    range(len(loop_tasks) - 1),
-                    loop_tasks,
-                    list(zip(*loop_orders, strict=True)),
-                    strict=False,
-                )
-            )
-        ):
+        tasks_indices = zip(
+            loop_tasks[:-1], zip(*loop_orders, strict=True), strict=True
+        )
+        for i, ((pre_task, task), indices) in reversed(list(enumerate(tasks_indices))):
             s = loop_over(preloops.get(i, "") + pre_task, s + task, indices, i)
 
     s += loop_tasks[-1]