No public description

PiperOrigin-RevId: 570964066

vmoe/evaluate/fewshot.py

@@ -39,7 +39,7 @@
 DatasetIterator = vmoe.data.input_pipeline.DatasetIterator
 PartitionSpec = jax.sharding.PartitionSpec
 PyTree = Any
-PRNGKey = jax.random.KeyArray
+PRNGKey = jax.Array
 
 BIAS_CONSTANT = 100.0
 VALID_KEY = vmoe.data.input_pipeline.VALID_KEY
@@ -333,7 +333,7 @@ def _make_fewshot_step_pjit(
     rng_keys: Sequence[str],
 ):
   """Wraps _fewshot_step with pjit."""
-  state_axis_resources = FewShotState(
+  state_axis_resources = FewShotState(  # pytype: disable=wrong-arg-types
       rngs={key: PartitionSpec() for key in rng_keys})
   fewshot_step_pjit = jax.experimental.pjit.pjit(
       functools.partial(_fewshot_step, apply_fn=apply_fn),

vmoe/train/train_state.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 
 """TrainState and other related classes."""
-from typing import Any, Callable, Dict, Mapping, Tuple, Union
+from typing import Any, Callable, Dict, Mapping, Tuple
 
 import flax.training.train_state
 import jax
 import optax
 
-PRNGKey = Union[jax.numpy.ndarray, jax.random.KeyArray]
+PRNGKey = jax.Array
 
 
 class TrainState(flax.training.train_state.TrainState):

vmoe/train/trainer.py

@@ -62,7 +62,7 @@
 Mesh = partitioning.Mesh
 NamedSharding = jax.sharding.NamedSharding
 PartitionSpec = partitioning.PartitionSpec
-PRNGKey = Union[jax.numpy.ndarray, jax.random.KeyArray]
+PRNGKey = jax.Array
 PyTree = Any
 ReportProgress = train_periodic_actions.ReportProgress
 SingleProcessPeriodicAction = train_periodic_actions.SingleProcessPeriodicAction

vmoe/utils.py

@@ -21,7 +21,7 @@
 import jax
 import jax.numpy as jnp
 
-PRNGKey = jax.random.KeyArray
+PRNGKey = jax.Array
 
 
 def make_rngs(rng_keys: Tuple[str, ...], seed: int) -> Dict[str, PRNGKey]: