Rewrite for/append as list comprehensions

Author: Armavica

pytensor/compile/debugmode.py

@@ -906,11 +906,10 @@ def _get_preallocated_maps(
                 name = f"strided{tuple(steps)}"
                 for r in considered_outputs:
                     if r in init_strided:
-                        strides = []
-                        shapes = []
-                        for i, size in enumerate(r_vals[r].shape):
-                            shapes.append(slice(None, size, None))
-                            strides.append(slice(None, None, steps[i]))
+                        shapes = [slice(None, size, None) for size in r_vals[r].shape]
+                        strides = [
+                            slice(None, None, steps[i]) for i in range(r_vals[r].ndim)
+                        ]
 
                         r_buf = init_strided[r]
 

pytensor/compile/function/__init__.py

@@ -247,18 +247,10 @@ def opt_log1p(node):
 
     """
     if isinstance(outputs, dict):
-        output_items = list(outputs.items())
+        assert all(isinstance(k, str) for k in outputs)
 
-        for item_pair in output_items:
-            assert isinstance(item_pair[0], str)
-
-        output_items_sorted = sorted(output_items)
-
-        output_keys = []
-        outputs = []
-        for pair in output_items_sorted:
-            output_keys.append(pair[0])
-            outputs.append(pair[1])
+        output_keys = sorted(outputs)
+        outputs = [outputs[key] for key in output_keys]
 
     else:
         output_keys = None

pytensor/compile/function/types.py

@@ -212,18 +212,14 @@ def std_fgraph(
 
         found_updates.extend(map(SymbolicOutput, updates))
     elif fgraph is None:
-        input_vars = []
-
         # If one of the inputs is non-atomic (i.e. has a non-`None` `Variable.owner`),
         # then we need to create/clone the graph starting at these inputs.
         # The result will be atomic versions of the given inputs connected to
         # the same outputs.
         # Otherwise, when all the inputs are already atomic, there's no need to
         # clone the graph.
-        clone = force_clone
-        for spec in input_specs:
-            input_vars.append(spec.variable)
-            clone |= spec.variable.owner is not None
+        input_vars = [spec.variable for spec in input_specs]
+        clone = force_clone or any(var.owner is not None for var in input_vars)
 
         fgraph = FunctionGraph(
             input_vars,

pytensor/compile/profiling.py

@@ -1204,8 +1204,7 @@ def min_memory_generator(executable_nodes, viewed_by, view_of):
                         compute_map[var][0] = 0
 
                     for k_remove, v_remove in viewedby_remove.items():
-                        for i in v_remove:
-                            viewed_by[k_remove].append(i)
+                        viewed_by[k_remove].extend(v_remove)
 
                     for k_add, v_add in viewedby_add.items():
                         for i in v_add:
@@ -1215,15 +1214,16 @@ def min_memory_generator(executable_nodes, viewed_by, view_of):
                         del view_of[k]
 
             # two data structure used to mimic Python gc
-            viewed_by = {}  # {var1: [vars that view var1]}
+            # * {var1: [vars that view var1]}
             # The len of the list is the value of python ref
             # count. But we use a list, not just the ref count value.
-            # This is more safe to help detect potential bug  in the algo
-            for var in fgraph.variables:
-                viewed_by[var] = []
-            view_of = {}  # {var1: original var viewed by var1}
+            # This is more safe to help detect potential bug in the algo
+            viewed_by = {var: [] for var in fgraph.variables}
+
+            # * {var1: original var viewed by var1}
             # The original mean that we don't keep track of all the intermediate
             # relationship in the view.
+            view_of = {}
 
             min_memory_generator(executable_nodes, viewed_by, view_of)
 

pytensor/graph/basic.py

@@ -1474,9 +1474,8 @@ def _compute_deps_cache_(io):
 
     _clients: dict[T, list[T]] = {}
     sources: deque[T] = deque()
-    search_res_len: int = 0
+    search_res_len = len(search_res)
     for snode, children in search_res:
-        search_res_len += 1
         if children:
             for child in children:
                 _clients.setdefault(child, []).append(snode)

pytensor/graph/fg.py

@@ -270,8 +270,10 @@ def remove_client(
 
                     self.execute_callbacks("on_prune", apply_node, reason)
 
-                    for i, in_var in enumerate(apply_node.inputs):
-                        removal_stack.append((in_var, (apply_node, i)))
+                    removal_stack.extend(
+                        (in_var, (apply_node, i))
+                        for i, in_var in enumerate(apply_node.inputs)
+                    )
 
                     if remove_if_empty:
                         del clients[var]

pytensor/graph/rewriting/basic.py

@@ -479,9 +479,9 @@ def merge_profile(prof1, prof2):
             new_sub_profile.append(p[6][idx])
 
         new_rewrite = SequentialGraphRewriter(*new_l)
-        new_nb_nodes = []
-        for p1, p2 in zip(prof1[8], prof2[8]):
-            new_nb_nodes.append((p1[0] + p2[0], p1[1] + p2[1]))
+        new_nb_nodes = [
+            (p1[0] + p2[0], p1[1] + p2[1]) for p1, p2 in zip(prof1[8], prof2[8])
+        ]
         new_nb_nodes.extend(prof1[8][len(new_nb_nodes) :])
         new_nb_nodes.extend(prof2[8][len(new_nb_nodes) :])
 
@@ -960,9 +960,9 @@ def register(self, rewriter: NodeRewriter, tag_list: IterableType[str]):
 
         tracks = rewriter.tracks()
         if tracks:
+            self._tracks.extend(tracks)
             for c in tracks:
                 self.track_dict[c].append(rewriter)
-                self._tracks.append(c)
 
         for tag in tag_list:
             self.tag_dict[tag].append(rewriter)

pytensor/link/basic.py

@@ -524,12 +524,13 @@ def make_thunk(self, **kwargs):
         thunk_groups = list(zip(*thunk_lists))
         order = [x[0] for x in zip(*order_lists)]
 
-        to_reset = []
-        for thunks, node in zip(thunk_groups, order):
-            for j, output in enumerate(node.outputs):
-                if output in no_recycling:
-                    for thunk in thunks:
-                        to_reset.append(thunk.outputs[j])
+        to_reset = [
+            thunk.outputs[j]
+            for thunks, node in zip(thunk_groups, order)
+            for j, output in enumerate(node.outputs)
+            if output in no_recycling
+            for thunk in thunks
+        ]
 
         wrapper = self.wrapper
         pre = self.pre
@@ -692,18 +693,16 @@ def make_all(self, input_storage=None, output_storage=None, storage_map=None):
         computed, last_user = gc_helper(nodes)
 
         if self.allow_gc:
-            post_thunk_old_storage = []
-
-            for node in nodes:
-                post_thunk_old_storage.append(
-                    [
-                        storage_map[input]
-                        for input in node.inputs
-                        if (input in computed)
-                        and (input not in fgraph.outputs)
-                        and (node == last_user[input])
-                    ]
-                )
+            post_thunk_old_storage = [
+                [
+                    storage_map[input]
+                    for input in node.inputs
+                    if (input in computed)
+                    and (input not in fgraph.outputs)
+                    and (node == last_user[input])
+                ]
+                for node in nodes
+            ]
         else:
             post_thunk_old_storage = None
 

pytensor/link/c/basic.py

@@ -1129,19 +1129,18 @@ def __compile__(
         )
 
     def get_init_tasks(self):
-        init_tasks = []
-        tasks = []
+        vars = [v for v in self.variables if v not in self.consts]
         id = 1
-        for v in self.variables:
-            if v in self.consts:
-                continue
-            init_tasks.append((v, "init", id))
-            tasks.append((v, "get", id + 1))
-            id += 2
-        for node in self.node_order:
-            tasks.append((node, "code", id))
-            init_tasks.append((node, "init", id + 1))
-            id += 2
+        init_tasks = [(v, "init", id + 2 * i) for i, v in enumerate(vars)]
+        tasks = [(v, "get", id + 2 * i + 1) for i, v in enumerate(vars)]
+
+        id += 2 * len(vars)
+        tasks.extend(
+            (node, "code", id + 2 * i) for i, node in enumerate(self.node_order)
+        )
+        init_tasks.extend(
+            (node, "init", id + 2 * i + 1) for i, node in enumerate(self.node_order)
+        )
         return init_tasks, tasks
 
     def make_thunk(
@@ -1492,12 +1491,11 @@ def in_sig(i, topological_pos, i_idx):
         # graph's information used to compute the key. If we mistakenly
         # pretend that inputs with clients don't have any, were are only using
         # those inputs more than once to compute the key.
-        for ipos, var in [
-            (i, var)
-            for i, var in enumerate(fgraph.inputs)
+        sig.extend(
+            (var.type, in_sig(var, -1, ipos))
+            for ipos, var in enumerate(fgraph.inputs)
             if not len(fgraph.clients[var])
-        ]:
-            sig.append((var.type, in_sig(var, -1, ipos)))
+        )
 
         # crystalize the signature and version
         sig = tuple(sig)

pytensor/link/c/op.py

@@ -220,12 +220,7 @@ def prepare_node(self, node, storage_map, compute_map, impl):
 
 def lquote_macro(txt: str) -> str:
     """Turn the last line of text into a ``\\``-commented line."""
-    res = []
-    spl = txt.split("\n")
-    for l in spl[:-1]:
-        res.append(l + " \\")
-    res.append(spl[-1])
-    return "\n".join(res)
+    return " \\\n".join(txt.split("\n"))
 
 
 def get_sub_macros(sub: dict[str, str]) -> tuple[str, str]:
@@ -240,21 +235,17 @@ def get_sub_macros(sub: dict[str, str]) -> tuple[str, str]:
     return "\n".join(define_macros), "\n".join(undef_macros)
 
 
-def get_io_macros(
-    inputs: list[str], outputs: list[str]
-) -> tuple[list[str]] | tuple[str, str]:
-    define_macros = []
-    undef_macros = []
+def get_io_macros(inputs: list[str], outputs: list[str]) -> tuple[str, str]:
+    define_inputs = [f"#define INPUT_{int(i)} {inp}" for i, inp in enumerate(inputs)]
+    define_outputs = [f"#define OUTPUT_{int(i)} {out}" for i, out in enumerate(outputs)]
 
-    for i, inp in enumerate(inputs):
-        define_macros.append(f"#define INPUT_{int(i)} {inp}")
-        undef_macros.append(f"#undef INPUT_{int(i)}")
+    undef_inputs = [f"#undef INPUT_{int(i)}" for i in range(len(inputs))]
+    undef_outputs = [f"#undef OUTPUT_{int(i)}" for i in range(len(outputs))]
 
-    for i, out in enumerate(outputs):
-        define_macros.append(f"#define OUTPUT_{int(i)} {out}")
-        undef_macros.append(f"#undef OUTPUT_{int(i)}")
+    define_all = "\n".join(define_inputs + define_outputs)
+    undef_all = "\n".join(undef_inputs + undef_outputs)
 
-    return "\n".join(define_macros), "\n".join(undef_macros)
+    return define_all, undef_all
 
 
 class ExternalCOp(COp):
@@ -560,9 +551,10 @@ def get_c_macros(
         define_macros.append(define_template % ("APPLY_SPECIFIC(str)", f"str##_{name}"))
         undef_macros.append(undef_template % "APPLY_SPECIFIC")
 
-        for n, v in self.__get_op_params():
-            define_macros.append(define_template % (n, v))
-            undef_macros.append(undef_template % (n,))
+        define_macros.extend(
+            define_template % (n, v) for n, v in self.__get_op_params()
+        )
+        undef_macros.extend(undef_template % (n,) for n, _ in self.__get_op_params())
 
         return "\n".join(define_macros), "\n".join(undef_macros)