[triton] Python typing improvements

[Mogball]

No description provided.

[Jokeren]

Thanks. To confirm it hasn't fixed constexpr problems like the following, right?

    block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 4], [1, 32], [4, 1], [1, 0])

Incompatible types in assignment (expression has type "BlockedLayout", variable has type "constexpr")Mypyassignment

[Mogball]

Thanks. To confirm it hasn't fixed constexpr problems like the following, right?

    block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 4], [1, 32], [4, 1], [1, 0])

Incompatible types in assignment (expression has type "BlockedLayout", variable has type "constexpr")Mypyassignment

No, I haven't gotten to that yet.

[saagarjha]

We also have

diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py
index eaaffdfbd..0a5460e01 100644
--- a/python/triton/runtime/jit.py
+++ b/python/triton/runtime/jit.py
@@ -10,7 +10,7 @@ import textwrap
 from collections import defaultdict
 from dataclasses import dataclass
 from functools import cached_property
-from typing import Callable, Generic, Iterable, Optional, ParamSpec, TypeVar, overload, Dict, Any, Tuple
+from typing import Callable, Concatenate, Generic, Iterable, Optional, ParamSpec, TYPE_CHECKING, TypeVar, overload, Dict, Any, Tuple

 from triton.backends import BaseBackend
 from types import ModuleType
@@ -29,6 +29,7 @@ INDENT_PATTERN = re.compile(r"^(?P<indent>[ \t]*)def\s+\w+\s*\(", re.MULTILINE)
 T = TypeVar("T")
 P = ParamSpec("P")
 R = TypeVar("R")
+U = TypeVar("U")

 # -----------------------------------------------------------------------------
 # Dependencies Finder
@@ -907,6 +908,20 @@ class JITFunction(JITCallable, KernelInterface[T]):
     def __call__(self: "JITFunction[Callable[P, R]]", *args: P.args, **kwargs: P.kwargs) -> R:
         raise RuntimeError("Cannot call @triton.jit'd outside of the scope of a kernel")

+    if TYPE_CHECKING:
+
+        @overload
+        def __get__(self: "JITFunction[Callable[Concatenate[U, P], R]]", instance: U,
+                    owner: Optional[type] = None) -> Callable[P, R]:
+            ...
+
+        @overload
+        def __get__(self, instance: None, owner: Optional[type] = None) -> "JITFunction[T]":
+            ...
+
+        def __get__(self, instance, owner=None):
+            ...
+
     def __repr__(self):
         return f"JITFunction({self.module}:{self.fn.__qualname__})"

to fix member functions on aggregates

[jeffniu-openai]

PR plz