Refactor check_duplicate_key and remove mod.cu

Author: Armavica

pytensor/link/c/cmodule.py

@@ -641,7 +641,7 @@ class ModuleCache:
     The cache contains one directory for each module, containing:
     - the dynamic library file itself (e.g. ``.so/.pyd``),
     - an empty ``__init__.py`` file, so Python can import it,
-    - a file containing the source code for the module (e.g. ``mod.cpp/mod.cu``),
+    - a file containing the source code for the module (e.g. ``mod.cpp``),
     - a ``key.pkl`` file, containing a KeyData object with all the keys
     associated with that module,
     - possibly a ``delete.me`` file, meaning this directory has been marked

pytensor/misc/check_duplicate_key.py

@@ -1,6 +1,7 @@
 import os
 import pickle
 import sys
+from collections import Counter
 
 from pytensor.configdefaults import config
 
@@ -15,77 +16,63 @@
 else:
     dirs = os.listdir(config.compiledir)
     dirs = [os.path.join(config.compiledir, d) for d in dirs]
-keys: dict = {}  # key -> nb seen
+keys: Counter[bytes] = Counter()  # key -> nb seen
 mods: dict = {}
 for dir in dirs:
     key = None
     try:
-        with open(os.path.join(dir, "key.pkl")) as f:
+        with open(os.path.join(dir, "key.pkl"), "rb") as f:
             key = f.read()
-        keys.setdefault(key, 0)
         keys[key] += 1
         del f
     except OSError:
         # print dir, "don't have a key.pkl file"
         pass
     try:
         path = os.path.join(dir, "mod.cpp")
-        if not os.path.exists(path):
-            path = os.path.join(dir, "mod.cu")
-        with open(path) as f:
-            mod = f.read()
+        with open(path) as fmod:
+            mod = fmod.read()
         mods.setdefault(mod, ())
         mods[mod] += (key,)
         del mod
-        del f
+        del fmod
         del path
     except OSError:
-        print(dir, "don't have a mod.{cpp,cu} file")
+        print(dir, "don't have a mod.cpp file")
 
 if DISPLAY_DUPLICATE_KEYS:
     for k, v in keys.items():
         if v > 1:
             print("Duplicate key (%i copies): %s" % (v, pickle.loads(k)))
 
-nbs_keys: dict = {}  # nb seen -> now many key
-for val in keys.values():
-    nbs_keys.setdefault(val, 0)
-    nbs_keys[val] += 1
+# nb seen -> how many keys
+nbs_keys = Counter(val for val in keys.values())
 
-nbs_mod: dict = {}  # nb seen -> how many key
-nbs_mod_to_key = {}  # nb seen -> keys
-more_than_one = 0
-for mod, kk in mods.items():
-    val = len(kk)
-    nbs_mod.setdefault(val, 0)
-    nbs_mod[val] += 1
-    if val > 1:
-        more_than_one += 1
-    nbs_mod_to_key[val] = kk
+# nb seen -> how many keys
+nbs_mod = Counter(len(kk) for kk in mods.values())
+# nb seen -> keys
+nbs_mod_to_key = {len(kk): kk for kk in mods.values()}
+more_than_one = sum(len(kk) > 1 for kk in mods.values())
 
 if DISPLAY_MOST_FREQUENT_DUPLICATE_CCODE:
-    m = max(nbs_mod.keys())
-    print("The keys associated to the mod.{cpp,cu} with the most number of copy:")
+    m = max(nbs_mod)
+    print("The keys associated to the mod.cpp with the most number of copy:")
     for kk in nbs_mod_to_key[m]:
         kk = pickle.loads(kk)
         print(kk)
 
 print("key.pkl histograph")
-l = list(nbs_keys.items())
-l.sort()
-print(l)
+print(sorted(nbs_keys.items()))
 
-print("mod.{cpp,cu} histogram")
-l = list(nbs_mod.items())
-l.sort()
-print(l)
+print("mod.cpp histogram")
+print(sorted(nbs_mod.items()))
 
-total = sum(len(k) for k in list(mods.values()))
+total = sum(len(k) for k in mods.values())
 uniq = len(mods)
 useless = total - uniq
-print("mod.{cpp,cu} total:", total)
-print("mod.{cpp,cu} uniq:", uniq)
-print("mod.{cpp,cu} with more than 1 copy:", more_than_one)
-print("mod.{cpp,cu} useless:", useless, float(useless) / total * 100, "%")
+print("mod.cpp total:", total)
+print("mod.cpp uniq:", uniq)
+print("mod.cpp with more than 1 copy:", more_than_one)
+print("mod.cpp useless:", useless, float(useless) / total * 100, "%")
 
 print("nb directory", len(dirs))