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: 559395727

haiku/_src/module.py

@@ -661,15 +661,15 @@ def __init__(self, name: Optional[str] = None):
   # Support @dataclass annotated modules.
   __post_init__ = __init__
 
-  def params_dict(self) -> Mapping[str, jnp.array]:
+  def params_dict(self) -> Mapping[str, jnp.ndarray]:
     """Returns parameters keyed by name for this module and submodules."""
     if not base.frame_stack:
       raise ValueError(
           "`module.params_dict()` must be used as part of an `hk.transform`.")
 
     return params_or_state_dict(self.module_name, self._submodules, "params")
 
-  def state_dict(self) -> Mapping[str, jnp.array]:
+  def state_dict(self) -> Mapping[str, jnp.ndarray]:
     """Returns state keyed by name for this module and submodules."""
     if not base.frame_stack:
       raise ValueError(
@@ -682,7 +682,7 @@ def params_or_state_dict(
     module_name: str,
     submodules: Set[str],
     which: str,
-) -> Mapping[str, jnp.array]:
+) -> Mapping[str, jnp.ndarray]:
   """Returns module parameters or state for the given module or submodules."""
   assert which in ("params", "state")
   out = {}