Fix PERF401: list comprehensions when appropriate

Author: Armavica

pytensor/compile/debugmode.py

@@ -756,10 +756,7 @@ def _get_preallocated_maps(
     # TODO: Sparse? Scalar does not really make sense.
 
     # Do not preallocate memory for outputs that actually work inplace
-    considered_outputs = []
-    for r in node.outputs:
-        if r not in inplace_outs:
-            considered_outputs.append(r)
+    considered_outputs = [r for r in node.outputs if r not in inplace_outs]
 
     # Output storage that was initially present in the storage_map
     if "initial" in prealloc_modes or "ALL" in prealloc_modes:

pytensor/link/c/basic.py

@@ -1450,12 +1450,16 @@ def in_sig(i, topological_pos, i_idx):
             if props:
                 version.append(props)
 
-            for i in node.inputs:
-                if isinstance(i.type, CLinkerObject):
-                    version.append(i.type.c_code_cache_version())
-            for o in node.outputs:
-                if isinstance(o.type, CLinkerObject):
-                    version.append(o.type.c_code_cache_version())
+            version.extend(
+                i.type.c_code_cache_version()
+                for i in node.inputs
+                if isinstance(i.type, CLinkerObject)
+            )
+            version.extend(
+                o.type.c_code_cache_version()
+                for o in node.outputs
+                if isinstance(o.type, CLinkerObject)
+            )
 
             # add the signature for this node
             sig.append(

pytensor/link/c/cmodule.py

@@ -2131,9 +2131,11 @@ def get_lines(cmd, parse=True):
                             or "-march=native" in line
                         ):
                             continue
-                        for reg in ("-march=", "-mtune=", "-target-cpu", "-mabi="):
-                            if reg in line:
-                                selected_lines.append(line.strip())
+                        selected_lines.extend(
+                            line.strip()
+                            for reg in ("-march=", "-mtune=", "-target-cpu", "-mabi=")
+                            if reg in line
+                        )
                     lines = list(set(selected_lines))  # to remove duplicate
 
                 return lines

pytensor/link/vm.py

@@ -1270,15 +1270,16 @@ def make_all(
         if self.allow_gc:
             post_thunk_clear = []
             for node in order:
-                clear_after_this_thunk = []
-                for input in node.inputs:
+                clear_after_this_thunk = [
+                    storage_map[input]
+                    for input in node.inputs
                     if (
                         input in computed
                         and input not in fgraph.outputs
                         and node == last_user[input]
                         and input not in reallocated_vars
-                    ):
-                        clear_after_this_thunk.append(storage_map[input])
+                    )
+                ]
                 post_thunk_clear.append(clear_after_this_thunk)
         else:
             post_thunk_clear = None

pytensor/tensor/conv/abstract_conv.py

@@ -585,10 +585,11 @@ def assert_shape(x, expected_shape, msg="Unexpected shape."):
     if expected_shape is None or not config.conv__assert_shape:
         return x
     shape = x.shape
-    tests = []
-    for i in range(x.ndim):
-        if expected_shape[i] is not None:
-            tests.append(pt.eq(shape[i], expected_shape[i]))
+    tests = [
+        pt.eq(shape[i], expected_shape[i])
+        for i in range(x.ndim)
+        if expected_shape[i] is not None
+    ]
     if tests:
         return Assert(msg)(x, *tests)
     else:

pytensor/tensor/subtensor.py

@@ -2107,10 +2107,7 @@ def perform(self, node, inp, out_):
         out[0] = x.take(i, axis=0, out=o)
 
     def connection_pattern(self, node):
-        rval = [[True]]
-
-        for ipt in node.inputs[1:]:
-            rval.append([False])
+        rval = [[True], *([False] for _ in node.inputs[1:])]
 
         return rval
 

tests/sparse/test_basic.py

@@ -1784,13 +1784,14 @@ def f_b(z, a, x, y):
                     )
                     == len(topo) - 5
                 )
-                new_topo = []
-                for node in topo:
+                new_topo = [
+                    node
+                    for node in topo
                     if not (
                         isinstance(node.op, Elemwise)
                         and isinstance(node.op.scalar_op, pytensor.scalar.basic.Cast)
-                    ):
-                        new_topo.append(node)
+                    )
+                ]
                 topo = new_topo
                 assert len(topo) == 5, topo
                 # Usmm is tested at the same time in debugmode

tests/tensor/test_math.py

@@ -3477,8 +3477,9 @@ def test_grad_useless_sum():
     old_values_eq_approx = staticmethod(TensorType.values_eq_approx)
     TensorType.values_eq_approx = staticmethod(values_eq_approx_remove_nan)
     try:
-        for test_value in test_values:
-            outputs.append(f(np.array([test_value]).astype("float32")))
+        outputs.extend(
+            f(np.array([test_value]).astype("float32")) for test_value in test_values
+        )
     finally:
         TensorType.values_eq_approx = old_values_eq_approx
 

tests/tensor/test_subtensor.py

@@ -1143,9 +1143,11 @@ def test_grad_list(self):
         data = random(4)
         data = np.asarray(data, dtype=self.dtype)
         idxs = [[i] for i in range(data.shape[0])]
-        for i in range(data.shape[0]):
-            for j in range(0, data.shape[0], 2):
-                idxs.append([i, j, (i + 1) % data.shape[0]])
+        idxs.extend(
+            [i, j, (i + 1) % data.shape[0]]
+            for i in range(data.shape[0])
+            for j in range(0, data.shape[0], 2)
+        )
         self.grad_list_(idxs, data)
 
         data = random(4, 3)

tests/typed_list/test_type.py

@@ -78,9 +78,7 @@ def test_load_alot(self):
         myType = TypedListType(TensorType(pytensor.config.floatX, shape=(None, None)))
 
         x = random_ranged(-1000, 1000, [10, 10])
-        testList = []
-        for i in range(10000):
-            testList.append(x)
+        testList = [x for _ in range(10000)]
 
         assert np.array_equal(myType.filter(testList), testList)