Merge pull request #3354 from chiamp:dunder

PiperOrigin-RevId: 567429834
google · Sep 21, 2023 · d33a33e · d33a33e
2 parents 79915d2 + dbc9254
commit d33a33e
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 8 deletions.
diff --git a/docs/guides/dropout.rst b/docs/guides/dropout.rst
@@ -212,7 +212,7 @@ the training step function. Refer to the
   from flax.training import train_state
 
   class TrainState(train_state.TrainState): #!
-    key: jax.random.KeyArray #!
+    key: jax.Array #!
 
   state = TrainState.create( #!
     apply_fn=my_model.apply,

diff --git a/examples/seq2seq/models.py b/examples/seq2seq/models.py
@@ -26,7 +26,7 @@
 import numpy as np
 
 Array = jax.Array
-PRNGKey = jax.random.KeyArray
+PRNGKey = jax.Array
 LSTMCarry = Tuple[Array, Array]
 
 

diff --git a/flax/linen/module.py b/flax/linen/module.py
@@ -74,7 +74,7 @@
 
 traceback_util.register_exclusion(__file__)
 
-KeyArray = Union[jax.Array, jax.random.KeyArray]  # pylint: disable=invalid-name
+KeyArray = jax.Array
 RNGSequences = Dict[str, KeyArray]
 Array = Any  # pylint: disable=invalid-name
 

diff --git a/flax/linen/recurrent.py b/flax/linen/recurrent.py
@@ -42,7 +42,7 @@
 from typing_extensions import Protocol
 
 A = TypeVar('A')
-PRNGKey = jax.random.KeyArray
+PRNGKey = jax.Array
 Shape = Tuple[int, ...]
 Dtype = Any  # this could be a real type?
 Array = jax.Array
@@ -750,7 +750,7 @@ def __call__(
       inputs: jax.Array,
       *,
       initial_carry: Optional[Carry] = None,
-      init_key: Optional[random.KeyArray] = None,
+      init_key: Optional[PRNGKey] = None,
       seq_lengths: Optional[Array] = None,
       return_carry: Optional[bool] = None,
       time_major: Optional[bool] = None,
@@ -976,7 +976,7 @@ def __call__(
       inputs: jax.Array,
       *,
       initial_carry: Optional[Carry] = None,
-      init_key: Optional[random.KeyArray] = None,
+      init_key: Optional[PRNGKey] = None,
       seq_lengths: Optional[Array] = None,
       return_carry: Optional[bool] = None,
       time_major: Optional[bool] = None,
@@ -1000,7 +1000,7 @@ def __call__(
       inputs: jax.Array,
       *,
       initial_carry: Optional[Carry] = None,
-      init_key: Optional[random.KeyArray] = None,
+      init_key: Optional[PRNGKey] = None,
       seq_lengths: Optional[Array] = None,
       return_carry: Optional[bool] = None,
       time_major: Optional[bool] = None,

diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py
@@ -25,7 +25,7 @@
 from jax import random
 import jax.numpy as jnp
 
-KeyArray = Union[jax.Array, jax.random.KeyArray]
+KeyArray = jax.Array
 
 
 class Dropout(Module):