Replace uses of jnp.array in types with jnp.ndarray.

`jnp.array` is a function, not a type:
https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.array.html
so it never makes sense to use `jnp.array` in a type annotation. Presumably the intent was to write `jnp.ndarray` aka `jax.Array`.

PiperOrigin-RevId: 559406793

scenic/model_lib/layers/masked_layers.py

@@ -331,7 +331,7 @@ class GroupNorm(nn.Module):
   def __call__(
       self,
       x: jnp.ndarray,
-      spatial_shape: Optional[jnp.array] = None,
+      spatial_shape: Optional[jnp.ndarray] = None,
   ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]:
     """Applies group normalization to the input (arxiv.org/abs/1803.08494).
 

scenic/projects/mtv/train_utils.py

@@ -145,9 +145,11 @@ def eval_step(
     return_logits_and_labels: bool = False,
     return_confusion_matrix: bool = False,
     debug: Optional[bool] = False
-) -> Union[Tuple[Dict[str, Tuple[float, int]], jnp.ndarray, jnp.array],
-           Tuple[Dict[str, Tuple[float, int]], jnp.ndarray],
-           Dict[str, Tuple[float, int]]]:
+) -> Union[
+    Tuple[Dict[str, Tuple[float, int]], jnp.ndarray, jnp.ndarray],
+    Tuple[Dict[str, Tuple[float, int]], jnp.ndarray],
+    Dict[str, Tuple[float, int]],
+]:
   """Runs a single step of training.
 
   Note that in this code, the buffer of the second argument (batch) is donated
@@ -214,8 +216,10 @@ def test_step(
     return_logits_and_labels: bool = False,
     softmax_logits: bool = False,
     debug: bool = False
-) -> Union[Dict[str, Tuple[float, int]], Tuple[Dict[str, Tuple[float, int]],
-                                               jnp.array, jnp.array]]:
+) -> Union[
+    Dict[str, Tuple[float, int]],
+    Tuple[Dict[str, Tuple[float, int]], jnp.ndarray, jnp.ndarray],
+]:
   """Runs a single step of testing.
 
   For multi-crop testing, we assume that num_crops consecutive entries in the

scenic/projects/objectvivit/trainer.py

@@ -250,8 +250,10 @@ def test_step(
     debug: bool = False,
     learn_token_score: bool = False,
     add_boxes: bool = False,
-) -> Union[Dict[str, Tuple[float, int]], Tuple[Dict[str, Tuple[float, int]],
-                                               jnp.array, jnp.array]]:
+) -> Union[
+    Dict[str, Tuple[float, int]],
+    Tuple[Dict[str, Tuple[float, int]], jnp.ndarray, jnp.ndarray],
+]:
   """Runs a single step of testing.
 
   For multi-crop testing, we assume that num_crops consecutive entries in the

scenic/projects/owl_vit/matching_base_models.py

@@ -58,9 +58,9 @@ def compute_cost_matrix(self, predictions: ArrayDict,
     """
     ...
 
-  def matcher(self,
-              cost: jnp.ndarray,
-              n_cols: Optional[jnp.array] = None) -> jnp.ndarray:
+  def matcher(
+      self, cost: jnp.ndarray, n_cols: Optional[jnp.ndarray] = None
+  ) -> jnp.ndarray:
     """Implements a matching function.
 
     Matching functions match predicted detections against ground truth