Code cleanups

pytensor/compile/compiledir.py

@@ -7,6 +7,7 @@
 import os
 import pickle
 import shutil
+from collections import Counter
 
 import numpy as np
 
@@ -111,11 +112,11 @@
     compiledir = config.compiledir
     table = []
     table_multiple_ops = []
-    table_op_class = {}
+    table_op_class = Counter()
     zeros_op = 0
     big_key_files = []
     total_key_sizes = 0
-    nb_keys = {}
+    nb_keys = Counter()
     for dir in os.listdir(compiledir):
         filename = os.path.join(compiledir, dir, "key.pkl")
         if not os.path.exists(filename):
@@ -125,9 +126,7 @@
                 keydata = pickle.load(file)
                 ops = list({x for x in flatten(keydata.keys) if isinstance(x, Op)})
                 # Whatever the case, we count compilations for OP classes.
-                for op_class in {op.__class__ for op in ops}:
-                    table_op_class.setdefault(op_class, 0)
-                    table_op_class[op_class] += 1
+                table_op_class.update({op.__class__ for op in ops})
                 if len(ops) == 0:
                     zeros_op += 1
                 else:
@@ -159,7 +158,6 @@
                 if size > max_key_file_size:
                     big_key_files.append((dir, size, ops))
 
-                nb_keys.setdefault(len(keydata.keys), 0)
                 nb_keys[len(keydata.keys)] += 1
             except OSError:
                 pass
@@ -198,8 +196,7 @@
         ),
         underline="+",
     )
-    table_op_class = sorted(table_op_class.items(), key=lambda t: t[1])
-    for op_class, nb in table_op_class:
+    for op_class, nb in reversed(table_op_class.most_common()):
         print(op_class, nb)
 
     if big_key_files:

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

@@ -14,6 +14,7 @@
 import pytensor.compile.profiling
 from pytensor.compile.io import In, SymbolicInput, SymbolicOutput
 from pytensor.compile.ops import deep_copy_op, view_op
+from pytensor.compile.profiling import ProfileStats
 from pytensor.configdefaults import config
 from pytensor.graph.basic import (
     Constant,
@@ -212,18 +213,14 @@
 
         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,
@@ -557,11 +554,11 @@
 
     def copy(
         self,
-        share_memory=False,
-        swap=None,
-        delete_updates=False,
-        name=None,
-        profile=None,
+        share_memory: bool = False,
+        swap: dict | None = None,
+        delete_updates: bool = False,
+        name: str | None = None,
+        profile: bool | str | ProfileStats | None = None,
     ):
         """
         Copy this function. Copied function will have separated maker and
@@ -588,7 +585,7 @@
             If provided, will be the name of the new
             Function. Otherwise, it will be old + " copy"
 
-        profile :
+        profile : bool | str | ProfileStats | None
             as pytensor.function profile parameter
 
         Returns
@@ -727,14 +724,8 @@
         # reinitialize new maker and create new function
         if profile is None:
             profile = config.profile or config.print_global_stats
-            # profile -> True or False
         if profile is True:
-            if name:
-                message = name
-            else:
-                message = str(profile.message) + " copy"
-            profile = pytensor.compile.profiling.ProfileStats(message=message)
-            # profile -> object
+            profile = pytensor.compile.profiling.ProfileStats(message=name)
         elif isinstance(profile, str):
             profile = pytensor.compile.profiling.ProfileStats(message=profile)
 

pytensor/compile/profiling.py

@@ -14,9 +14,9 @@
 import operator
 import sys
 import time
-from collections import defaultdict
+from collections import Counter, defaultdict
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -204,8 +204,8 @@
         self.fct_call_time = 0.0
         self.fct_callcount = 0
         self.vm_call_time = 0.0
-        self.apply_time = {}
-        self.apply_callcount = {}
+        self.apply_time = defaultdict(float)
+        self.apply_callcount = Counter()
         # self.apply_cimpl = None
         # self.message = None
 
@@ -234,9 +234,9 @@
     # Total time spent in Function.vm.__call__
     #
 
-    apply_time: dict[Union["FunctionGraph", Variable], float] | None = None
+    apply_time: dict[tuple["FunctionGraph", Apply], float]
 
-    apply_callcount: dict[Union["FunctionGraph", Variable], int] | None = None
+    apply_callcount: dict[tuple["FunctionGraph", Apply], int]
 
     apply_cimpl: dict[Apply, bool] | None = None
     # dict from node -> bool (1 if c, 0 if py)
@@ -292,10 +292,9 @@
     # param is called flag_time_thunks because most other attributes with time
     # in the name are times *of* something, rather than configuration flags.
     def __init__(self, atexit_print=True, flag_time_thunks=None, **kwargs):
-        self.apply_callcount = {}
+        self.apply_callcount = Counter()
         self.output_size = {}
-        # Keys are `(FunctionGraph, Variable)`
-        self.apply_time = {}
+        self.apply_time = defaultdict(float)
         self.apply_cimpl = {}
         self.variable_shape = {}
         self.variable_strides = {}
@@ -320,37 +319,29 @@
 
         """
         # timing is stored by node, we compute timing by class on demand
-        rval = {}
-        for (fgraph, node), t in self.apply_time.items():
-            typ = type(node.op)
-            rval.setdefault(typ, 0)
-            rval[typ] += t
-        return rval
+        rval = defaultdict(float)
+        for (_fgraph, node), t in self.apply_time.items():
+            rval[type(node.op)] += t
+        return dict(rval)
 
     def class_callcount(self):
         """
         dict op -> total number of thunk calls
 
         """
         # timing is stored by node, we compute timing by class on demand
-        rval = {}
-        for (fgraph, node), count in self.apply_callcount.items():
-            typ = type(node.op)
-            rval.setdefault(typ, 0)
-            rval[typ] += count
+        rval = Counter()
+        for (_fgraph, node), count in self.apply_callcount.items():
+            rval[type(node.op)] += count
         return rval
 
-    def class_nodes(self):
+    def class_nodes(self) -> Counter:
         """
         dict op -> total number of nodes
 
         """
         # timing is stored by node, we compute timing by class on demand
-        rval = {}
-        for (fgraph, node), count in self.apply_callcount.items():
-            typ = type(node.op)
-            rval.setdefault(typ, 0)
-            rval[typ] += 1
+        rval = Counter(type(node.op) for _fgraph, node in self.apply_callcount)
         return rval
 
     def class_impl(self):
@@ -360,12 +351,9 @@
         """
         # timing is stored by node, we compute timing by class on demand
         rval = {}
-        for fgraph, node in self.apply_callcount:
+        for _fgraph, node in self.apply_callcount:
             typ = type(node.op)
-            if self.apply_cimpl[node]:
-                impl = "C "
-            else:
-                impl = "Py"
+            impl = "C " if self.apply_cimpl[node] else "Py"
             rval.setdefault(typ, impl)
             if rval[typ] != impl and len(rval[typ]) == 2:
                 rval[typ] += impl
@@ -377,11 +365,10 @@
 
         """
         # timing is stored by node, we compute timing by Op on demand
-        rval = {}
+        rval = defaultdict(float)
         for (fgraph, node), t in self.apply_time.items():
-            rval.setdefault(node.op, 0)
             rval[node.op] += t
-        return rval
+        return dict(rval)
 
     def fill_node_total_time(self, fgraph, node, total_times):
         """
@@ -414,9 +401,8 @@
 
         """
         # timing is stored by node, we compute timing by Op on demand
-        rval = {}
+        rval = Counter()
         for (fgraph, node), count in self.apply_callcount.items():
-            rval.setdefault(node.op, 0)
             rval[node.op] += count
         return rval
 
@@ -426,10 +412,7 @@
 
         """
         # timing is stored by node, we compute timing by Op on demand
-        rval = {}
-        for (fgraph, node), count in self.apply_callcount.items():
-            rval.setdefault(node.op, 0)
-            rval[node.op] += 1
+        rval = Counter(node.op for _fgraph, node in self.apply_callcount)
         return rval
 
     def op_impl(self):
@@ -1204,8 +1187,7 @@
                         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 +1197,16 @@
                         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/gradient.py

@@ -1041,13 +1041,12 @@ def access_term_cache(node):
             # list of bools indicating if each input is connected to the cost
             inputs_connected = [
                 (
-                    True
-                    in [
+                    any(
                         input_to_output and output_to_cost
                         for input_to_output, output_to_cost in zip(
                             input_to_outputs, outputs_connected
                         )
-                    ]
+                    )
                 )
                 for input_to_outputs in connection_pattern
             ]
@@ -1067,25 +1066,24 @@ def access_term_cache(node):
             # List of bools indicating if each input only has NullType outputs
             only_connected_to_nan = [
                 (
-                    True
-                    not in [
+                    not any(
                         in_to_out and out_to_cost and not out_nan
                         for in_to_out, out_to_cost, out_nan in zip(
                             in_to_outs, outputs_connected, ograd_is_nan
                         )
-                    ]
+                    )
                 )
                 for in_to_outs in connection_pattern
             ]
 
-            if True not in inputs_connected:
+            if not any(inputs_connected):
                 # All outputs of this op are disconnected so we can skip
                 # Calling the op's grad method and report that the inputs
                 # are disconnected
                 # (The op's grad method could do this too, but this saves the
                 # implementer the trouble of worrying about this case)
                 input_grads = [disconnected_type() for ipt in inputs]
-            elif False not in only_connected_to_nan:
+            elif all(only_connected_to_nan):
                 # All inputs are only connected to nan gradients, so we don't
                 # need to bother calling the grad method. We know the gradient
                 # with respect to all connected inputs is nan.

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)