Fix PERF403: dict comprehensions when appropriate

Author: Armavica

pytensor/link/vm.py

@@ -1055,12 +1055,7 @@ def make_vm(
             for v in self.fgraph.inputs + self.fgraph.outputs:
                 vars_idx.setdefault(v, len(vars_idx))
 
-            nodes_idx_inv = {}
-            vars_idx_inv = {}
-            for node, i in nodes_idx.items():
-                nodes_idx_inv[i] = node
-            for var, i in vars_idx.items():
-                vars_idx_inv[i] = var
+            vars_idx_inv = {i: var for var, i in vars_idx.items()}
 
             # put storage_map and compute_map into a int-based scheme
             storage_map_list = [

pytensor/scalar/basic.py

@@ -4434,8 +4434,7 @@ def apply(self, fgraph):
             )
             # make sure we don't produce any float16.
             assert not any(o.dtype == "float16" for o in new_node.outputs)
-            for o, no in zip(node.outputs, new_node.outputs):
-                mapping[o] = no
+            mapping.update(zip(node.outputs, new_node.outputs))
 
         new_ins = [mapping[inp] for inp in fgraph.inputs]
         new_outs = [mapping[out] for out in fgraph.outputs]

pytensor/scan/op.py

@@ -2240,8 +2240,7 @@ def infer_shape(self, fgraph, node, input_shapes):
         # Non-sequences have a direct equivalent from self.inner_inputs in
         # node.inputs
         inner_non_sequences = self.inner_inputs[len(seqs_shape) + len(outs_shape) :]
-        for in_ns, out_ns in zip(inner_non_sequences, node.inputs[offset:]):
-            out_equivalent[in_ns] = out_ns
+        out_equivalent.update(zip(inner_non_sequences, node.inputs[offset:]))
 
         if info.as_while:
             self_outs = self.inner_outputs[:-1]

pytensor/scan/rewriting.py

@@ -1623,9 +1623,7 @@ def scan_save_mem(fgraph, node):
         (inps, outs, info, node_ins, compress_map) = compress_outs(
             op, not_required, nw_inputs
         )
-        inv_compress_map = {}
-        for k, v in compress_map.items():
-            inv_compress_map[v] = k
+        inv_compress_map = {v: k for k, v in compress_map.items()}
 
         # 3.6 Compose the new scan
         # TODO: currently we don't support scan with 0 step. So