Add type hints to methods of JitWrapped

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>
jax-ml · Sep 18, 2024 · a79de59 · a79de59
1 parent 90047a8
commit a79de59
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
@@ -812,11 +812,11 @@ def ax_leaf(l):
 
 class JitWrapped(stages.Wrapped[_P, _OutT]):
 
-  def eval_shape(self, *args, **kwargs):
+  def eval_shape(self, *args: _P.args, **kwargs: _P.kwargs):
     """See ``jax.eval_shape``."""
     raise NotImplementedError
 
-  def trace(self, *args, **kwargs) -> stages.Traced:
+  def trace(self, *args: _P.args, **kwargs: _P.kwargs) -> stages.Traced:
     raise NotImplementedError