fix: changing all PRNGKey to key because its going to be deprecated soon

algorithmic_efficiency/spec.py

@@ -79,7 +79,7 @@ def __eq__(self, other):
                            nn.Module]
 ParameterTypeTree = Dict[ParameterKey, Dict[ParameterKey, ParameterType]]
 
-RandomState = Any  # Union[jax.random.PRNGKey, int, bytes, ...]
+RandomState = Any  # Union[jax.random.key, int, bytes, ...]
 
 OptimizerState = Union[Dict[str, Any], Tuple[Any, Any]]
 Hyperparameters = Any

algorithmic_efficiency/workloads/fastmri/input_pipeline.py

@@ -204,7 +204,7 @@ def process_example(example_index, example):
           process_rng, example_index)
     else:
       # NOTE(dsuo): we use fixed randomness for eval.
-      process_rng = tf.cast(jax.random.PRNGKey(_EVAL_SEED), tf.int64)
+      process_rng = tf.cast(jax.random.key(_EVAL_SEED), tf.int64)
     return _process_example(*example, process_rng)
 
   ds = ds.enumerate().map(process_example, num_parallel_calls=16)

algorithmic_efficiency/workloads/mnist/workload.py

@@ -23,7 +23,7 @@ def _normalize(image: spec.Tensor, mean: float, stddev: float) -> spec.Tensor:
 
 
 def _build_mnist_dataset(
-    data_rng: jax.random.PRNGKey,
+    data_rng: jax.random.key,
     num_train_examples: int,
     num_validation_examples: int,
     train_mean: float,

algorithmic_efficiency/workloads/ogbg/ogbg_pytorch/workload.py

@@ -69,7 +69,7 @@ def loss_fn(
     return loss_dict
 
   def _build_input_queue(self,
-                         data_rng: jax.random.PRNGKey,
+                         data_rng: jax.random.key,
                          split: str,
                          data_dir: str,
                          global_batch_size: int):

algorithmic_efficiency/workloads/ogbg/workload.py

@@ -95,7 +95,7 @@ def eval_period_time_sec(self) -> int:
     return 4 * 60
 
   def _build_input_queue(self,
-                         data_rng: jax.random.PRNGKey,
+                         data_rng: jax.random.key,
                          split: str,
                          data_dir: str,
                          global_batch_size: int):

algorithmic_efficiency/workloads/utils.py

@@ -6,7 +6,7 @@ def print_jax_model_summary(model, fake_inputs):
   """Prints a summary of the jax module."""
   tabulate_fn = nn.tabulate(
       model,
-      jax.random.PRNGKey(0),
+      jax.random.key(0),
       console_kwargs={
           'force_terminal': False, 'force_jupyter': False, 'width': 240
       },

algorithmic_efficiency/workloads/wmt/wmt_jax/workload.py

@@ -105,7 +105,7 @@ def initialize_cache(self,
     config = models.TransformerConfig(deterministic=True, decode=True)
     target_shape = (inputs.shape[0], max_decode_len) + inputs.shape[2:]
     initial_variables = models.Transformer(config).init(
-        jax.random.PRNGKey(0),
+        jax.random.key(0),
         jnp.ones(inputs.shape, jnp.float32),
         jnp.ones(target_shape, jnp.float32))
     return initial_variables['cache']

algorithmic_efficiency/workloads/wmt/wmt_pytorch/workload.py

@@ -238,7 +238,7 @@ def model_fn(
     return logits_batch, None
 
   def _build_input_queue(self,
-                         data_rng: jax.random.PRNGKey,
+                         data_rng: jax.random.key,
                          split: str,
                          data_dir: str,
                          global_batch_size: int,

algorithmic_efficiency/workloads/wmt/workload.py

@@ -116,7 +116,7 @@ def glu(self) -> bool:
     return False
 
   def _build_input_queue(self,
-                         data_rng: jax.random.PRNGKey,
+                         data_rng: jax.random.key,
                          split: str,
                          data_dir: str,
                          global_batch_size: int,

docker/scripts/check_gpu.py

@@ -2,7 +2,7 @@
 
 print('JAX identified %d GPU devices' % jax.local_device_count())
 print('Generating RNG seed for CUDA sanity check ... ')
-rng = jax.random.PRNGKey(0)
+rng = jax.random.key(0)
 data_rng, shuffle_rng = jax.random.split(rng, 2)
 
 if jax.local_device_count() == 8 and data_rng is not None:

tests/modeldiffs/criteo1tb/compare.py

@@ -71,7 +71,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/criteo1tb_embed_init/compare.py

@@ -70,7 +70,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/criteo1tb_layernorm/compare.py

@@ -82,7 +82,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       # mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/criteo1tb_resnet/compare.py

@@ -89,7 +89,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/diff.py

@@ -13,7 +13,7 @@ def torch2jax(jax_workload,
               key_transform=None,
               sd_transform=None,
               init_kwargs=dict(dropout_rate=0.0, aux_dropout_rate=0.0)):
-  jax_params, model_state = jax_workload.init_model_fn(jax.random.PRNGKey(0),
+  jax_params, model_state = jax_workload.init_model_fn(jax.random.key(0),
                                                        **init_kwargs)
   pytorch_model, _ = pytorch_workload.init_model_fn([0], **init_kwargs)
   jax_params = jax_utils.unreplicate(jax_params).unfreeze()

tests/modeldiffs/fastmri/compare.py

@@ -73,7 +73,7 @@ def sort_key(k):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/fastmri_layernorm/compare.py

@@ -80,7 +80,7 @@ def sort_key(k):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/fastmri_model_size/compare.py

@@ -73,7 +73,7 @@ def sort_key(k):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/fastmri_tanh/compare.py

@@ -73,7 +73,7 @@ def sort_key(k):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/imagenet_resnet/compare.py

@@ -90,7 +90,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/imagenet_resnet/gelu_compare.py

@@ -37,7 +37,7 @@
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/imagenet_resnet/silu_compare.py

@@ -37,7 +37,7 @@
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/imagenet_vit/compare.py

@@ -103,7 +103,7 @@ def key_transform(k):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/imagenet_vit_glu/compare.py

@@ -39,7 +39,7 @@
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/imagenet_vit_map/compare.py

@@ -50,7 +50,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/imagenet_vit_postln/compare.py

@@ -39,7 +39,7 @@
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/librispeech_conformer/compare.py

@@ -78,7 +78,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/librispeech_conformer_attention_temperature/compare.py

@@ -78,7 +78,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/librispeech_conformer_gelu/compare.py

@@ -78,7 +78,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/librispeech_conformer_layernorm/compare.py

@@ -78,7 +78,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/librispeech_deepspeech/compare.py

@@ -103,7 +103,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/librispeech_deepspeech_noresnet/compare.py

@@ -39,7 +39,7 @@
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/librispeech_deepspeech_normaug/compare.py

@@ -39,7 +39,7 @@
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/librispeech_deepspeech_tanh/compare.py

@@ -39,7 +39,7 @@
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/ogbg/compare.py

@@ -107,7 +107,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/ogbg_gelu/compare.py

@@ -107,7 +107,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/ogbg_model_size/compare.py

@@ -107,7 +107,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/ogbg_silu/compare.py

@@ -107,7 +107,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/wmt/compare.py

@@ -128,7 +128,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/wmt_attention_temp/compare.py

@@ -128,7 +128,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/wmt_glu_tanh/compare.py

@@ -128,7 +128,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/modeldiffs/wmt_post_ln/compare.py

@@ -128,7 +128,7 @@ def sd_transform(sd):
   jax_model_kwargs = dict(
       augmented_and_preprocessed_input_batch=jax_batch,
       mode=spec.ForwardPassMode.EVAL,
-      rng=jax.random.PRNGKey(0),
+      rng=jax.random.key(0),
       update_batch_norm=False)
 
   out_diff(

tests/test_param_shapes.py

@@ -116,7 +116,7 @@ def get_workload(workload):
     pytorch_workload = PyTorchWmtWorkload()
   else:
     raise ValueError(f'Workload {workload} is not available.')
-  _ = jax_workload.init_model_fn(jax.random.PRNGKey(0))
+  _ = jax_workload.init_model_fn(jax.random.key(0))
   _ = pytorch_workload.init_model_fn([0])
   return jax_workload, pytorch_workload
 

tests/test_param_types.py

@@ -220,7 +220,7 @@ def get_workload(workload_name):
     pytorch_workload = PyTorchWmtWorkload()
   else:
     raise ValueError(f'Workload {workload_name} is not available.')
-  _ = jax_workload.init_model_fn(jax.random.PRNGKey(0))
+  _ = jax_workload.init_model_fn(jax.random.key(0))
   _ = pytorch_workload.init_model_fn([0])
   return jax_workload, pytorch_workload