Add IDE friendly type annotation for jit

[YouJiacheng]

see #9917

[froystig]

Did #9950 take care of this?

[YouJiacheng]

@froystig Thanks for your reply! #9950 is awesome, but it didn't take care of function's signature. Can you make a generic-type protocol?

[YouJiacheng]

@froystig
See:

from typing import TypeVar, Callable
from typing_extensions import ParamSpec, Protocol

T = TypeVar('T')
P = ParamSpec('P')

class Wrapped(Protocol[P, T]):
  def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
    """Executes the wrapped function, lowering and compiling as needed."""

def jit(f: Callable[P, T]) -> Wrapped[P, T]:
    pass

def f(x: float, y: float):
    return x + y

May I create a PR for it?

[YouJiacheng]

#10311 Instead.