From b01fc7eac79289efd70f2475adf3f3a36fc5159a Mon Sep 17 00:00:00 2001
From: Cheng <zcbenz@gmail.com>
Date: Mon, 22 Dec 2025 09:42:20 +0900
Subject: [PATCH] Fix stubgen (#2942)

---
 setup.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/setup.py b/setup.py
index d26800be..61aae881 100644
--- a/setup.py
+++ b/setup.py
@@ -181,7 +181,7 @@ class GenerateStubs(Command):
         pass
 
     def run(self) -> None:
-        out_path = "python/mlx/core"
+        out_path = "python/mlx"
         stub_cmd = [
             "python",
             "-m",
@@ -192,9 +192,6 @@ class GenerateStubs(Command):
             "python/mlx/_stub_patterns.txt",
         ]
         subprocess.run(stub_cmd + ["-r", "-O", out_path])
-        # Run again without recursive to specify output file name
-        subprocess.run(["rm", f"{out_path}/mlx.pyi"])
-        subprocess.run(stub_cmd + ["-o", f"{out_path}/__init__.pyi"])
 
 
 class MLXBdistWheel(bdist_wheel):