Clean Ruff settings

pyproject.toml

@@ -210,7 +210,6 @@ module = ['networkx', 'optax', 'optax.contrib', 'array_api_compat']
 ignore_missing_imports = true
 
 [tool.ruff]
-target-version = 'py310'
 line-length = 100
 
 [tool.ruff.lint]
@@ -258,9 +257,6 @@ ignore = [
   'TID252', # Relative imports from parent modules are banned.
 ]
 
-[tool.ruff.lint.flake8-annotations]
-mypy-init-return = true
-
 [tool.ruff.lint.flake8-errmsg]
 max-string-length = 40
 

tests/test_flax_module.py

@@ -13,7 +13,7 @@
 @pytest.mark.skip
 def test_dataclass_module() -> None:
     class SomeModule(nnx.Module):
-        def __init__(self, epsilon: Array):
+        def __init__(self, epsilon: Array) -> None:
             super().__init__()
             self.epsilon = epsilon
 

tjax/_src/rng.py

@@ -9,7 +9,7 @@
 
 
 class RngStream:
-    def __init__(self, key: KeyArray, count: JaxIntegralArray | None = None):
+    def __init__(self, key: KeyArray, count: JaxIntegralArray | None = None) -> None:
         super().__init__()
         if count is None:
             count = jnp.zeros((), dtype=jnp.uint32)

tjax/_src/shims.py

@@ -42,7 +42,8 @@ class custom_vjp(Generic[P, R_co]):  # noqa: N801
     def __init__(self,
                  func: Callable[P, R_co],
                  *,
-                 static_argnums: tuple[int, ...] = ()):
+                 static_argnums: tuple[int, ...] = ()
+                 ) -> None:
         super().__init__()
         static_argnums = tuple(sorted(static_argnums))
         self.vjp = jax.custom_vjp(func, nondiff_argnums=static_argnums)
@@ -72,7 +73,8 @@ class custom_vjp_method(Generic[U, P, R_co]):  # noqa: N801
     def __init__(self,
                  func: Callable[Concatenate[U, P], R_co],
                  *,
-                 static_argnums: tuple[int, ...] = ()):
+                 static_argnums: tuple[int, ...] = ()
+                 ) -> None:
         super().__init__()
         static_argnums = tuple(sorted(static_argnums))
         self.vjp = jax.custom_vjp(func, nondiff_argnums=static_argnums)
@@ -117,7 +119,8 @@ class custom_jvp(Generic[P, R_co]):  # noqa: N801
     def __init__(self,
                  func: Callable[P, R_co],
                  *,
-                 nondiff_argnums: tuple[int, ...] = ()):
+                 nondiff_argnums: tuple[int, ...] = ()
+                 ) -> None:
         super().__init__()
         nondiff_argnums = tuple(sorted(nondiff_argnums))
         self.jvp = jax.custom_jvp(func, nondiff_argnums=nondiff_argnums)
@@ -147,7 +150,8 @@ class custom_jvp_method(Generic[U, P, R_co]):  # noqa: N801
     def __init__(self,
                  func: Callable[Concatenate[U, P], R_co],
                  *,
-                 nondiff_argnums: tuple[int, ...] = ()):
+                 nondiff_argnums: tuple[int, ...] = ()
+                 ) -> None:
         super().__init__()
         nondiff_argnums = tuple(sorted(nondiff_argnums))
         self.jvp = jax.custom_jvp(func, nondiff_argnums=nondiff_argnums)