Fixed errors from new pyright

patrick-kidger · Apr 14, 2024 · b146089 · b146089
1 parent 25531dc
commit b146089
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 4 deletions.
diff --git a/equinox/_ad.py b/equinox/_ad.py
@@ -566,9 +566,11 @@ def f(x, y):
         # Skip jaxpr tracing for efficiency.
         closure_converted = _TrivialClosureConvert(fn, in_dynamic_struct, in_static)
     else:
+        fn = cast(Callable[_P, _T], fn)
         closed_jaxpr, out_dynamic_struct, out_static = filter_make_jaxpr(fn)(
-            *args, **kwargs
-        )  # pyright: ignore
+            *args,  # pyright: ignore
+            **kwargs,
+        )
         jaxpr = closed_jaxpr.jaxpr
         consts = closed_jaxpr.consts
         out_dynamic_struct = jtu.tree_flatten(out_dynamic_struct)

diff --git a/equinox/internal/_omega.py b/equinox/internal/_omega.py
@@ -69,7 +69,7 @@ def at(self):
 
 
 if TYPE_CHECKING:
-    ω: Any
+    ω: Any = ...
 else:
     ω = _ω
 

diff --git a/equinox/nn/_misc.py b/equinox/nn/_misc.py
@@ -11,7 +11,7 @@
     # https://github.com/microsoft/pyright/issues/3450
     def all_sequences(
         x: Union[Sequence[Any], Sequence[_T]],
-    ) -> "te.StrictTypeGuard[Sequence[_T]]":
+    ) -> "te.TypeIs[Sequence[_T]]":
         ...
 
 else:

diff --git a/pyproject.toml b/pyproject.toml
@@ -52,4 +52,5 @@ order-by-type = false
 [tool.pyright]
 reportIncompatibleMethodOverride = true
 reportIncompatibleVariableOverride = false  # Incompatible with eqx.AbstractVar
+reportFunctionMemberAccess = false
 include = ["equinox", "tests"]