Fix typing annotations for @jax.named_call

PiperOrigin-RevId: 578852649
jax-ml · Nov 2, 2023 · 81ac67f · 81ac67f
1 parent 1c66ac5
commit 81ac67f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/jax/_src/api.py b/jax/_src/api.py
@@ -2835,10 +2835,10 @@ def eval_shape(fun, *args, **kwargs):
 
 
 def named_call(
-    fun: Callable[..., Any],
+    fun: F,
     *,
     name: str | None = None,
-  ) -> Callable[..., Any]:
+) -> F:
   """Adds a user specified name to a function when staging out JAX computations.
 
   When staging out computations for just-in-time compilation to XLA (or other
@@ -2867,6 +2867,7 @@ def named_call(
 
   return source_info_util.extend_name_stack(name)(fun)
 
+
 @contextmanager
 def named_scope(
     name: str,