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: 559400993
Change-Id: Ib2629da9b8794131f0dff1f3e149ba4896bea527

lightweight_mmm/models.py

@@ -290,8 +290,8 @@ def media_mix_model(
     custom_priors: MutableMapping[str, Prior],
     transform_kwargs: Optional[MutableMapping[str, Any]] = None,
     weekday_seasonality: bool = False,
-    extra_features: Optional[jnp.array] = None
-    ) -> None:
+    extra_features: Optional[jnp.ndarray] = None,
+) -> None:
   """Media mix model.
 
   Args: