Remove the Type annotation for jit an pmap as there are additional attributes on the returned callable.

[copybara-service]

Remove the Type annotation for jit an pmap as there are additional attributes on the returned callable.

Using the experimental jax.jit(lambda x: x+1).lower(...) is raising an error with pytype.

[hawkinsp]

@NeilGirdhar fyi

[NeilGirdhar]

@hawkinsp What error are you getting? Seems like a bug in pytype. Could you file one there?

[jakevdp]

@NeilGirdhar My understanding is that the type annotation is incorrect because the returned object adds some methods; e.g. type checkers will incorrectly fail on the following:

from jax import jit
@jit
def f(x):
  return x
lowered = f.lower(1.0)  # <--- original function has no `lower` method, so type checking fails

The full solution would proably be to define some custom protocol, but in the meantime Any is not incorrect.

[NeilGirdhar]

The full solution would proably be to define some custom protocol

Right, that would be a better solution.

[NeilGirdhar]

@jakevdp By the way, I looked into it, but MyPy has bugs relating to that protocol: python/mypy#12169. So I guess, I won't be able to get that protocol using ParamSpec until MyPy has fixes. Could we use a protocol that doesn't use ParamSpec?

Something like this:

T = TypeVar("T", covariant=True)

# TODO: Use ParamSpec to annotate function parameters.
class JittedFunction(Protocol, Generic[T]):
  def __call__(self, *args: Any, **kwargs: Any) -> T:
    ...

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

And then where you have Callable, you could have Callable[..., T] and you could return JittedFunction[T]. What do you think?

[jblespiau]

Note also that ParamSpec is in Python 3.10. Internally, we are still at 3.7

[NeilGirdhar]

@jblespiau You can import ParamSpec from typing_extensions.