[JAX] Replace uses of jax.interpreters.xla.DeviceArray with jax.Array.

PiperOrigin-RevId: 508609103
google · Feb 27, 2023 · a9a8fc0 · a9a8fc0
1 parent 156a760
commit a9a8fc0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dopamine/jax/networks.py b/dopamine/jax/networks.py
@@ -348,7 +348,7 @@ class NoisyNetwork(nn.Module):
     rng_key: jax.interpreters.xla.DeviceArray, key for JAX RNG.
     eval_mode: bool, whether to turn off noise during evaluation.
   """
-  rng_key: jax.interpreters.xla.DeviceArray
+  rng_key: jax.Array
   eval_mode: bool = False
 
   @staticmethod