Generic Protocol for transformed function

[YouJiacheng]

With ParamSpec we can completely annotate jit signature.
Thus type checking can be performed with transformed function. (and lower)

T = TypeVar("T")
P = ParamSpec("P")
def jit(fun: Callable[P, T], ...) -> stages.Wrapped[P, T]:
  ...

# stages.py
T_co = TypeVar("T_co", covariant=True)
P_contra = ParamSpec("P_contra", contravariant=True)
class Wrapped(Protocol[P_contra, T_co]):
  def __call__(self, *args: P_contra.args, **kwargs: P_contra.kwargs) -> T_co:
    ...

  def lower(self, *args: P_contra.args, **kwargs: P_contra.kwargs) -> Lowered:
    ...

Just like #9999 , and it works almost perfectly with Visual Studio Code + pylance, signature(including parameters name) of function is fully preserved.
However, mypy currently doesn't handle partial + function-transform correctly. python/mypy#12593
We can wait mypy fix this bug, or make some workaround.
Some workaround option:

make jit a polymorphism decorator + decorator factory.

toy example

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

T_co = TypeVar('T_co', covariant=True)
T = TypeVar('T')
P = ParamSpec('P')

class Wrapped(Protocol[P, T_co]):
  def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T_co:
    pass

@overload
def jit(f: Callable[P, T], *, a: bool) -> Wrapped[P, T]:
    ...

@overload
def jit(f: Literal[None] = None, *, a: bool) -> Callable[[Callable[P, T]], Wrapped[P, T]]:
    ...

def jit(f=None, *, a):
    if f is None:
        return lambda f: jit(f, a=a)
    # original code here
    return f

@jit(a=True)
def f(x: int, y: int):
    return x + y

f(0, 1) #mypy ok
f(0, 'a') # mypy error  Argument 2 to "__call__" of "Wrapped" has incompatible type "str"; expected "int"

This break "Our general philosophy is to make the JAX core API as simple and explicit as possible" @mattjj #184 (comment)
I personally like this option, but it breaks the philosophy of JAX.

add a dedicated decorator factory function and use it internally

toy example

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

T_co = TypeVar('T_co', covariant=True)
T = TypeVar('T')
P = ParamSpec('P')

class Wrapped(Protocol[P, T_co]):
  def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T_co:
    pass

def jit(f: Callable[P, T], *, a: bool) -> Wrapped[P, T]:
    # original code here
    return f

def _make_jit(*, a: bool) -> Callable[[Callable[P, T]], Wrapped[P, T]]:
    return lambda f: jit(f, a=a) # type: ignore # a mypy bug fixed in mypy-0.950

@_make_jit(a=True)
def f(x: int, y: int):
    return x + y

f(0, 0) # mypy ok
f(0, 'a') # mypy error  Argument 2 to "__call__" of "Wrapped" has incompatible type "str"; expected "int"

I think this option is satisfactory for everyone. (also can be a public API)
In addtion, we can write a function that can create decorator factory function without losing type information, but current(even the next release) mypy has bug in this case as well. python/mypy#12595

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

T_co = TypeVar('T_co', covariant=True)
T = TypeVar('T')
P = ParamSpec('P')

class Wrapped(Protocol[P, T_co]):
  def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T_co:
    pass

def jit(f: Callable[P, T], *, a: bool) -> Wrapped[P, T]:
    # original code here
    return f

_T = TypeVar('_T')
_P = ParamSpec('_P')

def _make_factory(
    transform: Callable[Concatenate[Callable[P, T], _P], _T] # type: ignore
    # mypy can't recognize Concatenate, mypy-0.950 has fixed it
) -> Callable[_P, Callable[[Callable[P, T]], _T]]:
    return lambda *args, **kwargs: lambda f: transform(f, *args, **kwargs)

_make_jit: Callable[..., Callable[[Callable[P, T]], Wrapped[P, T]]] = _make_factory(jit)
# mypy need a type annotation, but this annotation will erase signature of jit -> bad for vscode
# mypy-0.950 has not fixed it

@_make_jit(a=True)
def f(x: int, y: int):
    return x + y

f(0, 0) # mypy ok
f(0, 'a') # mypy error Argument 2 to "__call__" of "Wrapped" has incompatible type "str"; expected "int"

[YouJiacheng]

@jakevdp @mattjj @froystig
Could you please give some advice?

[froystig]

I commented in #10321 about option 2 ("dedicated factory function"), but it possibly applies to option 1 ("polymorphic jit") as well. I think we have a rather strict constraint to maintain: keep jit as it is (no double meanings based on how it is called), and allow for partial(jit, ...). We have many downstream dependencies on jit at this point (many of which take partial), and possibly many more to come, so we want to keep this as simple and explicit as possible.

These are both clever approaches and I'm learning simply by reading them. Thanks for thinking about this. Are there any other approaches you can think of? Following your comment here, does the issue with mypy+partial seem temporary? If so, can we make this work for jit without partial for now, without trying to work around mypy's partial bug, and anticipate a future fix?

[YouJiacheng]

Yes, I think the mypy issue is temporary.
Do you have any idea about "make this work fof jit without partial"? Can we make the jit with partial be specially typed?

[froystig]

My suggestion is that we pretend the mypy issue is already fixed, so we do not try to work around it (for now). In that case, what would be the best approach to annotation?

[YouJiacheng]

There is something surprising: as the last case of python/mypy#12593 ,mypy can't handle partial + jit with typevar annotated fun correctly. But IIRC, we use this annotation a few month ago.

[YouJiacheng]

In that case, maybe we can annotate jit return-type as Union[Wrapped[P, T], Wrapped], and annotated _make_jit return-type as Wrapped[P, T].
As a result, any downstream codebase won't be broken, and we can do type-checking for jitted JAX functions with _make_jit instead of partial.
And users using vscode can enjoy a nice jitted function signature hint.

[froystig]

When the mypy bug is fixed, would we be able to remove _make_jit and only annotate jit's return type as Wrapped[P, T]?

[YouJiacheng]

Yes. But I think it is not bad to have _make_jit
mypy won't report an error if partial pass wrong argument.

def f(x: int, y: int):
  return x + y
g = partial(f, z=0)
h = partial(f, 'a')
# no error with mypy

[YouJiacheng]

Ah, I find that mypy support for partial is too weak...

def jit(
    fun: Callable[P, T],
    *,
    static_argnums: Union[int, Iterable[int], None] = None,
    static_argnames: Union[str, Iterable[str], None] = None,
    device: Optional[xc.Device] = None,
    backend: Optional[str] = None,
    donate_argnums: Union[int, Iterable[int]] = (),
    inline: bool = False,
    abstracted_axes: Optional[Any] = None,
  ) -> Union[stages.GenericWrapped[P, T], stages.Wrapped]:

class GenericWrapped(Protocol[P, T_co]):
  @overload
  def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T_co:
    """Executes the wrapped function, lowering and compiling as needed."""
  @overload
  def __call__(self, *args: Any, **kwargs: Any) -> T_co: ... # without this mypy will expect P.args

  @overload
  def lower(self, *args: P.args, **kwargs: P.kwargs) -> Lowered:
    """Lower this function for the given arguments.

    A lowered function is staged out of Python and translated to a
    compiler's input language, possibly in a backend-dependent
    manner. It is ready for compilation but not yet compiled.

    Returns:
      A ``Lowered`` instance representing the lowering.
    """
  @overload
  def lower(self, *args: Any, **kwargs: Any) -> Lowered: ...

class Wrapped(Protocol):
  def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
  def lower(self, *args: Any, **kwargs: Any) -> Lowered: ...

And

from jax import jit
from functools import partial

@partial(jit)
def f(x: float, y: float):
    return x + y

x: float = f(0.0, 0.0)

error: Incompatible types in assignment (expression has type "Union[T, Any]", variable has type "float")  [assignment]

😭
This actually is caused by that mypy will erase the parameter specification of func for partial, only retain the return type.
So the TypeVar inside func is erased.

[YouJiacheng]

Due to the mypy‘s limitations, we can only leave a non-generic jit here, while using _make_jit internally. Or we can expose _make_jit as a public API for people who want to preserve function signature with JIT.