Type Enhancement for Func Transforms and Bug Fix (#3003)

Co-authored-by: Awni Hannun <awni@apple.com>
2026-01-19 22:19:57 +00:00
parent d96a2bdf57
commit 46d0fdc5ec
3 changed files with 24 additions and 6 deletions
@@ -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