Type Enhancement for Func Transforms and Bug Fix

python/mlx/_stub_patterns.txt

@@ -1,10 +1,12 @@
 mlx.core.__prefix__:
-  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, ParamSpec, TypeVar
   import sys
   if sys.version_info >= (3, 10):
     from typing import TypeAlias
   else:
     from typing_extensions import TypeAlias
+  P = ParamSpec("P")
+  R = TypeVar("R")
 
 mlx.core.__suffix__:
   from typing import Union

python/src/transforms.cpp

@@ -1333,7 +1333,7 @@ void init_transforms(nb::module_& m) {
       "argnums"_a = nb::none(),
       "argnames"_a = std::vector<std::string>{},
       nb::sig(
-          "def value_and_grad(fun: Callable, argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable"),
+          "def value_and_grad(fun: Callable[P, R], argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable[P, Tuple[R, Any]]"),
       R"pbdoc(
         Returns a function which computes the value and gradient of ``fun``.
 
@@ -1472,7 +1472,7 @@ void init_transforms(nb::module_& m) {
       "outputs"_a = nb::none(),
       "shapeless"_a = false,
       nb::sig(
-          "def compile(fun: Callable, inputs: Optional[object] = None, outputs: Optional[object] = None, shapeless: bool = False) -> Callable"),
+          "def compile(fun: Callable[P, R], inputs: Optional[object] = None, outputs: Optional[object] = None, shapeless: bool = False) -> Callable[P, R]"),
       R"pbdoc(
         Returns a compiled function which produces the same output as ``fun``.
 

setup.py

@@ -169,10 +169,15 @@ def run(self):
             self.copy_tree(regular_dir, inplace_dir)
 
         # Build type stubs.
+        ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name)  # type: ignore[no-untyped-call]
+        extdir = ext_fullpath.parent.parent.resolve()
         build_temp = Path(self.build_temp) / ext.name
+        env = os.environ.copy()
+        env["PYTHONPATH"] += f":{extdir}"
         subprocess.run(
             ["cmake", "--install", build_temp, "--component", "core_stub"],
             check=True,
+            env=env,
         )