hooks: fix for torch and triton backends (#2751)

marcelotduarte · Dec 29, 2024 · 962861d · 962861d
1 parent 181ce49
commit 962861d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/cx_Freeze/hooks/torch.py b/cx_Freeze/hooks/torch.py
@@ -128,3 +128,9 @@ def load_torch__dynamo_skipfiles(finder: ModuleFinder, module: Module) -> None:
         dont_inherit=True,
         optimize=finder.optimize,
     )
+
+
+def load_torch__numpy(finder: ModuleFinder, module: Module) -> None:
+    """Patch to work in Windows."""
+    finder.exclude_module("torch._numpy.testing")
+    finder.include_package(module.name)
diff --git a/cx_Freeze/hooks/triton.py b/cx_Freeze/hooks/triton.py
@@ -19,3 +19,10 @@ def load_triton(finder: ModuleFinder, module: Module) -> None:
     source_lib = module.file.parent / "_C"
     if source_lib.exists():
         finder.include_files(source_lib, f"lib/{module.name}/_C")
+
+    # exclude backends module, but include its source modules
+    finder.exclude_module("triton.backends")
+    backends_path = module.file.parent / "backends"
+    for source in backends_path.rglob("*.py"):  # type: Path
+        target = f"lib/{module.name}" / source.relative_to(backends_path)
+        finder.include_files(source, target)