[JAX] 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`. Change uses of `jnp.array` to `jnp.ndarray`.

PiperOrigin-RevId: 555234490

scenic/projects/mbt/model.py

@@ -570,10 +570,12 @@ def init_from_train_state(self,
                                                    restored_model_cfg,
                                                    restore_output_proj)
 
-  def loss_function(self,
-                    logits: jnp.array,
-                    batch: base_model.Batch,
-                    model_params: Optional[jnp.array] = None) -> float:
+  def loss_function(
+      self,
+      logits: jnp.ndarray,
+      batch: base_model.Batch,
+      model_params: Optional[jnp.ndarray] = None,
+  ) -> float:
     """Returns sigmoid cross entropy loss with an L2 penalty on the weights.
 
     Args: