Use a Counter in compiledir

Author: Armavica

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 @@ def print_compiledir_content():
     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 @@ def print_compiledir_content():
                 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 @@ def print_compiledir_content():
                 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 @@ def print_compiledir_content():
         ),
         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: