Replaces references to jax.numpy.DeviceArray with jax.Array.\n

PiperOrigin-RevId: 515673111 Change-Id: I87d1cbd2bb49aad7b8a4029b20d2b1bf058a573d
willmcgowan · Mar 13, 2023 · ca2b942 · ca2b942
1 parent 59bfea3
commit ca2b942
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/open_spiel/python/examples/bridge_supervised_learning.py b/open_spiel/python/examples/bridge_supervised_learning.py
@@ -129,7 +129,7 @@ def loss(
       params: Params,
       inputs: np.ndarray,
       targets: np.ndarray,
-  ) -> jnp.DeviceArray:
+  ) -> jax.Array:
     """Cross-entropy loss."""
     assert targets.dtype == np.int32
     log_probs = net.apply(params, inputs)
@@ -140,7 +140,7 @@ def accuracy(
       params: Params,
       inputs: np.ndarray,
       targets: np.ndarray,
-  ) -> jnp.DeviceArray:
+  ) -> jax.Array:
     """Classification accuracy."""
     predictions = net.apply(params, inputs)
     return jnp.mean(jnp.argmax(predictions, axis=-1) == targets)

diff --git a/open_spiel/python/examples/hearts_supervised_learning.py b/open_spiel/python/examples/hearts_supervised_learning.py
@@ -129,7 +129,7 @@ def loss(
       params: Params,
       inputs: np.ndarray,
       targets: np.ndarray,
-  ) -> jnp.DeviceArray:
+  ) -> jax.Array:
     """Cross-entropy loss."""
     assert targets.dtype == np.int32
     log_probs = net.apply(params, inputs)
@@ -140,7 +140,7 @@ def accuracy(
       params: Params,
       inputs: np.ndarray,
       targets: np.ndarray,
-  ) -> jnp.DeviceArray:
+  ) -> jax.Array:
     """Classification accuracy."""
     predictions = net.apply(params, inputs)
     return jnp.mean(jnp.argmax(predictions, axis=-1) == targets)

diff --git a/open_spiel/python/examples/meta_cfr/sequential_games/utils.py b/open_spiel/python/examples/meta_cfr/sequential_games/utils.py
@@ -27,7 +27,7 @@
 from open_spiel.python.examples.meta_cfr.sequential_games.typing import Params
 
 
-def get_batched_input(input_list: List[jax.numpy.DeviceArray],
+def get_batched_input(input_list: List[jax.Array],
                       infostate_list: List[InfostateNode],
                       illegal_action_list: List[List[int]], batch_size: int):
   """Returns list of function arguments extended to be consistent with batch size.
@@ -95,7 +95,7 @@ def filter_terminal_infostates(infostates_map: InfostateMapping):
 
 def get_network_output(net_apply: ApplyFn, net_params: Params,
                        net_input: np.ndarray, illegal_actions: List[int],
-                       key: hk.PRNGSequence) -> jax.numpy.DeviceArray:
+                       key: hk.PRNGSequence) -> jax.Array:
   """Returns policy generated as output of model.
 
   Args:
@@ -119,7 +119,7 @@ def get_network_output(net_apply: ApplyFn, net_params: Params,
 def get_network_output_batched(
     net_apply: ApplyFn, net_params: Params, net_input: np.ndarray,
     all_illegal_actions: List[List[int]],
-    key: hk.PRNGSequence) -> List[jax.numpy.DeviceArray]:
+    key: hk.PRNGSequence) -> List[jax.Array]:
   """Returns policy of batched input generated as output of model.
 
   Args: