ruff

jaxtyping/_decorator.py

@@ -70,15 +70,13 @@ def _apply_typechecker(typechecker, fn):
 def jaxtyped(
     *,
     typechecker=_sentinel,
-) -> Callable[[Callable[_Params, _Return]], Callable[_Params, _Return]]:
-    ...
+) -> Callable[[Callable[_Params, _Return]], Callable[_Params, _Return]]: ...
 
 
 @overload
 def jaxtyped(
     fn: Callable[_Params, _Return], *, typechecker=_sentinel
-) -> Callable[_Params, _Return]:
-    ...
+) -> Callable[_Params, _Return]: ...
 
 
 def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):

test/test_array.py

@@ -694,8 +694,7 @@ def test_extension(jaxtyp, typecheck, getkey):
     Y = Float[X, "b"]
 
     @jaxtyp(typecheck)
-    def f(a: X, b: Y):
-        ...
+    def f(a: X, b: Y): ...
 
     a = jr.normal(getkey(), (3, 4))
     b = jr.normal(getkey(), (4,))
@@ -708,8 +707,7 @@ def f(a: X, b: Y):
         f(a, a)
 
     @typecheck
-    def g(a: Shaped[PRNGKeyArray, "2"]):
-        ...
+    def g(a: Shaped[PRNGKeyArray, "2"]): ...
 
     with pytest.raises(ParamError):
         g(jr.PRNGKey(0))

test/test_decorator.py

@@ -14,8 +14,7 @@
 
 class M(metaclass=abc.ABCMeta):
     @jaxtyped(typechecker=None)
-    def f(self):
-        ...
+    def f(self): ...
 
     @jaxtyped(typechecker=None)
     @classmethod
@@ -39,13 +38,11 @@ def h2():
 
     @jaxtyped(typechecker=None)
     @abc.abstractmethod
-    def i1(self):
-        ...
+    def i1(self): ...
 
     @abc.abstractmethod
     @jaxtyped(typechecker=None)
-    def i2(self):
-        ...
+    def i2(self): ...
 
 
 class N:

test/test_pytree.py

@@ -166,8 +166,7 @@ class OtherCustomNamedTuple(NamedTuple):
         y: Float[jnp.ndarray, "b c"]
 
     @typecheck
-    def g(x: PyTree[CustomNamedTuple]):
-        ...
+    def g(x: PyTree[CustomNamedTuple]): ...
 
     g(
         CustomNamedTuple(