[LSC] change uses of jax.random.KeyArray and jax.random.PRNGKeyArray …

…to jax.Array

This change replaces uses of jax.random.KeyArray and jax.random.PRNGKeyArray in the context of type annotations with jax.Array, which is the correct annotation for JAX PRNG keys moving forward.

The purpose of this change is to remove references to KeyArray and PRNGKeyArray, which are deprecated (google/jax#17594) and will soon be removed from JAX. The design and thought process behind this is described in https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html.

Note that KeyArray and PRNGKeyArray have always been aliased to Any, so the new type annotation is far more specific than the old one.

PiperOrigin-RevId: 574195739
Change-Id: I703dbac5824d7497fa6228312a5722b96f0df665

acme/agents/jax/ail/losses_test.py

@@ -64,8 +64,13 @@ def dummy_discriminator(
     prior = .7
     loss_fn = losses.pugail_loss(
         positive_class_prior=prior, entropy_coefficient=0.)
-    loss, _ = loss_fn(dummy_discriminator, {}, one_transition,
-                      zero_transition, ())
+    loss, _ = loss_fn(  # pytype: disable=wrong-arg-types
+        dummy_discriminator,
+        {},
+        one_transition,
+        zero_transition,
+        (),
+    )
 
     d_one = jax.nn.sigmoid(dummy_discriminator({}, one_transition)[0])
     d_zero = jax.nn.sigmoid(dummy_discriminator({}, zero_transition)[0])

acme/jax/types.py

@@ -23,7 +23,7 @@
 import jax
 import jax.numpy as jnp
 
-PRNGKey = jax.random.KeyArray
+PRNGKey = jax.Array
 Networks = TypeVar('Networks')
 """Container for all agent network components."""
 Policy = TypeVar('Policy')

examples/baselines/imitation/run_iqlearn.py

@@ -83,7 +83,7 @@ def demonstration_dataset_factory(
         transitions_iterator, key=random_key, batch_size=batch_size
     )
 
-  return demonstration_dataset_factory
+  return demonstration_dataset_factory  # pytype: disable=bad-return-type
 
 
 def build_experiment_config() -> (
@@ -114,7 +114,7 @@ def build_experiment_config() -> (
 
   # Construct the agent
   iq_learn_config = iq_learn.IQLearnConfig(alpha=1.0)
-  iq_learn_builder = iq_learn.IQLearnBuilder(
+  iq_learn_builder = iq_learn.IQLearnBuilder(  # pytype: disable=wrong-arg-types
       config=iq_learn_config, make_demonstrations=make_demonstrations
   )