Fixed wrong PRNGKey annotation.

patrick-kidger · Nov 27, 2023 · dce2fa1 · dce2fa1
1 parent 588a8c3
commit dce2fa1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/equinox/nn/_batch_norm.py b/equinox/nn/_batch_norm.py
@@ -4,7 +4,7 @@
 import jax
 import jax.lax as lax
 import jax.numpy as jnp
-from jaxtyping import Array, Bool, Float
+from jaxtyping import Array, Bool, Float, PRNGKeyArray
 
 from .._module import field
 from ._sequential import StatefulLayer
@@ -111,7 +111,7 @@ def __call__(
         x: Array,
         state: State,
         *,
-        key: Optional["jax.random.PRNGKey"] = None,  # pyright: ignore
+        key: Optional[PRNGKeyArray] = None,
         inference: Optional[bool] = None,
     ) -> tuple[Array, State]:
         """**Arguments:**

diff --git a/equinox/nn/_spectral_norm.py b/equinox/nn/_spectral_norm.py
@@ -111,7 +111,7 @@ def __call__(
         x: Array,
         state: State,
         *,
-        key: Optional["jax.random.PRNGKey"] = None,  # pyright: ignore
+        key: Optional[PRNGKeyArray] = None,
         inference: Optional[bool] = None
     ) -> tuple[Array, State]:
         """**Arguments:**