Make non-strict zip strict in printing.py

pymc-devs · Jul 3, 2024 · 1218135 · 1218135
1 parent 921ce36
commit 1218135
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/pytensor/printing.py b/pytensor/printing.py
@@ -925,18 +925,17 @@ def process(self, output, pstate):
             )
         idx = node.outputs.index(output)
         pattern, precedences = self.patterns[idx]
-        precedences += (1000,) * len(node.inputs)
+        precedences += (1000,) * (len(node.inputs) - len(precedences))
 
         def pp_process(input, new_precedence):
             with set_precedence(pstate, new_precedence):
                 r = pprinter.process(input, pstate)
             return r
 
         d = {
-            str(i): x
-            for i, x in enumerate(
-                pp_process(input, precedence)
-                for input, precedence in zip(node.inputs, precedences, strict=False)
+            str(i): pp_process(input, precedence)
+            for i, (input, precedence) in enumerate(
+                zip(node.inputs, precedences, strict=True)
             )
         }
         r = pattern % d