Typing improvement: Preserve the wrapped callable's signature in `jax.jit`

[lebrice]

• [X] Check for duplicate requests.
• [X] Describe your goal, and if possible provide a code snippet with a motivating example.

---

Hello there. Simple feature request / bug report:

Currently, jax.jit drops the typing signature of the wrapped callable. For example this code shows no warnings in a code editor:

import jax

@jax.jit
def foo(a: jax.Array) -> jax.Array:
    return a

foo(bob=123)  # type-checker should display a warning!

Same goes for functions or methods annotated with a functools.partial of jax.jit: The signature of the wrapped callable is dropped:

import jax
import functools

@functools.partial(jax.jit, static_argnames=["some_static_arg"])
def foo_with_static_arg(a: jax.Array, some_static_arg: Any) -> jax.Array:
    return a

foo_with_static_arg(bob=123)  # type-checker should also display a warning here!

Is there something I'm not aware of that might make this undesirable for some reason?

In the meantime, I made #23720 to address this. Let me know what you think :)

[jakevdp]

Crossref #14688, which was a previous attempt to solve this.