diff --git a/numpyro/distributions/censored.py b/numpyro/distributions/censored.py
index 0368df806..7e40c458a 100644
--- a/numpyro/distributions/censored.py
+++ b/numpyro/distributions/censored.py
@@ -8,7 +8,7 @@
 import numpy as np
 
 import jax
-from jax import lax
+from jax import Array, lax
 import jax.numpy as jnp
 from jax.typing import ArrayLike
 
@@ -112,7 +112,7 @@ def __init__(
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self.base_dist.expand(self.batch_shape).sample(key, sample_shape)
 
     @constraints.dependent_property(is_discrete=False, event_dim=0)
@@ -120,7 +120,7 @@ def support(self) -> Constraint:
         return self._support
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         dtype = jnp.result_type(value, float)
         minval = 100.0 * jnp.finfo(dtype).tiny
 
@@ -228,7 +228,7 @@ def __init__(
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self.base_dist.expand(self.batch_shape).sample(key, sample_shape)
 
     @constraints.dependent_property(is_discrete=False, event_dim=0)
@@ -236,7 +236,7 @@ def support(self) -> Constraint:
         return self._support
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         dtype = jnp.result_type(value, float)
         eps = jnp.finfo(dtype).eps
 
@@ -363,7 +363,7 @@ def __init__(
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self.base_dist.expand(self.batch_shape).sample(key, sample_shape)
 
     @constraints.dependent_property(is_discrete=False, event_dim=1)
@@ -385,7 +385,7 @@ def _get_censoring_masks(self, value):
         return m_left, m_right, m_int, m_double, m_point
 
     @validate_sample
-    def log_prob(self, value):
+    def log_prob(self, value) -> Array:
         dtype = jnp.result_type(value, float)
         minval = 100.0 * jnp.finfo(dtype).tiny  # for values close to 0
         eps = jnp.finfo(dtype).eps  # otherwise
@@ -441,7 +441,7 @@ def log_prob(self, value):
         logp = jnp.where(m_double, lp_double, logp)
         return logp
 
-    def _validate_sample(self, value: ArrayLike) -> None:
+    def _validate_sample(self, value: ArrayLike) -> Array:
         if value.shape[-1] != 2:
             raise ValueError(
                 f"Expected last dimension of `value` to be 2 (lower, upper), but got shape {value.shape}"
diff --git a/numpyro/distributions/conjugate.py b/numpyro/distributions/conjugate.py
index 304a243a8..1a0958a0e 100644
--- a/numpyro/distributions/conjugate.py
+++ b/numpyro/distributions/conjugate.py
@@ -5,7 +5,7 @@
 from typing import Optional
 
 import jax
-from jax import lax, nn, random
+from jax import Array, lax, nn, random
 import jax.numpy as jnp
 from jax.scipy.special import betainc, betaln, gammaln
 from jax.typing import ArrayLike
@@ -72,7 +72,7 @@ def __init__(
         self._beta = Beta(concentration1, concentration0)
         super(BetaBinomial, self).__init__(batch_shape, validate_args=validate_args)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         key_beta, key_binom = random.split(key)
         probs = self._beta.sample(key_beta, sample_shape)
@@ -81,7 +81,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         return (
             -_log_beta_1(self.total_count - value + 1, value)
             + betaln(
@@ -92,11 +92,11 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return self._beta.mean * self.total_count
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return (
             self._beta.variance
             * self.total_count
@@ -161,7 +161,7 @@ def __init__(
             batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""If :math:`X \sim \mathrm{BetaNegativeBinomial}(\alpha, \beta, n)`, then the sampling
         procedure is:
 
@@ -182,7 +182,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return NegativeBinomialProbs(total_count=self.n, probs=probs).sample(key_nb)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""If :math:`X \sim \mathrm{BetaNegativeBinomial}(\alpha, \beta, n)`, then the log
         probability mass function is:
 
@@ -201,7 +201,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""If :math:`X \sim \mathrm{BetaNegativeBinomial}(\alpha, \beta, n)` and
         :math:`\beta > 1`, then the mean is:
 
@@ -217,7 +217,7 @@ def mean(self) -> ArrayLike:
         )
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""If :math:`X \sim \mathrm{BetaNegativeBinomial}(\alpha, \beta, n)` and
         :math:`\beta > 2`, then the variance is:
 
@@ -289,7 +289,7 @@ def __init__(
             validate_args=validate_args,
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         key_dirichlet, key_multinom = random.split(key)
         probs = self._dirichlet.sample(key_dirichlet, sample_shape)
@@ -300,18 +300,18 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         ).sample(key_multinom)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         alpha = self.concentration
         return _log_beta_1(alpha.sum(-1), value.sum(-1)) - _log_beta_1(
             alpha, value
         ).sum(-1)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return self._dirichlet.mean * jnp.expand_dims(self.total_count, -1)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         n = jnp.expand_dims(self.total_count, -1)
         alpha = self.concentration
         alpha_sum = self.concentration.sum(-1, keepdims=True)
@@ -362,7 +362,7 @@ def __init__(
             self._gamma.batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""If :math:`X \sim \mathrm{GammaPoisson}(\alpha, \lambda)`, then the sampling
         procedure is:
 
@@ -383,7 +383,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return Poisson(rate).sample(key_poisson)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""If :math:`X \sim \mathrm{GammaPoisson}(\alpha, \lambda)`, then the
         probability mass function is:
 
@@ -399,16 +399,16 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""If :math:`X \sim \mathrm{GammaPoisson}(\alpha, \lambda)`, then the mean is:
 
         .. math::
             \mathbb{E}[X] = \frac{\alpha}{\lambda}
         """
-        return self.concentration / self.rate
+        return jnp.asarray(self.concentration / self.rate)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""If :math:`X \sim \mathrm{GammaPoisson}(\alpha, \lambda)`, then the variance is:
 
         .. math::
@@ -416,7 +416,7 @@ def variance(self) -> ArrayLike:
         """
         return self.concentration / jnp.square(self.rate) * (1 + self.rate)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""If :math:`X \sim \mathrm{GammaPoisson}(\alpha, \lambda)`, then the cumulative
         distribution function is:
 
@@ -514,7 +514,7 @@ def __init__(
         super().__init__(concentration, rate, validate_args=validate_args)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""If :math:`X \sim \mathrm{NegativeBinomial}(r, \mathrm{logits}(p))`, then the log
         probability mass function is:
 
diff --git a/numpyro/distributions/constraints.py b/numpyro/distributions/constraints.py
index b71c46863..5fb9f9629 100644
--- a/numpyro/distributions/constraints.py
+++ b/numpyro/distributions/constraints.py
@@ -70,6 +70,7 @@
 import numpy as np
 
 import jax
+from jax import Array
 import jax.numpy as jnp
 from jax.tree_util import register_pytree_node
 from jax.typing import ArrayLike
@@ -699,8 +700,8 @@ def feasible_like(self, prototype: NumLike) -> NumLike:
 class _OrderedVector(_SingletonConstraint[NonScalarArray]):
     event_dim = 1
 
-    def __call__(self, x: NonScalarArray) -> ArrayLike:
-        return (x[..., 1:] > x[..., :-1]).all(axis=-1)
+    def __call__(self, x: NonScalarArray) -> Array:
+        return (x[..., 1:] > x[..., :-1]).all(axis=-1)  # type: ignore
 
     def feasible_like(self, prototype: NonScalarArray) -> NonScalarArray:
         return jnp.broadcast_to(jnp.arange(float(prototype.shape[-1])), prototype.shape)
@@ -757,7 +758,7 @@ class _PositiveOrderedVector(_SingletonConstraint[NonScalarArray]):
 
     event_dim = 1
 
-    def __call__(self, x: NonScalarArray) -> ArrayLike:
+    def __call__(self, x: NonScalarArray) -> Array:
         return jnp.logical_and(
             ordered_vector.check(x), jnp.all(positive.check(x), axis=-1)
         )
@@ -783,9 +784,9 @@ def feasible_like(self, prototype: NumLike) -> NumLike:
 
 
 class _Real(_SingletonConstraint[NumLike]):
-    def __call__(self, x: NumLike) -> ArrayLike:
+    def __call__(self, x: NumLike) -> Array:
         # XXX: consider to relax this condition to [-inf, inf] interval
-        return (x == x) & (x != float("inf")) & (x != float("-inf"))
+        return (x == x) & (x != float("inf")) & (x != float("-inf"))  # type: ignore
 
     def feasible_like(self, prototype: NumLike) -> NumLike:
         return jnp.zeros_like(prototype)
@@ -794,9 +795,9 @@ def feasible_like(self, prototype: NumLike) -> NumLike:
 class _Simplex(_SingletonConstraint[NonScalarArray]):
     event_dim = 1
 
-    def __call__(self, x: NonScalarArray) -> ArrayLike:
+    def __call__(self, x: NonScalarArray) -> Array:
         x_sum = x.sum(axis=-1)
-        return (x >= 0).all(axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)
+        return (x >= 0).all(axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)  # type: ignore
 
     def feasible_like(self, prototype: NonScalarArray) -> NonScalarArray:
         return jnp.full_like(prototype, 1 / prototype.shape[-1])
@@ -855,6 +856,7 @@ def __call__(self, x: NonScalarArray) -> ArrayLike:
         zerosum_true = True
         for dim in range(-self.event_dim, 0):
             zerosum_true = zerosum_true & xp.allclose(x.sum(dim), 0, atol=tol)
+        # FIXME: shape must match batch shape of `x`, not be a literal boolean.
         return zerosum_true
 
     def eq(self, other: object, static: bool = False) -> ArrayLike:
diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
index 57a3d2baa..c50983cfb 100644
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -128,27 +128,27 @@ def left_scale(self):
     def right_scale(self):
         return self.scale / self.asymmetry
 
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         if self._validate_args:
             self._validate_sample(value)
         z = value - self.loc
         z = -jnp.abs(z) / jnp.where(z < 0, self.left_scale, self.right_scale)
         return z - jnp.log(self.left_scale + self.right_scale)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         shape = (2,) + sample_shape + self.batch_shape + self.event_shape
         u, v = random.exponential(key, shape=shape)
         return self.loc - self.left_scale * u + self.right_scale * v
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         total_scale = self.left_scale + self.right_scale
         mean = self.loc + (self.right_scale**2 - self.left_scale**2) / total_scale
         return jnp.broadcast_to(mean, self.batch_shape)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         left = self.left_scale
         right = self.right_scale
         total = left + right
@@ -157,7 +157,7 @@ def variance(self) -> ArrayLike:
         variance = p * left**2 + q * right**2 + p * q * total**2
         return jnp.broadcast_to(variance, self.batch_shape)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         z = value - self.loc
         k = self.asymmetry
         return jnp.where(
@@ -166,7 +166,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
             k**2 / (1 + k**2) * jnp.exp(-jnp.abs(z) / self.left_scale),
         )
 
-    def icdf(self, value: ArrayLike) -> ArrayLike:
+    def icdf(self, value: ArrayLike) -> Array:
         k = self.asymmetry
         temp = k**2 / (1 + k**2)
         return jnp.where(
@@ -224,7 +224,7 @@ def __init__(
             jnp.stack([concentration1, concentration0], axis=-1)
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Generates samples from the distribution using the underlying Dirichlet implementation.
 
         Since a :math:`\mathrm{Beta}(\alpha, \beta)` distribution is equivalent to a
@@ -242,7 +242,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return self._dirichlet.sample(key, sample_shape)[..., 0]
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Calculates the log of the probability density function.
 
         To avoid `NaN` gradients at the boundaries :math:`x=0` or :math:`x=1`, this
@@ -281,23 +281,27 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""Calculates the analytical mean.
 
         .. math:: E[X] = \frac{\alpha}{\alpha + \beta}
         """
-        return self.concentration1 / (self.concentration1 + self.concentration0)
+        return jnp.asarray(
+            self.concentration1 / (self.concentration1 + self.concentration0)
+        )
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""Calculates the analytical variance.
 
         .. math:: Var(X) = \frac{\alpha \beta}{(\alpha + \beta)^2 (\alpha + \beta + 1)}
         """
         total = self.concentration1 + self.concentration0
-        return self.concentration1 * self.concentration0 / (total**2 * (total + 1))
+        return jnp.asarray(
+            self.concentration1 * self.concentration0 / (total**2 * (total + 1))
+        )
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""Cumulative distribution function using the regularized incomplete beta function.
 
         .. math:: I_x(\alpha, \beta) = \frac{\text{B}(x; \alpha, \beta)}{\text{B}(\alpha, \beta)}
@@ -307,7 +311,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         """
         return betainc(self.concentration1, self.concentration0, value)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""Inverse cumulative distribution function (Quantile function).
 
         :param q: Probability value in :math:`[0,1]`.
@@ -315,7 +319,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         """
         return betaincinv(self.concentration1, self.concentration0, q)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""Entropy of the Beta distribution.
 
         .. math::
@@ -372,7 +376,7 @@ def __init__(
             batch_shape=batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Generates samples using the inverse CDF method via :func:`~jax.random.cauchy`.
 
         :param key: JAX PRNGKey for reproducibility.
@@ -387,7 +391,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return self.loc + eps * self.scale
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Calculates the log of the probability density function.
 
         .. math::
@@ -406,7 +410,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""The mean of the Cauchy distribution is undefined.
 
         Returns ``NaN`` for all batch elements.
@@ -414,14 +418,14 @@ def mean(self) -> ArrayLike:
         return jnp.full(self.batch_shape, jnp.nan)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""The variance of the Cauchy distribution is undefined.
 
         Returns ``NaN`` for all batch elements.
         """
         return jnp.full(self.batch_shape, jnp.nan)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""Cumulative distribution function.
 
         .. math::
@@ -434,7 +438,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         scaled = (value - self.loc) / self.scale
         return jnp.arctan(scaled) / jnp.pi + 0.5
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""Inverse cumulative distribution function (Quantile function).
 
         .. math::
@@ -446,7 +450,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         """
         return self.loc + self.scale * jnp.tan(jnp.pi * (q - 0.5))
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""Entropy of the Cauchy distribution.
 
         .. math::
@@ -480,7 +484,7 @@ def __init__(
             validate_args=validate_args,
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         shape = sample_shape + self.batch_shape
         samples = random.dirichlet(key, self.concentration, shape=shape)
@@ -489,7 +493,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         normalize_term = jnp.sum(gammaln(self.concentration), axis=-1) - gammaln(
             jnp.sum(self.concentration, axis=-1)
         )
@@ -499,11 +503,11 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return self.concentration / jnp.sum(self.concentration, axis=-1, keepdims=True)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         con0 = jnp.sum(self.concentration, axis=-1, keepdims=True)
         return self.concentration * (con0 - self.concentration) / (con0**2 * (con0 + 1))
 
@@ -513,7 +517,7 @@ def infer_shapes(concentration):
         event_shape = concentration[-1:]
         return batch_shape, event_shape
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         (n,) = self.event_shape
         total = self.concentration.sum(axis=-1)
         return (
@@ -569,7 +573,7 @@ def __init__(
     def support(self) -> constraints.Constraint:
         return constraints.independent(constraints.real, self.event_dim)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         batch_shape = sample_shape + self.batch_shape
 
@@ -609,7 +613,7 @@ def scan_fn(init, noise, tm1, dt):
         return sde_out
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         sample_shape = lax.broadcast_shapes(
             value.shape[: -self.event_dim], self.batch_shape
         )
@@ -695,7 +699,7 @@ def __init__(
             batch_shape=jnp.shape(rate), validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Generates samples by scaling standard exponential draws by the
         inverse rate: :math:`X = E / \lambda`, where :math:`E \sim \mathrm{Exp}(1)`.
 
@@ -712,7 +716,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Calculates the log of the probability density function.
 
         .. math::
@@ -726,7 +730,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return jnp.log(self.rate) - self.rate * value
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""Calculates the analytical mean.
 
         .. math:: E[X] = \frac{1}{\lambda}
@@ -734,14 +738,14 @@ def mean(self) -> ArrayLike:
         return jnp.reciprocal(self.rate)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""Calculates the analytical variance.
 
         .. math:: \mathrm{Var}(X) = \frac{1}{\lambda^2}
         """
         return jnp.reciprocal(self.rate**2)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""Cumulative distribution function.
 
         .. math::
@@ -752,7 +756,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         """
         return -jnp.expm1(-self.rate * value)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""Inverse cumulative distribution function (Quantile function).
 
         .. math::
@@ -763,7 +767,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         """
         return -jnp.log1p(-q) / self.rate
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""Entropy of the Exponential distribution.
 
         .. math::
@@ -801,7 +805,7 @@ def __init__(
             batch_shape=batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Method to generate samples :math:`X \sim \mathrm{Gamma}(\alpha, \lambda)`.
         It uses :func:`~jax.random.gamma` under the hood to generate samples.
         """
@@ -810,7 +814,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return random.gamma(key, self.concentration, shape=shape) / self.rate
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""If :math:`X \sim \mathrm{Gamma}(\alpha, \lambda)`, then
 
         .. math::
@@ -832,16 +836,16 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""If :math:`X \sim \mathrm{Gamma}(\alpha, \lambda)`, then
 
         .. math::
             \mathbb{E}[X] = \frac{\alpha}{\lambda}
         """
-        return self.concentration / self.rate
+        return jnp.asarray(self.concentration / self.rate)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""If :math:`X \sim \mathrm{Gamma}(\alpha, \lambda)`, then
 
         .. math::
@@ -849,7 +853,7 @@ def variance(self) -> ArrayLike:
         """
         return self.concentration / jnp.power(self.rate, 2)
 
-    def cdf(self, x):
+    def cdf(self, x) -> Array:
         r"""If :math:`X \sim \mathrm{Gamma}(\alpha, \lambda)`, then
 
         .. math::
@@ -862,7 +866,7 @@ def cdf(self, x):
         """
         return gammainc(self.concentration, self.rate * x)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""If :math:`X \sim \mathrm{Gamma}(\alpha, \lambda)`, then
 
         .. math::
@@ -875,7 +879,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         """
         return gammaincinv(self.concentration, q) / self.rate
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""If :math:`X \sim \mathrm{Gamma}(\alpha, \lambda)`, then
 
         .. math::
@@ -1018,7 +1022,7 @@ def initial_value(self) -> Array:
         return self._initial_value
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         # If there's no initial value, the mean is zero (base distribution mean).
         if self._initial_value is None:
             return self.base_dist.mean
@@ -1035,7 +1039,7 @@ def propagate(z, _):
         return jnp.moveaxis(means, 0, -2)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         # Given z_t = z_0 + \sum_{k=1}^t A^{t-k} \epsilon_t, the covariance of the state
         # vector at step t is E[z_t transpose(z_t)] = \sum_{k,k'}^t A^{t-k}
         # E[\epsilon_k transpose(\epsilon_{k'})] transpose(A^{t-k'}). We only have
@@ -1094,14 +1098,14 @@ def __init__(
             batch_shape, event_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         shape = sample_shape + self.batch_shape + self.event_shape
         walks = random.normal(key, shape=shape)
         return jnp.cumsum(walks, axis=-1) * jnp.expand_dims(self.scale, axis=-1)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         init_prob = Normal(0.0, self.scale).log_prob(value[..., 0])
         scale = jnp.expand_dims(self.scale, -1)
         # Normal is location-invariant, so evaluate the increments under a
@@ -1112,11 +1116,11 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return init_prob + jnp.sum(step_probs, axis=-1)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.zeros(self.batch_shape + self.event_shape)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.broadcast_to(
             jnp.expand_dims(self.scale, -1) ** 2 * jnp.arange(1, self.num_steps + 1),
             self.batch_shape + self.event_shape,
@@ -1141,26 +1145,26 @@ def __init__(
             batch_shape=jnp.shape(scale), validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         return jnp.abs(self._cauchy.sample(key, sample_shape))
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         return self._cauchy.log_prob(value) + jnp.log(2)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         return self._cauchy.cdf(value) * 2 - 1
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         return self._cauchy.icdf((q + 1) / 2)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.full(self.batch_shape, jnp.inf)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.full(self.batch_shape, jnp.inf)
 
 
@@ -1187,22 +1191,22 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         return jnp.abs(self._normal.sample(key, sample_shape))
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         return self._normal.log_prob(value) + jnp.log(2)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         return self._normal.cdf(value) * 2 - 1
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         return self._normal.icdf((q + 1) / 2)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.sqrt(2 / jnp.pi) * self.scale
 
     @property
-    def variance(self) -> ArrayLike:
-        return (1 - 2 / jnp.pi) * self.scale**2
+    def variance(self) -> Array:
+        return jnp.asarray((1 - 2 / jnp.pi) * self.scale**2)
 
 
 class InverseGamma(TransformedDistribution):
@@ -1234,18 +1238,18 @@ def __init__(
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         # mean is inf for alpha <= 1
         a = self.rate / (self.concentration - 1)
         return jnp.where(self.concentration <= 1, jnp.inf, a)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         # var is inf for alpha <= 2
         a = (self.rate / (self.concentration - 1)) ** 2 / (self.concentration - 2)
         return jnp.where(self.concentration <= 2, jnp.inf, a)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         return (
             self.concentration
             + jnp.log(self.rate)
@@ -1290,14 +1294,14 @@ def __init__(
             validate_args=validate_args,
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         random_shape = sample_shape + self.batch_shape + self.event_shape
         unifs = random.uniform(key, shape=random_shape)
         return self.icdf(unifs)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         scaled_value = value * self.rate
         return (
             jnp.log(self.concentration)
@@ -1306,14 +1310,14 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
             - self.concentration * jnp.expm1(scaled_value)
         )
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         return -jnp.expm1(-self.concentration * jnp.expm1(value * self.rate))
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         return jnp.log1p(-jnp.log1p(-q) / self.concentration) / self.rate
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return -jnp.exp(self.concentration) * expi(-self.concentration) / self.rate
 
 
@@ -1357,7 +1361,7 @@ def __init__(
             batch_shape=batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples from the Gumbel distribution via the location-scale
         transform :math:`X = \mu + \beta Z`, where
         :math:`Z \sim \mathrm{Gumbel}(0, 1)` is drawn from
@@ -1374,7 +1378,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return self.loc + self.scale * standard_gumbel_sample
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability density function at ``value``.
 
         Letting :math:`z = (x - \mu)/\beta`,
@@ -1389,7 +1393,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return -(z + jnp.exp(-z)) - jnp.log(self.scale)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""Mean of the Gumbel distribution:
 
         .. math::
@@ -1403,7 +1407,7 @@ def mean(self) -> ArrayLike:
         )
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""Variance of the Gumbel distribution:
 
         .. math::
@@ -1411,7 +1415,7 @@ def variance(self) -> ArrayLike:
         """
         return jnp.broadcast_to(jnp.pi**2 / 6.0 * self.scale**2, self.batch_shape)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""Cumulative Distribution Function (CDF) of the Gumbel distribution:
 
         .. math::
@@ -1422,7 +1426,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         """
         return jnp.exp(-jnp.exp((self.loc - value) / self.scale))
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""Inverse CDF (quantile function) of the Gumbel distribution:
 
         .. math::
@@ -1464,7 +1468,7 @@ def __init__(
         )
         super().__init__(batch_shape=batch_shape, validate_args=validate_args)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         finfo = jnp.finfo(jnp.result_type(float))
         u = random.uniform(
@@ -1475,7 +1479,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return jnp.clip(jnp.exp(log_sample), finfo.tiny, 1 - finfo.eps)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         finfo = jnp.finfo(jnp.result_type(float))
         normalize_term = jnp.log(self.concentration0) + jnp.log(self.concentration1)
         value_con1 = jnp.clip(value**self.concentration1, None, 1 - finfo.eps)
@@ -1486,12 +1490,12 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         log_beta = betaln(1 + 1 / self.concentration1, self.concentration0)
         return self.concentration0 * jnp.exp(log_beta)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         log_beta = betaln(1 + 2 / self.concentration1, self.concentration0)
         return self.concentration0 * jnp.exp(log_beta) - jnp.square(self.mean)
 
@@ -1535,7 +1539,7 @@ def __init__(
             batch_shape=batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples via the location-scale transform
         :math:`X = \mu + b Z`, where :math:`Z \sim \mathrm{Laplace}(0, 1)` is
         drawn from :func:`~jax.random.laplace`.
@@ -1551,7 +1555,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return self.loc + eps * self.scale
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability density function at ``value``:
 
         .. math::
@@ -1565,7 +1569,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return -value_scaled - normalize_term
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""Mean of the Laplace distribution:
 
         .. math::
@@ -1574,7 +1578,7 @@ def mean(self) -> ArrayLike:
         return jnp.broadcast_to(self.loc, self.batch_shape)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""Variance of the Laplace distribution:
 
         .. math::
@@ -1582,7 +1586,7 @@ def variance(self) -> ArrayLike:
         """
         return jnp.broadcast_to(2 * self.scale**2, self.batch_shape)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""Cumulative Distribution Function (CDF) of the Laplace distribution.
         Letting :math:`z = (x - \mu)/b`,
 
@@ -1599,7 +1603,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         scaled = (value - self.loc) / self.scale
         return 0.5 - 0.5 * jnp.sign(scaled) * jnp.expm1(-jnp.abs(scaled))
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""Inverse CDF (quantile function) of the Laplace distribution:
 
         .. math::
@@ -1613,7 +1617,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         a = q - 0.5
         return self.loc - self.scale * jnp.sign(a) * jnp.log1p(-2 * jnp.abs(a))
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""Differential entropy of the Laplace distribution:
 
         .. math::
@@ -1693,7 +1697,7 @@ def __init__(
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.broadcast_to(
             jnp.identity(self.dimension),
             self.batch_shape + (self.dimension, self.dimension),
@@ -1855,7 +1859,7 @@ def _onion(self, key: jax.Array, size):
         diag = jnp.ones(cholesky.shape[:-1]).at[..., 1:].set(jnp.sqrt(1 - beta_sample))
         return add_diag(cholesky, diag)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         if self.sample_method == "onion":
             return self._onion(key, sample_shape)
@@ -1863,7 +1867,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
             return self._cvine(key, sample_shape)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         # Note about computing Jacobian of the transformation from Cholesky factor to
         # correlation matrix:
         #
@@ -1927,14 +1931,14 @@ def __init__(
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.exp(self.loc + self.scale**2 / 2)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return (jnp.exp(self.scale**2) - 1) * jnp.exp(2 * self.loc + self.scale**2)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         return (1 + jnp.log(2 * jnp.pi)) / 2 + self.loc + jnp.log(self.scale)
 
 
@@ -1978,7 +1982,7 @@ def __init__(
         batch_shape = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
         super(Logistic, self).__init__(batch_shape, validate_args=validate_args)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples via the location-scale transform
         :math:`X = \mu + s Z`, where :math:`Z \sim \mathrm{Logistic}(0, 1)` is
         drawn from :func:`~jax.random.logistic`.
@@ -1994,7 +1998,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return self.loc + z * self.scale
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability density function at ``value``.
 
         Letting :math:`u = (\mu - x)/s`, the log PDF is
@@ -2014,7 +2018,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return log_exponent - log_denominator
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""Mean of the Logistic distribution:
 
         .. math::
@@ -2023,7 +2027,7 @@ def mean(self) -> ArrayLike:
         return jnp.broadcast_to(self.loc, self.batch_shape)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""Variance of the Logistic distribution:
 
         .. math::
@@ -2032,7 +2036,7 @@ def variance(self) -> ArrayLike:
         var = (self.scale**2) * (jnp.pi**2) / 3
         return jnp.broadcast_to(var, self.batch_shape)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""Cumulative Distribution Function (CDF) of the Logistic distribution.
         Letting :math:`z = (x - \mu)/s`,
 
@@ -2048,7 +2052,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         scaled = (value - self.loc) / self.scale
         return expit(scaled)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""Inverse CDF (quantile function) of the Logistic distribution:
 
         .. math::
@@ -2062,7 +2066,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         """
         return self.loc + self.scale * logit(q)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""Differential entropy of the Logistic distribution:
 
         .. math::
@@ -2095,17 +2099,17 @@ def support(self) -> constraints.Constraint:
         return self._support
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return (self.high - self.low) / jnp.log(self.high / self.low)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return (
             0.5 * (self.high**2 - self.low**2) / jnp.log(self.high / self.low)
             - self.mean**2
         )
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         log_low = jnp.log(self.low)
         log_high = jnp.log(self.high)
         return (log_low + log_high) / 2 + jnp.log(log_high - log_low)
@@ -2216,10 +2220,10 @@ def __init__(
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.broadcast_to(self.loc, self.shape())
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         eps = random.normal(
             key, shape=sample_shape + self.batch_shape + self.event_shape
         )
@@ -2230,7 +2234,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return samples
 
     @validate_sample
-    def log_prob(self, values):
+    def log_prob(self, values) -> Array:
         n, p = self.event_shape
 
         row_log_det = tri_logabsdet(self.scale_tril_row)
@@ -2350,7 +2354,7 @@ def __init__(
             validate_args=validate_args,
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         eps = random.normal(
             key, shape=sample_shape + self.batch_shape + self.event_shape
@@ -2360,7 +2364,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         M = _batch_mahalanobis(self.scale_tril, value - self.loc)
         half_log_det = tri_logabsdet(self.scale_tril)
         normalize_term = half_log_det + 0.5 * self.scale_tril.shape[-1] * jnp.log(
@@ -2380,11 +2384,11 @@ def precision_matrix(self):
         return cho_solve((self.scale_tril, True), identity)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.broadcast_to(self.loc, self.shape())
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.broadcast_to(
             jnp.sum(self.scale_tril**2, axis=-1), self.batch_shape + self.event_shape
         )
@@ -2405,7 +2409,7 @@ def infer_shapes(
                 event_shape = lax.broadcast_shapes(event_shape, matrix[-1:])
                 return batch_shape, event_shape
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         (n,) = self.event_shape
         half_log_det = tri_logabsdet(self.scale_tril)
         return n * (jnp.log(2 * np.pi) + 1) / 2 + half_log_det
@@ -2528,13 +2532,13 @@ def __init__(
                     self.adj_matrix, np.swapaxes(self.adj_matrix, -2, -1)
                 ), "adjacency matrix must be symmetric"
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         # TODO: look into a sparse sampling method
         mvn = MultivariateNormal(self.mean, precision_matrix=self.precision_matrix)
         return mvn.sample(key, sample_shape=sample_shape)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         phi = value - self.loc
         adj_matrix = self.adj_matrix
 
@@ -2581,7 +2585,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return 0.5 * (-n * jnp.log(2 * jnp.pi) + logprec + logdet - logquad)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.broadcast_to(self.loc, self.shape())
 
     @lazy_property
@@ -2667,7 +2671,7 @@ def __init__(
             validate_args=validate_args,
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         key_normal, key_chi2 = random.split(key)
         std_normal = random.normal(
@@ -2681,7 +2685,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         n = self.scale_tril.shape[-1]
         Z = (
             tri_logabsdet(self.scale_tril)
@@ -2707,7 +2711,7 @@ def precision_matrix(self) -> Array:
         return cho_solve((self.scale_tril, True), identity)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         # for df <= 1. should be jnp.nan (keeping jnp.inf for consistency with scipy)
         return jnp.broadcast_to(
             jnp.where(jnp.expand_dims(self.df, -1) <= 1, jnp.inf, self.loc),
@@ -2715,7 +2719,7 @@ def mean(self) -> ArrayLike:
         )
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         df = jnp.expand_dims(self.df, -1)
         var = jnp.power(self.scale_tril, 2).sum(-1) * (df / (df - 2))
         var = jnp.where(df > 2, var, jnp.inf)
@@ -2836,7 +2840,7 @@ def __init__(
 
     @property
     def mean(self) -> Array:
-        return self.loc
+        return jnp.asarray(self.loc)
 
     @lazy_property
     def variance(self) -> Array:
@@ -2877,7 +2881,7 @@ def precision_matrix(self) -> Array:
         inverse_cov_diag = jnp.reciprocal(self.cov_diag)
         return add_diag(-jnp.matmul(jnp.swapaxes(A, -1, -2), A), inverse_cov_diag)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         key_W, key_D = random.split(key)
         batch_shape = sample_shape + self.batch_shape
@@ -2892,7 +2896,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         diff = value - self.loc
         M = _batch_lowrank_mahalanobis(
             self.cov_factor, self.cov_diag, diff, self._capacitance_tril
@@ -2902,7 +2906,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
         return -0.5 * (self.loc.shape[-1] * jnp.log(2 * jnp.pi) + log_det + M)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         log_det = _batch_lowrank_logdet(
             self.cov_factor, self.cov_diag, self._capacitance_tril
         )
@@ -2954,7 +2958,7 @@ def __init__(
             batch_shape=batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Generates samples via the reparameterization trick:
         :math:`X = \mu + \sigma \epsilon`, where :math:`\epsilon \sim \mathcal{N}(0,1)`.
 
@@ -2972,7 +2976,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return self.loc + eps * self.scale
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Calculates the log of the probability density function.
 
         .. math::
@@ -2988,7 +2992,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         value_scaled = (value - self.loc) / self.scale
         return -0.5 * value_scaled**2 - normalize_term
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""Cumulative distribution function.
 
         .. math::
@@ -3004,7 +3008,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         scaled = (value - self.loc) / self.scale
         return ndtr(scaled)
 
-    def log_cdf(self, value: ArrayLike) -> ArrayLike:
+    def log_cdf(self, value: ArrayLike) -> Array:
         r"""Log of the cumulative distribution function. Implementation
         calls :func:`jax.scipy.stats.norm.logcdf`.
 
@@ -3013,7 +3017,7 @@ def log_cdf(self, value: ArrayLike) -> ArrayLike:
         """
         return jax_norm.logcdf(value, loc=self.loc, scale=self.scale)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""Inverse cumulative distribution function (Quantile function).
 
         .. math::
@@ -3029,7 +3033,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         return self.loc + self.scale * ndtri(q)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""Calculates the analytical mean.
 
         .. math:: E[X] = \mu
@@ -3037,14 +3041,14 @@ def mean(self) -> ArrayLike:
         return jnp.broadcast_to(self.loc, self.batch_shape)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""Calculates the analytical variance.
 
         .. math:: \mathrm{Var}(X) = \sigma^2
         """
         return jnp.broadcast_to(self.scale**2, self.batch_shape)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""Entropy of the Normal distribution.
 
         .. math::
@@ -3077,13 +3081,13 @@ def __init__(
         super(Pareto, self).__init__(base_dist, transforms, validate_args=validate_args)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         # mean is inf for alpha <= 1
         a = jnp.divide(self.alpha * self.scale, (self.alpha - 1))
         return jnp.where(self.alpha <= 1, jnp.inf, a)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         # var is inf for alpha <= 2
         a = jnp.divide(
             (self.scale**2) * self.alpha, (self.alpha - 1) ** 2 * (self.alpha - 2)
@@ -3095,7 +3099,7 @@ def variance(self) -> ArrayLike:
     def support(self) -> constraints.Constraint:
         return constraints.greater_than(self.scale)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         return jnp.log(self.scale / self.alpha) + 1 + 1 / self.alpha
 
 
@@ -3160,11 +3164,11 @@ def __init__(
         super().__init__(batch_shape=batch_shape, validate_args=validate_args)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         z = (value - self.loc) / self.scale
         return jnp.log(2 / jnp.pi) - jnp.log(self.scale) - jnp.logaddexp(z, -z)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         dtype = jnp.result_type(float)
         finfo = jnp.finfo(dtype)
@@ -3173,20 +3177,20 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return self.icdf(u)
 
     # TODO: refactor validate_sample to only does validation check and use it here
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         z = (value - self.loc) / self.scale
         return jnp.arctan(jnp.exp(z)) * (2 / jnp.pi)
 
-    def icdf(self, value: ArrayLike) -> ArrayLike:
+    def icdf(self, value: ArrayLike) -> Array:
         return jnp.log(jnp.tan(value * (jnp.pi / 2))) * self.scale + self.loc
 
     @property
-    def mean(self) -> ArrayLike:
-        return self.loc
+    def mean(self) -> Array:
+        return jnp.asarray(self.loc)
 
     @property
-    def variance(self) -> ArrayLike:
-        return (jnp.pi / 2 * self.scale) ** 2
+    def variance(self) -> Array:
+        return jnp.asarray((jnp.pi / 2 * self.scale) ** 2)
 
 
 class StudentT(Distribution):
@@ -3217,7 +3221,7 @@ def __init__(
         self._chi2 = Chi2(df)
         super(StudentT, self).__init__(batch_shape, validate_args=validate_args)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         key_normal, key_chi2 = random.split(key)
         std_normal = random.normal(key_normal, shape=sample_shape + self.batch_shape)
@@ -3226,7 +3230,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return self.loc + self.scale * y
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         y = (value - self.loc) / self.scale
         z = (
             jnp.log(self.scale)
@@ -3238,21 +3242,21 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return -0.5 * (self.df + 1.0) * jnp.log1p(y**2.0 / self.df) - z
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         # for df <= 1. should be jnp.nan (keeping jnp.inf for consistency with scipy)
         return jnp.broadcast_to(
             jnp.where(self.df <= 1, jnp.inf, self.loc), self.batch_shape
         )
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         var = jnp.where(
             self.df > 2, jnp.divide(self.scale**2 * self.df, self.df - 2.0), jnp.inf
         )
         var = jnp.where(self.df <= 1, jnp.nan, var)
         return jnp.broadcast_to(var, self.batch_shape)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         # Ref: https://en.wikipedia.org/wiki/Student's_t-distribution#Related_distributions
         # X^2 ~ F(1, df) -> df / (df + X^2) ~ Beta(df/2, 0.5)
         scaled = (value - self.loc) / self.scale
@@ -3267,13 +3271,13 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
             - jnp.sign(scaled) * betainc(0.5 * self.df, 0.5, beta_value)
         )
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         beta_value = betaincinv(0.5 * self.df, 0.5, 1 - jnp.abs(1 - 2 * q))
         scaled_squared = self.df * (1 / beta_value - 1)
         scaled = jnp.sign(q - 0.5) * jnp.sqrt(scaled_squared)
         return scaled * self.scale + self.loc
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         return jnp.broadcast_to(
             (self.df + 1) / 2 * (digamma((self.df + 1) / 2) - digamma(self.df / 2))
             + jnp.log(self.df) / 2
@@ -3307,29 +3311,29 @@ def __init__(
     def support(self) -> constraints.Constraint:
         return self._support
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         shape = sample_shape + self.batch_shape
         return random.uniform(key, shape=shape, minval=self.low, maxval=self.high)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         shape = lax.broadcast_shapes(jnp.shape(value), self.batch_shape)
         return -jnp.broadcast_to(jnp.log(self.high - self.low), shape)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         cdf = (value - self.low) / (self.high - self.low)
         return jnp.clip(cdf, 0.0, 1.0)
 
-    def icdf(self, value: ArrayLike) -> ArrayLike:
-        return self.low + value * (self.high - self.low)
+    def icdf(self, value: ArrayLike) -> Array:
+        return jnp.asarray(self.low + value * (self.high - self.low))
 
     @property
-    def mean(self) -> ArrayLike:
-        return self.low + (self.high - self.low) / 2.0
+    def mean(self) -> Array:
+        return jnp.asarray(self.low + (self.high - self.low) / 2.0)
 
     @property
-    def variance(self) -> ArrayLike:
-        return (self.high - self.low) ** 2 / 12.0
+    def variance(self) -> Array:
+        return jnp.asarray((self.high - self.low) ** 2 / 12.0)
 
     @staticmethod
     def infer_shapes(
@@ -3339,7 +3343,7 @@ def infer_shapes(
         event_shape: tuple[int, ...] = ()
         return batch_shape, event_shape
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         return jnp.log(self.high - self.low)
 
 
@@ -3362,7 +3366,7 @@ def __init__(
         batch_shape = lax.broadcast_shapes(jnp.shape(concentration), jnp.shape(scale))
         super().__init__(batch_shape=batch_shape, validate_args=validate_args)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         return random.weibull_min(
             key,
@@ -3372,28 +3376,28 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         ll = -jnp.power(value / self.scale, self.concentration)
         ll += jnp.log(self.concentration)
         ll += (self.concentration - 1.0) * jnp.log(value)
         ll -= self.concentration * jnp.log(self.scale)
         return ll
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         return 1 - jnp.exp(-((value / self.scale) ** self.concentration))
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return self.scale * jnp.exp(gammaln(1.0 + 1.0 / self.concentration))
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return self.scale**2 * (
             jnp.exp(gammaln(1.0 + 2.0 / self.concentration))
             - jnp.exp(gammaln(1.0 + 1.0 / self.concentration)) ** 2
         )
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         return (
             jnp.euler_gamma * (1 - 1 / self.concentration)
             + jnp.log(self.scale / self.concentration)
@@ -3501,26 +3505,26 @@ def __init__(
         scale_classic = scale * asymmetry / quantile
         self._ald = AsymmetricLaplace(loc=loc, scale=scale_classic, asymmetry=asymmetry)
 
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         if self._validate_args:
             self._validate_sample(value)
         return self._ald.log_prob(value)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         return self._ald.sample(key, sample_shape=sample_shape)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return self._ald.mean
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return self._ald.variance
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         return self._ald.cdf(value)
 
-    def icdf(self, value: ArrayLike) -> ArrayLike:
+    def icdf(self, value: ArrayLike) -> Array:
         return self._ald.icdf(value)
 
 
@@ -3610,11 +3614,11 @@ def support(self) -> constraints.Constraint:
         return constraints.zero_sum(len(self.event_shape))
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.zeros(self.batch_shape + self.event_shape)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         event_ndim = len(self.event_shape)
         zero_sum_axes = tuple(range(-event_ndim, 0))
         theoretical_var = jnp.square(self.scale)
@@ -3688,11 +3692,11 @@ def scale_tril(self):
         return self.base_dist.scale_tril
 
     @lazy_property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return self.concentration[..., None, None] * self.scale_matrix
 
     @lazy_property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         diag = jnp.diagonal(self.scale_matrix, axis1=-1, axis2=-2)
         return self.concentration[..., None, None] * (
             self.scale_matrix**2 + diag[..., :, None] * diag[..., None, :]
@@ -3706,7 +3710,7 @@ def infer_shapes(
             concentration, scale_matrix, rate_matrix, scale_tril
         )
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         p = self.event_shape[-1]
         return (
             (p + 1) * tri_logabsdet(self.scale_tril)
@@ -3783,7 +3787,7 @@ def __init__(
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         # The log density of the Wishart distribution includes a term
         # t = trace(rate_matrix @ cov). Here, value = cholesky(cov) such that
         # t = trace(value.T @ rate_matrix @ value) by the cyclical property of the
@@ -3820,7 +3824,7 @@ def rate_matrix(self):
         )
         return cho_solve((self.scale_tril, True), identity)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         # Sample using the Bartlett decomposition
         # (https://en.wikipedia.org/wiki/Wishart_distribution#Bartlett_decomposition).
@@ -3843,7 +3847,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return jnp.matmul(*jnp.broadcast_arrays(self.scale_tril, latent))
 
     @lazy_property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         # The mean follows from the Bartlett decomposition sampling. All off-diagonal
         # elements of the latent variable have zero expectation. The diagonal are the
         # expected square roots of chi^2 variables which can be expressed in terms of
@@ -3854,7 +3858,7 @@ def mean(self) -> ArrayLike:
         return self.scale_tril * sqrtchi2[..., None, :]
 
     @lazy_property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         # We have the same as for the mean except now the lower off-diagonals are one
         # due to the standard normal noise, and the diagonals are equal to the dof of
         # the chi^2 variables.
@@ -3970,7 +3974,7 @@ def scale_tril(self):
         return self.base_dist.scale_tril
 
     @lazy_property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         # Mean exists only when concentration > p + 1
         p = self.scale_matrix.shape[-1]
         return jnp.where(
@@ -3980,12 +3984,12 @@ def mean(self) -> ArrayLike:
         )
 
     @lazy_property
-    def mode(self) -> ArrayLike:
+    def mode(self) -> Array:
         p = self.scale_matrix.shape[-1]
         return self.scale_matrix / (self.concentration[..., None, None] + p + 1)
 
     @lazy_property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         # Variance of entry (i,j) for nu > p + 3
         # Var(X_ij) = (Psi_ij^2 + Psi_ii * Psi_jj) / ((nu - p - 1)^2 * (nu - p - 3))
         p = self.scale_matrix.shape[-1]
@@ -4078,7 +4082,7 @@ def __init__(
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         # L = value (Cholesky factor), X = L @ L^T ~ InverseWishart(Psi, nu)
         # log p(X) = (nu/2) log|Psi| - (nu*p/2) log(2) - log Gamma_p(nu/2)
         #            - ((nu+p+1)/2) log|X| - tr(Psi @ X^{-1}) / 2
@@ -4110,7 +4114,7 @@ def rate_matrix(self):
         )
         return cho_solve((self.scale_tril, True), identity)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         # Sample from standard InverseWishartCholesky using Bartlett decomposition
         # Ref: https://nbviewer.org/gist/fehiepsi/5ef8e09e61604f10607380467eb82006#Precision-to-scale_tril
@@ -4140,7 +4144,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return jnp.matmul(self.scale_tril, L_inv_std)
 
     @lazy_property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         # Approximate: chol(E[X]) where E[X] = Psi / (nu - p - 1) for nu > p + 1
         p = self.scale_tril.shape[-1]
         mean_x = jnp.where(
@@ -4157,7 +4161,7 @@ def mean(self) -> ArrayLike:
         )
 
     @lazy_property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         # Variance of Cholesky factor is complex; return NaN for now
         return jnp.full(self.batch_shape + self.event_shape, jnp.nan)
 
@@ -4215,7 +4219,7 @@ def support(self) -> constraints.Constraint:
         return self._support
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Compute the log probability density function of the Lévy distribution.
 
         .. math::
@@ -4231,12 +4235,12 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
             jnp.log(2.0 * jnp.pi) - jnp.log(self.scale) + self.scale / shifted_value
         ) - 1.5 * jnp.log(shifted_value)
 
-    def sample(self, key: ArrayLike, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: ArrayLike, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         u = random.uniform(key, shape=sample_shape + self.batch_shape)
         return self.icdf(u)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""
         The inverse cumulative distribution function of Lévy distribution is given by,
 
@@ -4250,7 +4254,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         """
         return self.loc + self.scale * jnp.power(ndtri(1 - 0.5 * q), -2)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""The cumulative distribution function of Lévy distribution is given by,
 
         .. math::
@@ -4265,14 +4269,14 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         return 2.0 - 2.0 * ndtr(jnp.sqrt(inv_standardized))
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.broadcast_to(jnp.inf, self.batch_shape)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.broadcast_to(jnp.inf, self.batch_shape)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""If :math:`X \sim \text{Levy}(\mu, c)`, then the entropy of :math:`X` is given by,
 
         .. math::
@@ -4385,15 +4389,15 @@ def __init__(
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.broadcast_to(self.loc, self.shape())
 
     @lazy_property
-    def covariance_row(self) -> ArrayLike:
+    def covariance_row(self) -> Array:
         return jnp.fft.irfft(self.covariance_rfft, n=self.event_shape[-1])
 
     @lazy_property
-    def covariance_matrix(self) -> ArrayLike:
+    def covariance_matrix(self) -> Array:
         *leading_shape, n = self.covariance_row.shape
         if leading_shape:
             # `toeplitz` flattens the input, and we need to broadcast manually.
@@ -4405,7 +4409,7 @@ def covariance_matrix(self) -> ArrayLike:
             return toeplitz(self.covariance_row)
 
     @lazy_property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.broadcast_to(self.covariance_row[..., 0, None], self.shape())
 
     @staticmethod
@@ -4421,7 +4425,7 @@ def infer_shapes(
                 event_shape = loc[-1:]
                 return batch_shape, event_shape
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         (n,) = self.event_shape
         log_abs_det_jacobian = 2 * jnp.log(2) * ((n - 1) // 2) - jnp.log(n) * n
         return self.base_dist.entropy() + log_abs_det_jacobian / 2
@@ -4468,7 +4472,7 @@ def __init__(
         super().__init__(batch_shape=batch_shape, validate_args=validate_args)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         a_ln_x_m_ln_b = xlogy(self.sharpness, value) - xlogy(self.sharpness, self.scale)
         return (
             jnp.log(self.sharpness)
@@ -4478,7 +4482,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
             - (self.concentration + 1.0) * nn.softplus(a_ln_x_m_ln_b)
         )
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         return jnp.exp(
             -self.concentration
             * nn.softplus(
@@ -4486,7 +4490,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
             )
         )
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         q_root_p = jnp.power(q, -jnp.reciprocal(self.concentration))
         return self.scale * jnp.power(q_root_p - 1.0, -jnp.reciprocal(self.sharpness))
 
@@ -4495,7 +4499,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> jnp.ndar
         return self.icdf(random.uniform(key, shape=self.shape(sample_shape)))
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         safe_a = jnp.where(self.sharpness > 1.0, self.sharpness, 2.0)
         return jnp.where(
             self.sharpness > 1.0,
@@ -4505,7 +4509,7 @@ def mean(self) -> ArrayLike:
         )
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         safe_a = jnp.where(self.sharpness > 2.0, self.sharpness, 3.0)
         return jnp.where(
             self.sharpness > 2.0,
diff --git a/numpyro/distributions/copula.py b/numpyro/distributions/copula.py
index bad6cc028..38d4ade23 100644
--- a/numpyro/distributions/copula.py
+++ b/numpyro/distributions/copula.py
@@ -66,7 +66,7 @@ def __init__(
             validate_args=validate_args,
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
 
         shape = sample_shape + self.batch_shape
@@ -75,7 +75,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return self.marginal_dist.icdf(cdf)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         # Ref: https://en.wikipedia.org/wiki/Copula_(probability_theory)#Gaussian_copula
         # see also https://github.com/pyro-ppl/numpyro/pull/1506#discussion_r1037525015
         marginal_lps = self.marginal_dist.log_prob(value)
@@ -90,11 +90,11 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return copula_lp + marginal_lps.sum(axis=-1)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.broadcast_to(self.marginal_dist.mean, self.shape())
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.broadcast_to(self.marginal_dist.variance, self.shape())
 
     @constraints.dependent_property(is_discrete=False, event_dim=1)
diff --git a/numpyro/distributions/directional.py b/numpyro/distributions/directional.py
index e1b806f34..5c1da2e33 100644
--- a/numpyro/distributions/directional.py
+++ b/numpyro/distributions/directional.py
@@ -123,7 +123,7 @@ def __init__(
             batch_shape=batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         """Generate sample from von Mises distribution
 
         :param key: random number generator key
@@ -140,20 +140,20 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return samples
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         return -(
             jnp.log(2 * jnp.pi) + jnp.log(i0e(self.concentration))
         ) + self.concentration * (jnp.cos((value - self.loc) % (2 * jnp.pi)) - 1)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         """Computes circular mean of distribution. NOTE: same as location when mapped to support [-pi, pi]"""
         return jnp.broadcast_to(
             (self.loc + jnp.pi) % (2.0 * jnp.pi) - jnp.pi, self.batch_shape
         )
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         """Computes circular variance of distribution"""
         return jnp.broadcast_to(
             1.0 - i1e(self.concentration) / i0e(self.concentration), self.batch_shape
@@ -265,7 +265,7 @@ def __repr__(self) -> str:
             + ")"
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         base_key, skew_key = random.split(key)
         bd = self.base_dist
         ys = bd.sample(base_key, sample_shape)
@@ -281,7 +281,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         ) - jnp.pi
         return samples
 
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         if self._validate_args:
             self._validate_sample(value)
         if self.base_dist._validate_args:
@@ -296,7 +296,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return self.base_dist.log_prob(value) + skew_prob
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         """Mean of the base distribution"""
         return self.base_dist.mean
 
@@ -432,7 +432,7 @@ def norm_const(self) -> Array:
         return norm_const.reshape(jnp.shape(self.phi_loc))
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         indv = self.phi_concentration * jnp.cos(
             value[..., 0] - self.phi_loc
         ) + self.psi_concentration * jnp.cos(value[..., 1] - self.psi_loc)
@@ -443,7 +443,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
         return indv + corr - self.norm_const
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         """
         ** References: **
             1. A New Unified Approach for the Simulation of a Wide Class of Directional Distributions
@@ -544,7 +544,7 @@ def cond_fn(curr: PhiMarginalState) -> Array:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         """Computes circular mean of distribution. Note: same as location when mapped to support [-pi, pi]"""
         mean = (jnp.stack((self.phi_loc, self.psi_loc), axis=-1) + jnp.pi) % (
             2.0 * jnp.pi
@@ -596,7 +596,7 @@ def __init__(self, concentration: Array, *, validate_args: Optional[bool] = None
         super().__init__(batch_shape, event_shape, validate_args=validate_args)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         """
         Note this is the mean in the sense of a centroid in the submanifold
         that minimizes expected squared geodesic distance.
@@ -604,15 +604,15 @@ def mean(self) -> ArrayLike:
         return safe_normalize(self.concentration)
 
     @property
-    def mode(self):
+    def mode(self) -> jax.Array:
         return safe_normalize(self.concentration)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         shape = sample_shape + self.batch_shape + self.event_shape
         eps = random.normal(key, shape=shape)
         return safe_normalize(self.concentration + eps)
 
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         if self._validate_args:
             event_shape = value.shape[-1:]
             if event_shape != self.event_shape:
@@ -661,7 +661,7 @@ def _dot(x, y):
 
 
 def _projected_normal_log_prob_3(concentration, value):
-    def _dot(x: Array, y: Array) -> ArrayLike:
+    def _dot(x: Array, y: Array) -> Array:
         return (x[..., None, :] @ y[..., None])[..., 0, 0]
 
     # We integrate along a ray, factorizing the integrand as a product of:
diff --git a/numpyro/distributions/discrete.py b/numpyro/distributions/discrete.py
index 3eb6208a5..30273ba4f 100644
--- a/numpyro/distributions/discrete.py
+++ b/numpyro/distributions/discrete.py
@@ -54,20 +54,20 @@
 from numpyro.util import is_prng_key, not_jax_tracer
 
 
-def _to_probs_bernoulli(logits: ArrayLike) -> ArrayLike:
+def _to_probs_bernoulli(logits: ArrayLike) -> Array:
     return expit(logits)
 
 
-def _to_logits_bernoulli(probs: ArrayLike) -> ArrayLike:
+def _to_logits_bernoulli(probs: ArrayLike) -> Array:
     ps_clamped = clamp_probs(probs)
     return jnp.log(ps_clamped) - jnp.log1p(-ps_clamped)
 
 
-def _to_probs_multinom(logits: ArrayLike) -> ArrayLike:
+def _to_probs_multinom(logits: ArrayLike) -> Array:
     return softmax(logits, axis=-1)
 
 
-def _to_logits_multinom(probs: ArrayLike) -> ArrayLike:
+def _to_logits_multinom(probs: ArrayLike) -> Array:
     minval = jnp.finfo(jnp.result_type(probs)).min
     return jnp.clip(jnp.log(probs), minval)
 
@@ -102,7 +102,7 @@ def __init__(self, probs: ArrayLike, *, validate_args: Optional[bool] = None):
             batch_shape=jnp.shape(self.probs), validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples from the Bernoulli distribution.
 
         This method invokes :func:`~jax.random.bernoulli` directly, which generates
@@ -120,7 +120,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return samples.astype(jnp.result_type(samples, int))
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability mass function at specified binary
         configurations.
 
@@ -144,7 +144,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return xlogy(value, ps_clamped) + xlog1py(1 - value, -ps_clamped)
 
     @lazy_property
-    def logits(self) -> ArrayLike:
+    def logits(self) -> Array:
         r"""The log-odds (logits) parameter of the Bernoulli distribution is given by
         the logit transformation of the success probability:
 
@@ -154,7 +154,7 @@ def logits(self) -> ArrayLike:
         return _to_logits_bernoulli(self.probs)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""The mean of the Bernoulli distribution is given by the success probability
         parameter:
 
@@ -164,10 +164,10 @@ def mean(self) -> ArrayLike:
         :return: The mean of the Bernoulli distribution, which is equal to the success
             probability :attr:`probs`.
         """
-        return self.probs
+        return jnp.asarray(self.probs)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""The variance of the Bernoulli distribution is given by:
 
         .. math::
@@ -176,15 +176,15 @@ def variance(self) -> ArrayLike:
         :return: The variance of the Bernoulli distribution, which is the product of
             the success probability and its complement.
         """
-        return self.probs * (1 - self.probs)
+        return jnp.asarray(self.probs * (1 - self.probs))
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         values = jnp.arange(2).reshape((-1,) + (1,) * len(self.batch_shape))
         if expand:
             values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape)
         return values
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""The entropy of the Bernoulli distribution is given by:
 
         .. math::
@@ -223,7 +223,7 @@ def __init__(self, logits: ArrayLike, *, validate_args: Optional[bool] = None):
             batch_shape=jnp.shape(self.logits), validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples from the Bernoulli distribution.
 
         The method first converts :attr:`logits` to probabilities via the sigmoid
@@ -241,7 +241,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return samples.astype(jnp.result_type(samples, int))
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability mass function at specified binary configurations.
 
         The log probability mass function leverages the numerically-stable
@@ -263,7 +263,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return -binary_cross_entropy_with_logits(self.logits, value)
 
     @lazy_property
-    def probs(self) -> ArrayLike:
+    def probs(self) -> Array:
         r"""The success probability parameter of the Bernoulli distribution is given by
         the sigmoid of the log-odds parameter:
 
@@ -273,7 +273,7 @@ def probs(self) -> ArrayLike:
         return _to_probs_bernoulli(self.logits)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""The mean of the Bernoulli distribution is given by the sigmoid of the
         log-odds parameter:
 
@@ -283,7 +283,7 @@ def mean(self) -> ArrayLike:
         return self.probs
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""The variance of the Bernoulli distribution is given by:
 
         .. math::
@@ -291,13 +291,13 @@ def variance(self) -> ArrayLike:
         """
         return self.probs * (1 - self.probs)
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         values = jnp.arange(2).reshape((-1,) + (1,) * len(self.batch_shape))
         if expand:
             values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape)
         return values
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""The entropy of the Bernoulli distribution is given by:
 
         .. math::
@@ -337,6 +337,8 @@ def Bernoulli(
         return BernoulliProbs(probs, validate_args=validate_args)
     elif logits is not None:
         return BernoulliLogits(logits, validate_args=validate_args)
+    else:
+        raise NotImplementedError
 
 
 class BinomialProbs(Distribution):
@@ -377,7 +379,7 @@ def __init__(
             batch_shape=batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples from the Binomial distribution.
 
         This method uses the internal :func:`~numpyro.distributions.util.binomial`
@@ -393,7 +395,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability mass function at specified count configurations.
 
         The log probability mass function is fully evaluated in log-space to prevent
@@ -430,7 +432,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @lazy_property
-    def logits(self) -> ArrayLike:
+    def logits(self) -> Array:
         r"""The log-odds (logits) parameter of the Binomial distribution is given by
         the logit transformation of the success probability:
 
@@ -440,7 +442,7 @@ def logits(self) -> ArrayLike:
         return _to_logits_bernoulli(self.probs)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""The mean of the Binomial distribution is given by:
 
         .. math::
@@ -449,7 +451,7 @@ def mean(self) -> ArrayLike:
         return jnp.broadcast_to(self.total_count * self.probs, self.batch_shape)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""The variance of the Binomial distribution is given by:
 
         .. math::
@@ -466,7 +468,7 @@ def support(self) -> constraints.Constraint:
         """
         return constraints.integer_interval(0, self.total_count)
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         if not_jax_tracer(self.total_count):
             total_count = np.amax(self.total_count)
             # NB: the error can't be raised if inhomogeneous issue happens when tracing
@@ -522,7 +524,7 @@ def __init__(
             batch_shape=batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples from the Binomial distribution.
 
         The method first converts :attr:`logits` to probabilities via the sigmoid function
@@ -539,7 +541,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability mass function at specified count
         configurations.
 
@@ -576,7 +578,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @lazy_property
-    def probs(self) -> ArrayLike:
+    def probs(self) -> Array:
         r"""The success probability per trial of the Binomial distribution is given by
         the sigmoid of the log-odds parameter:
 
@@ -586,7 +588,7 @@ def probs(self) -> ArrayLike:
         return _to_probs_bernoulli(self.logits)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""The mean of the Binomial distribution is given by:
 
         .. math::
@@ -595,7 +597,7 @@ def mean(self) -> ArrayLike:
         return jnp.broadcast_to(self.total_count * self.probs, self.batch_shape)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""The variance of the Binomial distribution is given by:
 
         .. math::
@@ -672,7 +674,7 @@ def __init__(self, probs: Array, *, validate_args: Optional[bool] = None):
             batch_shape=jnp.shape(self.probs)[:-1], validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples from the Categorical distribution.
 
         This method delegates to :func:`~numpyro.distributions.util.categorical`, which
@@ -688,7 +690,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return categorical(key, self.probs, shape=sample_shape + self.batch_shape)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability mass function at specified category indices.
 
         .. math::
@@ -709,7 +711,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return jnp.take_along_axis(log_pmf, value, axis=-1)[..., 0]
 
     @lazy_property
-    def logits(self) -> ArrayLike:
+    def logits(self) -> Array:
         r"""The log-probability (logits) parameter of the Categorical distribution is
         the (already-normalized) log of the category probabilities:
 
@@ -719,7 +721,7 @@ def logits(self) -> ArrayLike:
         return _to_logits_multinom(self.probs)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""The mean of a Categorical distribution over arbitrary unordered categories
         is not well-defined. This property therefore returns ``NaN``.
 
@@ -728,7 +730,7 @@ def mean(self) -> ArrayLike:
         return jnp.full(self.batch_shape, jnp.nan, dtype=jnp.result_type(self.probs))
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""The variance of a Categorical distribution over arbitrary unordered
         categories is not well-defined. This property therefore returns ``NaN``.
 
@@ -744,7 +746,7 @@ def support(self) -> constraints.Constraint:
         """
         return constraints.integer_interval(0, jnp.shape(self.probs)[-1] - 1)
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         r"""Enumerate all values in the support of the Categorical distribution.
 
         :param expand: Whether to broadcast the enumerated values across the batch
@@ -759,7 +761,7 @@ def enumerate_support(self, expand: bool = True) -> ArrayLike:
             values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape)
         return values
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""The entropy of the Categorical distribution is given by:
 
         .. math::
@@ -805,7 +807,7 @@ def __init__(self, logits: Array, *, validate_args: Optional[bool] = None):
             batch_shape=jnp.shape(logits)[:-1], validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples from the Categorical distribution.
 
         This method invokes :func:`~jax.random.categorical` directly, which samples in
@@ -823,7 +825,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability mass function at specified category indices.
 
         .. math::
@@ -847,7 +849,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return jnp.take_along_axis(log_pmf, value, -1)[..., 0]
 
     @lazy_property
-    def probs(self) -> ArrayLike:
+    def probs(self) -> Array:
         r"""The probability vector of the Categorical distribution is given by the
         softmax of the logits:
 
@@ -858,7 +860,7 @@ def probs(self) -> ArrayLike:
         return _to_probs_multinom(self.logits)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""The mean of a Categorical distribution over arbitrary unordered categories
         is not well-defined. This property therefore returns ``NaN``.
 
@@ -867,7 +869,7 @@ def mean(self) -> ArrayLike:
         return jnp.full(self.batch_shape, jnp.nan, dtype=jnp.result_type(self.logits))
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""The variance of a Categorical distribution over arbitrary unordered
         categories is not well-defined. This property therefore returns ``NaN``.
 
@@ -883,7 +885,7 @@ def support(self) -> constraints.Constraint:
         """
         return constraints.integer_interval(0, jnp.shape(self.logits)[-1] - 1)
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         r"""Enumerate all values in the support of the Categorical distribution.
 
         :param expand: Whether to broadcast the enumerated values across the batch
@@ -898,7 +900,7 @@ def enumerate_support(self, expand: bool = True) -> ArrayLike:
             values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape)
         return values
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""The entropy of the Categorical distribution is given by:
 
         .. math::
@@ -971,7 +973,7 @@ def support(self) -> constraints.Constraint:
         """
         return self._support
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         r"""Draw samples from the discrete uniform distribution.
 
         This method invokes :func:`~jax.random.randint` directly, which generates
@@ -987,7 +989,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return random.randint(key, shape=shape, minval=self.low, maxval=self.high + 1)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Evaluate the log probability mass function at specified integer values.
 
         .. math::
@@ -1003,7 +1005,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         shape = lax.broadcast_shapes(jnp.shape(value), self.batch_shape)
         return -jnp.broadcast_to(jnp.log(self.high + 1 - self.low), shape)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""Evaluate the cumulative distribution function (CDF) of the discrete
         uniform distribution.
 
@@ -1017,7 +1019,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         cdf = (jnp.floor(value) + 1 - self.low) / (self.high - self.low + 1)
         return jnp.clip(cdf, 0.0, 1.0)
 
-    def icdf(self, value: ArrayLike) -> ArrayLike:
+    def icdf(self, value: ArrayLike) -> Array:
         r"""Evaluate the inverse cumulative distribution function (quantile function)
         of the discrete uniform distribution.
 
@@ -1027,10 +1029,10 @@ def icdf(self, value: ArrayLike) -> ArrayLike:
         :param value: Quantile level(s) :math:`u \in [0, 1]`.
         :return: The inverse CDF evaluated at ``value``.
         """
-        return self.low + value * (self.high - self.low + 1) - 1
+        return jnp.asarray(self.low + value * (self.high - self.low + 1) - 1)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         r"""The mean of the discrete uniform distribution is the midpoint of the
         support:
 
@@ -1039,10 +1041,10 @@ def mean(self) -> ArrayLike:
 
         :return: The mean of the discrete uniform distribution.
         """
-        return self.low + (self.high - self.low) / 2.0
+        return jnp.asarray(self.low + (self.high - self.low) / 2.0)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         r"""The variance of the discrete uniform distribution is given by:
 
         .. math::
@@ -1050,9 +1052,9 @@ def variance(self) -> ArrayLike:
 
         :return: The variance of the discrete uniform distribution.
         """
-        return ((self.high - self.low + 1) ** 2 - 1) / 12.0
+        return jnp.asarray(((self.high - self.low + 1) ** 2 - 1) / 12.0)
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         r"""Enumerate all values in the support of the discrete uniform distribution.
 
         Both :attr:`low` and :attr:`high` must be concrete (non-JAX-tracer) values and
@@ -1084,7 +1086,7 @@ def enumerate_support(self, expand: bool = True) -> ArrayLike:
             values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape)
         return values
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         r"""The entropy of the discrete uniform distribution is given by:
 
         .. math::
@@ -1136,7 +1138,7 @@ def infer_shapes(predictor, cutpoints):
         event_shape = ()
         return batch_shape, event_shape
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         raise NotImplementedError
 
 
@@ -1170,7 +1172,7 @@ def __init__(
             validate_args=validate_args,
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         return multinomial(
             key,
@@ -1181,22 +1183,22 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         value = jnp.array(value, jnp.result_type(float))
         return gammaln(self.total_count + 1) + jnp.sum(
             xlogy(value, self.probs) - gammaln(value + 1), axis=-1
         )
 
     @lazy_property
-    def logits(self) -> ArrayLike:
+    def logits(self) -> Array:
         return _to_logits_multinom(self.probs)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return self.probs * jnp.expand_dims(self.total_count, -1)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.expand_dims(self.total_count, -1) * self.probs * (1 - self.probs)
 
     @constraints.dependent_property(is_discrete=True, event_dim=1)
@@ -1244,7 +1246,7 @@ def __init__(
             validate_args=validate_args,
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         return multinomial(
             key,
@@ -1255,7 +1257,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         )
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         if self._validate_args:
             self._validate_sample(value)
         normalize_term = self.total_count * logsumexp(self.logits, axis=-1) - gammaln(
@@ -1266,15 +1268,15 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         )
 
     @lazy_property
-    def probs(self) -> ArrayLike:
+    def probs(self) -> Array:
         return _to_probs_multinom(self.logits)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.expand_dims(self.total_count, -1) * self.probs
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.expand_dims(self.total_count, -1) * self.probs * (1 - self.probs)
 
     @constraints.dependent_property(is_discrete=True, event_dim=1)
@@ -1353,12 +1355,12 @@ def __init__(
         self.is_sparse = is_sparse
         super(Poisson, self).__init__(jnp.shape(rate), validate_args=validate_args)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         return random.poisson(key, self.rate, shape=sample_shape + self.batch_shape)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         # Using an integer vs. floating-point `rate` leads to differing results.
         # To ensure consistent behavior, `rate` is explicitly cast to a floating-point type.
         # See: https://github.com/pyro-ppl/numpyro/issues/2181
@@ -1388,14 +1390,14 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return xlogy(_value, rate) - gammaln(_value + 1.0) - rate
 
     @property
-    def mean(self) -> ArrayLike:
-        return self.rate
+    def mean(self) -> Array:
+        return jnp.asarray(self.rate)
 
     @property
-    def variance(self) -> ArrayLike:
-        return self.rate
+    def variance(self) -> Array:
+        return jnp.asarray(self.rate)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         k = jnp.floor(value) + 1
         return gammaincc(k, self.rate)
 
@@ -1427,7 +1429,7 @@ def __init__(
             batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         key_bern, key_base = random.split(key)
         shape = sample_shape + self.batch_shape
@@ -1436,7 +1438,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return jnp.where(mask, 0, samples)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         log_prob = jnp.log1p(-self.gate) + self.base_dist.log_prob(value)
         return jnp.where(value == 0, jnp.log(self.gate + jnp.exp(log_prob)), log_prob)
 
@@ -1445,11 +1447,11 @@ def support(self) -> constraints.Constraint:
         return self.base_dist.support
 
     @lazy_property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return (1 - self.gate) * self.base_dist.mean
 
     @lazy_property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return (1 - self.gate) * (
             self.base_dist.mean**2 + self.base_dist.variance
         ) - self.mean**2
@@ -1458,7 +1460,7 @@ def variance(self) -> ArrayLike:
     def has_enumerate_support(self):
         return self.base_dist.has_enumerate_support
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         return self.base_dist.enumerate_support(expand=expand)
 
 
@@ -1478,7 +1480,7 @@ def __init__(
         super().__init__(base_dist, gate, validate_args=validate_args)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         log_prob_minus_log_gate = -self.gate_logits + self.base_dist.log_prob(value)
         log_gate = -softplus(-self.gate_logits)
         log_prob = log_prob_minus_log_gate + log_gate
@@ -1609,18 +1611,18 @@ def __init__(
     def support(self) -> constraints.Constraint:
         return self.base_dist.support
 
-    def _log_one_minus_p_zero(self) -> ArrayLike:
+    def _log_one_minus_p_zero(self) -> Array:
         # log(1 - B(0)) for the discrete base, used to renormalize the truncated PMF.
         log_p0 = self.base_dist.log_prob(jnp.zeros((), dtype=jnp.result_type(int)))
         return jax.nn.log1mexp(-log_p0)
 
-    def _log_gate(self) -> ArrayLike:
+    def _log_gate(self) -> Array:
         return jnp.log(self.gate)
 
-    def _log_one_minus_gate(self) -> ArrayLike:
+    def _log_one_minus_gate(self) -> Array:
         return jnp.log1p(-self.gate)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         key_bern, key_base = random.split(key)
         shape = sample_shape + self.batch_shape
@@ -1631,14 +1633,12 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
             samples = self.base_dist(rng_key=key_base, sample_shape=sample_shape)
         return jnp.where(zero_mask, jnp.zeros_like(samples), samples)
 
-    def _sample_truncated(
-        self, key: jax.Array, sample_shape: tuple[int, ...]
-    ) -> ArrayLike:
+    def _sample_truncated(self, key: jax.Array, sample_shape: tuple[int, ...]) -> Array:
         # Rejection sampling from the zero-truncated base distribution: redraw any
         # element that came back as 0 until all elements are strictly positive.
         first = self.base_dist(rng_key=key, sample_shape=sample_shape)
 
-        def cond_fun(state: tuple) -> ArrayLike:
+        def cond_fun(state: tuple) -> Array:
             _, current = state
             return jnp.any(current == 0)
 
@@ -1653,7 +1653,7 @@ def body_fun(state: tuple) -> tuple:
         return samples
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         log_gate = self._log_gate()
         log_one_minus_gate = self._log_one_minus_gate()
         # Replace zeros with ones before evaluating the base log_prob to avoid
@@ -1669,7 +1669,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return jnp.where(value == 0, log_gate, log_nonzero)
 
     @lazy_property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         if self._is_discrete:
             trunc = -jnp.expm1(
                 self.base_dist.log_prob(jnp.zeros((), dtype=jnp.result_type(int)))
@@ -1678,7 +1678,7 @@ def mean(self) -> ArrayLike:
         return (1 - self.gate) * self.base_dist.mean
 
     @lazy_property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         if self._is_discrete:
             trunc = -jnp.expm1(
                 self.base_dist.log_prob(jnp.zeros((), dtype=jnp.result_type(int)))
@@ -1730,10 +1730,10 @@ def __init__(
         (self.gate_logits,) = promote_shapes(gate_logits, shape=batch_shape)
         super().__init__(base_dist, gate, validate_args=validate_args)
 
-    def _log_gate(self) -> ArrayLike:
+    def _log_gate(self) -> Array:
         return -softplus(-self.gate_logits)
 
-    def _log_one_minus_gate(self) -> ArrayLike:
+    def _log_one_minus_gate(self) -> Array:
         return -softplus(self.gate_logits)
 
 
@@ -1823,7 +1823,7 @@ def __init__(self, probs: ArrayLike, *, validate_args: Optional[bool] = None):
             batch_shape=jnp.shape(self.probs), validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         probs = self.probs
         dtype = jnp.result_type(probs)
@@ -1832,23 +1832,23 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return jnp.floor(jnp.log1p(-u) / jnp.log1p(-probs))
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         probs = jnp.where((self.probs == 1) & (value == 0), 0, self.probs)
         return value * jnp.log1p(-probs) + jnp.log(probs)
 
     @lazy_property
-    def logits(self) -> ArrayLike:
+    def logits(self) -> Array:
         return _to_logits_bernoulli(self.probs)
 
     @property
-    def mean(self) -> ArrayLike:
-        return 1.0 / self.probs - 1.0
+    def mean(self) -> Array:
+        return jnp.asarray(1.0 / self.probs - 1.0)
 
     @property
-    def variance(self) -> ArrayLike:
-        return (1.0 / self.probs - 1.0) / self.probs
+    def variance(self) -> Array:
+        return jnp.asarray((1.0 / self.probs - 1.0) / self.probs)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         return -(1 - self.probs) * jnp.log1p(-self.probs) / self.probs - jnp.log(
             self.probs
         )
@@ -1865,10 +1865,10 @@ def __init__(self, logits: ArrayLike, *, validate_args: Optional[bool] = None):
         )
 
     @lazy_property
-    def probs(self) -> ArrayLike:
+    def probs(self) -> Array:
         return _to_probs_bernoulli(self.logits)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         logits = self.logits
         dtype = jnp.result_type(logits)
@@ -1877,18 +1877,18 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return jnp.floor(jnp.log1p(-u) / -softplus(logits))
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         return (-value - 1) * softplus(self.logits) + self.logits
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return 1.0 / self.probs - 1.0
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return (1.0 / self.probs - 1.0) / self.probs
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         logq = -jax.nn.softplus(self.logits)
         logp = -jax.nn.softplus(-self.logits)
         p = jax.scipy.special.expit(self.logits)
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
index f43a85d7e..e899fc0db 100644
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -346,7 +346,7 @@ def has_rsample(self) -> bool:
 
     def rsample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         if self.has_rsample:
             return self.sample(key, sample_shape=sample_shape)
 
@@ -369,7 +369,7 @@ def shape(self, sample_shape: tuple[int, ...] = ()) -> tuple[int, ...]:
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         """
         Returns a sample from the distribution having shape given by
         `sample_shape + batch_shape + event_shape`. Note that when `sample_shape` is non-empty,
@@ -385,7 +385,7 @@ def sample(
 
     def sample_with_intermediates(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> tuple[ArrayLike, list[Any]]:
+    ) -> tuple[Array, list[Any]]:
         """
         Same as ``sample`` except that any intermediate computations are
         returned (useful for `TransformedDistribution`).
@@ -397,7 +397,7 @@ def sample_with_intermediates(
         """
         return self.sample(key, sample_shape=sample_shape), []
 
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         """
         Evaluates the log probability density for a batch of samples given by
         `value`.
@@ -409,19 +409,23 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         raise NotImplementedError
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         """
         Mean of the distribution.
         """
         raise NotImplementedError
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         """
         Variance of the distribution.
         """
         raise NotImplementedError
 
+    @property
+    def mode(self) -> Array:
+        raise NotImplementedError
+
     def _validate_sample(self, value: ArrayLike) -> ArrayLike:
         assert self.support is not None
         mask = self.support(value)
@@ -442,7 +446,7 @@ def __call__(
         sample_shape: tuple[int, ...] = ...,
         sample_intermediates: Literal[False] = ...,
         **kwargs: Any,
-    ) -> ArrayLike: ...
+    ) -> Array: ...
 
     @overload
     def __call__(
@@ -452,13 +456,13 @@ def __call__(
         sample_shape: tuple[int, ...] = ...,
         sample_intermediates: Literal[True] = ...,
         **kwargs: Any,
-    ) -> tuple[ArrayLike, list[Any]]: ...
+    ) -> tuple[Array, list[Any]]: ...
 
     def __call__(
         self,
         *args: Any,
         **kwargs: Any,
-    ) -> Union[ArrayLike, tuple[ArrayLike, list[Any]]]:
+    ) -> Union[Array, tuple[Array, list[Any]]]:
         key = kwargs.pop("rng_key")
         sample_intermediates = kwargs.pop("sample_intermediates", False)
         if sample_intermediates:
@@ -483,14 +487,14 @@ def to_event(
             return self
         return Independent(self, reinterpreted_batch_ndims)
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         """
         Returns an array with shape `len(support) x batch_shape`
         containing all values in the support.
         """
         raise NotImplementedError
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         """
         Returns the entropy of the distribution.
         """
@@ -627,7 +631,7 @@ def infer_shapes(
         event_shape = ()
         return batch_shape, event_shape
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         """
         The cumulative distribution function of this distribution.
 
@@ -636,7 +640,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         """
         raise NotImplementedError
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         """
         The inverse cumulative distribution function of this distribution.
 
@@ -731,10 +735,10 @@ def has_rsample(self) -> bool:
 
     def _sample(
         self,
-        sample_fn: Callable[..., tuple[ArrayLike, list[ArrayLike]]],
+        sample_fn: Callable[..., tuple[Array, list[Array]]],
         key: Optional[jax.Array],
         sample_shape: tuple[int, ...] = (),
-    ) -> tuple[ArrayLike, list[ArrayLike]]:
+    ) -> tuple[Array, list[Array]]:
         interstitial_sizes = tuple(self._interstitial_sizes.values())
         expanded_sizes = tuple(self._expanded_sizes.values())
         batch_shape = expanded_sizes + interstitial_sizes
@@ -755,7 +759,7 @@ def _sample(
         for dim1, dim2 in zip(interstitial_dims, interstitial_sample_dims):
             permutation[dim1], permutation[dim2] = permutation[dim2], permutation[dim1]
 
-        def reshape_sample(x: ArrayLike) -> ArrayLike:
+        def reshape_sample(x: ArrayLike) -> Array:
             """
             Reshapes samples and intermediates to ensure that the output
             shape is correct: This implicitly replaces the interstitial dims
@@ -772,7 +776,7 @@ def reshape_sample(x: ArrayLike) -> ArrayLike:
 
     def rsample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self._sample(
             lambda *args, **kwargs: (self.base_dist.rsample(*args, **kwargs), []),
             key,
@@ -786,17 +790,17 @@ def support(self) -> constraints.Constraint:
 
     def sample_with_intermediates(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> tuple[ArrayLike, list[ArrayLike]]:
+    ) -> tuple[Array, list[Array]]:
         return self._sample(self.base_dist.sample_with_intermediates, key, sample_shape)
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self.sample_with_intermediates(key, sample_shape)[0]
 
     def log_prob(
         self, value: ArrayLike, intermediates: Optional[list[Any]] = None
-    ) -> ArrayLike:
+    ) -> Array:
         # TODO: utilize `intermediates`
         shape = lax.broadcast_shapes(
             self.batch_shape,
@@ -805,7 +809,7 @@ def log_prob(
         log_prob = self.base_dist.log_prob(value)
         return jnp.broadcast_to(log_prob, shape)
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         samples = jnp.asarray(self.base_dist.enumerate_support(expand=False))
         enum_shape = samples.shape[:1]
         samples = samples.reshape(enum_shape + (1,) * len(self.batch_shape))
@@ -816,18 +820,18 @@ def enumerate_support(self, expand: bool = True) -> ArrayLike:
         return samples
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return jnp.broadcast_to(
             self.base_dist.mean, self.batch_shape + self.event_shape
         )
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.broadcast_to(
             self.base_dist.variance, self.batch_shape + self.event_shape
         )
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         return jnp.broadcast_to(self.base_dist.entropy(), self.batch_shape)
 
 
@@ -900,7 +904,7 @@ def __init__(
         super().__init__(batch_shape, event_shape, validate_args=validate_args)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         batch_shape = jnp.shape(value)[: jnp.ndim(value) - len(self.event_shape)]
         batch_shape = lax.broadcast_shapes(batch_shape, self.batch_shape)
         return jnp.zeros(batch_shape)
@@ -975,11 +979,11 @@ def support(self) -> constraints.Constraint:
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return self.base_dist.mean
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return self.base_dist.variance
 
     @property
@@ -988,15 +992,15 @@ def has_rsample(self) -> bool:
 
     def rsample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self.base_dist.rsample(key, sample_shape=sample_shape)
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self.base_dist.sample(key, sample_shape)
 
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         log_prob = self.base_dist.log_prob(value)
         return sum_rightmost(log_prob, self.reinterpreted_batch_ndims)
 
@@ -1009,7 +1013,7 @@ def expand(self, batch_shape: Sequence[int]) -> Distribution:
             self.reinterpreted_batch_ndims
         )
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         axes = range(-self.reinterpreted_batch_ndims, 0)
         return jnp.asarray(self.base_dist.entropy()).sum(axes)
 
@@ -1053,7 +1057,7 @@ def has_rsample(self) -> bool:
 
     def rsample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self.base_dist.rsample(key, sample_shape=sample_shape)
 
     @property
@@ -1063,10 +1067,10 @@ def support(self) -> constraints.Constraint:
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self.base_dist.sample(key, sample_shape)
 
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         if self._mask is False:
             shape = lax.broadcast_shapes(
                 tuple(self.base_dist.batch_shape),
@@ -1087,15 +1091,15 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
             value = jnp.where(mask, value, default_value)
         return jnp.where(self._mask, self.base_dist.log_prob(value), 0.0)
 
-    def enumerate_support(self, expand: bool = True) -> ArrayLike:
+    def enumerate_support(self, expand: bool = True) -> Array:
         return self.base_dist.enumerate_support(expand=expand)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         return self.base_dist.mean
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return self.base_dist.variance
 
     def tree_flatten(self) -> tuple[tuple[Any, ...], tuple[Any, ...]]:
@@ -1213,11 +1217,11 @@ def has_rsample(self) -> bool:
 
     def rsample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         x = self.base_dist.rsample(key, sample_shape=sample_shape)
         for transform in self.transforms:
             x = transform(x)
-        return x
+        return x  # type: ignore
 
     @property
     def support(self) -> constraints.Constraint:
@@ -1233,27 +1237,29 @@ def support(self) -> constraints.Constraint:
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         x = self.base_dist.sample(key, sample_shape)
         for transform in self.transforms:
             x = transform(x)
-        return x
+        # ``Transform.__call__`` is typed to return the wide ``ArrayLike``, so ``x`` is
+        # widened here even though it is a jax array at runtime.
+        return x  # type: ignore
 
     def sample_with_intermediates(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> tuple[ArrayLike, list[Any]]:
+    ) -> tuple[Array, list[Any]]:
         x = self.base_dist.sample(key, sample_shape)
         intermediates: list[Any] = []
         for transform in self.transforms:
             x_tmp = x
             x, t_inter = transform.call_with_intermediates(x)
             intermediates.append([x_tmp, t_inter])
-        return x, intermediates
+        return x, intermediates  # type: ignore
 
     @validate_sample
     def log_prob(
         self, value: ArrayLike, intermediates: Optional[list[Any]] = None
-    ) -> ArrayLike:
+    ) -> Array:
         if intermediates is not None:
             if len(intermediates) != len(self.transforms):
                 raise ValueError(
@@ -1279,14 +1285,14 @@ def log_prob(
         return log_prob
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         raise NotImplementedError
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         raise NotImplementedError
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         sign: Any = 1
         for transform in reversed(self.transforms):
             sign = sign * transform.sign
@@ -1294,7 +1300,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         q = self.base_dist.cdf(value)
         return jnp.where(sign < 0, 1 - q, q)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         sign: Any = 1
         for transform in self.transforms:
             sign = sign * transform.sign
@@ -1302,7 +1308,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         value = self.base_dist.icdf(q)
         for transform in self.transforms:
             value = transform(value)
-        return value
+        return value  # type: ignore
 
 
 class FoldedDistribution(TransformedDistribution):
@@ -1323,7 +1329,7 @@ def __init__(
         super().__init__(base_dist, AbsTransform(), validate_args=validate_args)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         dim = max(len(self.batch_shape), jnp.ndim(value))
         plus_minus = jnp.array([1.0, -1.0]).reshape((2,) + (1,) * dim)
         return logsumexp(self.base_dist.log_prob(plus_minus * value), axis=0)
@@ -1366,27 +1372,27 @@ def support(self) -> constraints.Constraint:
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         if not sample_shape:
-            return self.v
+            return jnp.asarray(self.v)
         shape = sample_shape + self.batch_shape + self.event_shape
         return jnp.broadcast_to(self.v, shape)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         log_prob = jnp.where(value == self.v, 0, -jnp.inf)
         log_prob = sum_rightmost(log_prob, len(self.event_shape))
         return log_prob + self.log_density
 
     @property
-    def mean(self) -> ArrayLike:
-        return self.v
+    def mean(self) -> Array:
+        return jnp.asarray(self.v)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         return jnp.zeros(self.batch_shape + self.event_shape)
 
-    def entropy(self) -> ArrayLike:
+    def entropy(self) -> Array:
         return -jnp.broadcast_to(self.log_density, self.batch_shape)
 
 
@@ -1412,9 +1418,9 @@ def __init__(self, log_factor: ArrayLike, *, validate_args: bool = False):
 
     def sample(
         self, key: Optional[jax.Array], sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return jnp.empty(sample_shape + self.batch_shape + self.event_shape)
 
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         shape = lax.broadcast_shapes(self.batch_shape, jnp.shape(value)[:-1])
         return jnp.broadcast_to(self.log_factor, shape)
diff --git a/numpyro/distributions/flows.py b/numpyro/distributions/flows.py
index fc8663afd..f699550e4 100644
--- a/numpyro/distributions/flows.py
+++ b/numpyro/distributions/flows.py
@@ -174,7 +174,7 @@ def log_abs_det_jacobian(
     def tree_flatten(self):
         return (), ((), {"bn_arn": self.bn_arn})
 
-    def eq(self, other: Transform, static: bool = False) -> ArrayLike:
+    def eq(self, other: Transform, static: bool = False) -> bool:
         return (
             isinstance(other, BlockNeuralAutoregressiveTransform)
             and self.bn_arn is other.bn_arn
diff --git a/numpyro/distributions/mixtures.py b/numpyro/distributions/mixtures.py
index 0c0d26fdc..6c97922b5 100644
--- a/numpyro/distributions/mixtures.py
+++ b/numpyro/distributions/mixtures.py
@@ -63,22 +63,22 @@ class _MixtureBase(Distribution):
     """
 
     @property
-    def component_mean(self) -> ArrayLike:
+    def component_mean(self) -> Array:
         raise NotImplementedError
 
     @property
-    def component_variance(self) -> ArrayLike:
+    def component_variance(self) -> Array:
         raise NotImplementedError
 
-    def component_log_probs(self, value: ArrayLike) -> ArrayLike:
+    def component_log_probs(self, value: ArrayLike) -> Array:
         raise NotImplementedError
 
     def component_sample(
         self, key: jax.Array, sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         raise NotImplementedError
 
-    def component_cdf(self, samples: ArrayLike) -> ArrayLike:
+    def component_cdf(self, samples: ArrayLike) -> Array:
         raise NotImplementedError
 
     @property
@@ -100,14 +100,14 @@ def mixture_dim(self) -> int:
         return -self.event_dim - 1
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         probs = self.mixing_distribution.probs
         probs = probs.reshape(probs.shape + (1,) * self.event_dim)
         weighted_component_means = probs * self.component_mean
         return jnp.sum(weighted_component_means, axis=self.mixture_dim)
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         probs = self.mixing_distribution.probs
         probs = probs.reshape(probs.shape + (1,) * self.event_dim)
         mean_cond_var = jnp.sum(probs * self.component_variance, axis=self.mixture_dim)
@@ -117,7 +117,7 @@ def variance(self) -> ArrayLike:
         var_cond_mean = jnp.sum(probs * sq_deviation, axis=self.mixture_dim)
         return mean_cond_var + var_cond_mean
 
-    def cdf(self, samples: ArrayLike) -> ArrayLike:
+    def cdf(self, samples: ArrayLike) -> Array:
         """The cumulative distribution function
 
         :param value: samples from this distribution.
@@ -131,7 +131,7 @@ def cdf(self, samples: ArrayLike) -> ArrayLike:
 
     def sample_with_intermediates(
         self, key: jax.Array, sample_shape: tuple[int, ...] = ()
-    ) -> tuple[ArrayLike, list[ArrayLike]]:
+    ) -> tuple[Array, list[Array]]:
         """
         A version of ``sample`` that also returns the sampled component indices
 
@@ -161,11 +161,11 @@ def sample_with_intermediates(
         # Final sample shape (*sample_shape, *batch_shape, *event_shape)
         return jnp.squeeze(samples_selected, axis=self.mixture_dim), [indices]
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         return self.sample_with_intermediates(key=key, sample_shape=sample_shape)[0]
 
     @validate_sample
-    def log_prob(self, value: ArrayLike, intermediates=None) -> ArrayLike:
+    def log_prob(self, value: ArrayLike, intermediates=None) -> Array:
         del intermediates
         sum_log_probs = self.component_log_probs(value)
         safe_sum_log_probs = jnp.where(
@@ -271,26 +271,26 @@ def is_discrete(self) -> bool:
         return self.component_distribution.is_discrete
 
     @property
-    def component_mean(self) -> ArrayLike:
+    def component_mean(self) -> Array:
         return self.component_distribution.mean
 
     @property
-    def component_variance(self) -> ArrayLike:
+    def component_variance(self) -> Array:
         return self.component_distribution.variance
 
-    def component_cdf(self, samples: ArrayLike) -> ArrayLike:
+    def component_cdf(self, samples: ArrayLike) -> Array:
         return self.component_distribution.cdf(
             jnp.expand_dims(samples, axis=self.mixture_dim)
         )
 
     def component_sample(
         self, key: jax.Array, sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         return self.component_distribution.expand(
             sample_shape + self.batch_shape + (self.mixture_size,)
         ).sample(key)
 
-    def component_log_probs(self, value: ArrayLike) -> ArrayLike:
+    def component_log_probs(self, value: ArrayLike) -> Array:
         value = jnp.expand_dims(value, self.mixture_dim)
         component_log_probs = self.component_distribution.log_prob(value)
         return jax.nn.log_softmax(self.mixing_distribution.logits) + component_log_probs
@@ -439,13 +439,13 @@ def is_discrete(self) -> bool:
         return self.component_distributions[0].is_discrete
 
     @property
-    def component_mean(self) -> ArrayLike:
+    def component_mean(self) -> Array:
         return jnp.stack(
             [d.mean for d in self.component_distributions], axis=self.mixture_dim
         )
 
     @property
-    def component_variance(self) -> ArrayLike:
+    def component_variance(self) -> Array:
         return jnp.stack(
             [d.variance for d in self.component_distributions], axis=self.mixture_dim
         )
@@ -458,14 +458,14 @@ def component_cdf(self, samples: ArrayLike) -> Array:
 
     def component_sample(
         self, key: jax.Array, sample_shape: tuple[int, ...] = ()
-    ) -> ArrayLike:
+    ) -> Array:
         keys = jax.random.split(key, self.mixture_size)
         samples = []
         for k, d in zip(keys, self.component_distributions):
             samples.append(d.expand(sample_shape + self.batch_shape).sample(k))
         return jnp.stack(samples, axis=self.mixture_dim)
 
-    def component_log_probs(self, value: ArrayLike) -> ArrayLike:
+    def component_log_probs(self, value: ArrayLike) -> Array:
         component_log_probs = []
         for d in self.component_distributions:
             log_prob = d.log_prob(value)
diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py
index 8a2c9b603..b939328b0 100644
--- a/numpyro/distributions/transforms.py
+++ b/numpyro/distributions/transforms.py
@@ -448,9 +448,11 @@ def eq(self, other: object, static: bool = False) -> ArrayLike:
             return all(
                 p1.eq(p2, static=True) for p1, p2 in zip(self.parts, other.parts)
             )
-        result = jnp.array(True)
+        result = True
         for p1, p2 in zip(self.parts, other.parts):
-            result = result & p1.eq(p2, static=False)
+            # eq is boolean-valued; ArrayLike advertises float/complex members
+            # that don't support `&`, so `&` is safe here in practice.
+            result = result & p1.eq(p2, static=False)  # ty: ignore[unsupported-operator]
         return result
 
 
diff --git a/numpyro/distributions/truncated.py b/numpyro/distributions/truncated.py
index ba4075fcb..140077009 100644
--- a/numpyro/distributions/truncated.py
+++ b/numpyro/distributions/truncated.py
@@ -5,7 +5,7 @@
 from typing import Optional, Union
 
 import jax
-from jax import lax
+from jax import Array, lax
 import jax.numpy as jnp
 import jax.random as random
 from jax.scipy.special import logsumexp
@@ -72,7 +72,7 @@ def _tail_prob_at_high(self):
         # if low < loc, returns cdf(high) = 1; otherwise returns 1 - cdf(high) = 0
         return jnp.where(self.low <= self.base_dist.loc, 1.0, 0.0)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         dtype = jnp.result_type(float)
         finfo = jnp.finfo(dtype)
@@ -80,7 +80,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         u = random.uniform(key, shape=sample_shape + self.batch_shape, minval=minval)
         return self.icdf(u)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         loc = self.base_dist.loc
         sign = jnp.where(loc >= self.low, 1.0, -1.0)
         ppf = (1 - sign) * loc + sign * self.base_dist.icdf(
@@ -88,7 +88,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         )
         return jnp.where(q < 0, jnp.nan, ppf)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         # For left truncated distribution: CDF(x) = (F(x) - F(low)) / (1 - F(low))
         # where F is the base distribution CDF
         base_cdf_value = self.base_dist.cdf(value)
@@ -103,14 +103,14 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         return result
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         sign = jnp.where(self.base_dist.loc >= self.low, 1.0, -1.0)
         return self.base_dist.log_prob(value) - jnp.log(
             sign * (self._tail_prob_at_high - self._tail_prob_at_low)
         )
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         if isinstance(self.base_dist, Normal):
             low_prob = jnp.exp(self.log_prob(self.low))
             return self.base_dist.loc + low_prob * self.base_dist.scale**2
@@ -120,7 +120,7 @@ def mean(self) -> ArrayLike:
             raise NotImplementedError("mean only available for Normal and Cauchy")
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         if isinstance(self.base_dist, Normal):
             low_prob = jnp.exp(self.log_prob(self.low))
             return (self.base_dist.scale**2) * (
@@ -164,10 +164,10 @@ def support(self) -> Constraint:
         return self._support
 
     @lazy_property
-    def _cdf_at_high(self) -> ArrayLike:
+    def _cdf_at_high(self) -> Array:
         return self.base_dist.cdf(self.high)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         dtype = jnp.result_type(float)
         finfo = jnp.finfo(dtype)
@@ -175,11 +175,11 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         u = random.uniform(key, shape=sample_shape + self.batch_shape, minval=minval)
         return self.icdf(u)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         ppf = self.base_dist.icdf(q * self._cdf_at_high)
         return jnp.where(q > 1, jnp.nan, ppf)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         # For right truncated distribution: CDF(x) = F(x) / F(high)
         # where F is the base distribution CDF
         base_cdf_value = self.base_dist.cdf(value)
@@ -194,11 +194,11 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         return result
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         return self.base_dist.log_prob(value) - jnp.log(self._cdf_at_high)
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         if isinstance(self.base_dist, Normal):
             high_prob = jnp.exp(self.log_prob(self.high))
             return self.base_dist.loc - high_prob * self.base_dist.scale**2
@@ -208,7 +208,7 @@ def mean(self) -> ArrayLike:
             raise NotImplementedError("mean only available for Normal and Cauchy")
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         if isinstance(self.base_dist, Normal):
             high_prob = jnp.exp(self.log_prob(self.high))
             return (self.base_dist.scale**2) * (
@@ -259,21 +259,21 @@ def support(self) -> Constraint:
         return self._support
 
     @lazy_property
-    def _tail_prob_at_low(self) -> ArrayLike:
+    def _tail_prob_at_low(self) -> Array:
         # if low < loc, returns cdf(low); otherwise returns 1 - cdf(low)
         loc = self.base_dist.loc
         sign = jnp.where(loc >= self.low, 1.0, -1.0)
         return self.base_dist.cdf(loc - sign * (loc - self.low))
 
     @lazy_property
-    def _tail_prob_at_high(self) -> ArrayLike:
+    def _tail_prob_at_high(self) -> Array:
         # if low < loc, returns cdf(high); otherwise returns 1 - cdf(high)
         loc = self.base_dist.loc
         sign = jnp.where(loc >= self.low, 1.0, -1.0)
         return self.base_dist.cdf(loc - sign * (loc - self.high))
 
     @lazy_property
-    def _log_diff_tail_probs(self) -> ArrayLike:
+    def _log_diff_tail_probs(self) -> Array:
         # use log_cdf method, if available, to avoid inf's in log_prob
         # fall back to cdf, if log_cdf not available
         log_cdf = getattr(self.base_dist, "log_cdf", None)
@@ -289,7 +289,7 @@ def _log_diff_tail_probs(self) -> ArrayLike:
             sign = jnp.where(loc >= self.low, 1.0, -1.0)
             return jnp.log(sign * (self._tail_prob_at_high - self._tail_prob_at_low))
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         dtype = jnp.result_type(float)
         finfo = jnp.finfo(dtype)
@@ -297,7 +297,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         u = random.uniform(key, shape=sample_shape + self.batch_shape, minval=minval)
         return self.icdf(u)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         # NB: we use a more numerically stable formula for a symmetric base distribution
         #   A = icdf(cdf(low) + (cdf(high) - cdf(low)) * q) = icdf[(1 - q) * cdf(low) + q * cdf(high)]
         # will suffer by precision issues when low is large;
@@ -312,7 +312,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         )
         return jnp.where(jnp.logical_or(q < 0, q > 1), jnp.nan, ppf)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         # For two-sided truncated distribution: CDF(x) = (F(x) - F(low)) / (F(high) - F(low))
         # where F is the base distribution CDF
         base_cdf_value = self.base_dist.cdf(value)
@@ -334,7 +334,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         return result
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         # NB: we use a more numerically stable formula for a symmetric base distribution
         # if low < loc
         #   cdf(high) - cdf(low) = as-is
@@ -343,7 +343,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         return self.base_dist.log_prob(value) - self._log_diff_tail_probs
 
     @property
-    def mean(self) -> ArrayLike:
+    def mean(self) -> Array:
         if isinstance(self.base_dist, Normal):
             low_prob = jnp.exp(self.log_prob(self.low))
             high_prob = jnp.exp(self.log_prob(self.high))
@@ -354,7 +354,7 @@ def mean(self) -> ArrayLike:
             raise NotImplementedError("mean only available for Normal and Cauchy")
 
     @property
-    def variance(self) -> ArrayLike:
+    def variance(self) -> Array:
         if isinstance(self.base_dist, Normal):
             low_prob = jnp.exp(self.log_prob(self.low))
             high_prob = jnp.exp(self.log_prob(self.high))
@@ -447,7 +447,7 @@ def __init__(
             batch_shape, validate_args=validate_args
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         denom = jnp.square(jnp.arange(0.5, self.num_gamma_variates))
         x = random.gamma(
@@ -457,7 +457,7 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
         return jnp.clip(x * (0.5 / jnp.pi**2), None, self.truncation_point)
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         value = value[..., None]
         all_indices = jnp.arange(0, self.num_log_prob_terms)
         two_n_plus_one = 2.0 * all_indices + 1.0
@@ -525,7 +525,7 @@ def support(self) -> Constraint:
         return self._support
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         r"""Logarithmic probability distribution:
 
         Z inequal minus one:
@@ -544,14 +544,12 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
         """
 
         @jax.custom_jvp
-        def f(
-            x: ArrayLike, alpha: ArrayLike, low: ArrayLike, high: ArrayLike
-        ) -> ArrayLike:
+        def f(x: ArrayLike, alpha: ArrayLike, low: ArrayLike, high: ArrayLike) -> Array:
             neq_neg1_mask = jnp.not_equal(alpha, -1.0)
             neq_neg1_alpha = jnp.where(neq_neg1_mask, alpha, 0.0)
             # eq_neg1_alpha = jnp.where(~neq_neg1_mask, alpha, -1.0)
 
-            def neq_neg1_fn() -> ArrayLike:
+            def neq_neg1_fn() -> Array:
                 one_more_alpha = 1.0 + neq_neg1_alpha
                 return jnp.log(
                     jnp.power(x, neq_neg1_alpha)
@@ -559,7 +557,7 @@ def neq_neg1_fn() -> ArrayLike:
                     / (jnp.power(high, one_more_alpha) - jnp.power(low, one_more_alpha))
                 )
 
-            def eq_neg1_fn() -> ArrayLike:
+            def eq_neg1_fn() -> Array:
                 return -jnp.log(x) - jnp.log(jnp.log(high) - jnp.log(low))
 
             return jnp.where(neq_neg1_mask, neq_neg1_fn(), eq_neg1_fn())
@@ -586,7 +584,7 @@ def f_jvp(
 
             # Alpha tangent with approximation
             # Variable part for all values alpha unequal -1
-            def alpha_tangent_variable(alpha: ArrayLike) -> ArrayLike:
+            def alpha_tangent_variable(alpha: ArrayLike) -> Array:
                 one_more_alpha = 1.0 + alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 high_pow_one_more_alpha = jnp.power(high, one_more_alpha)
@@ -649,7 +647,7 @@ def alpha_tangent_variable(alpha: ArrayLike) -> ArrayLike:
 
         return f(value, self.alpha, self.low, self.high)
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         r"""Cumulated probability distribution:
         Z inequal minus one:
 
@@ -667,20 +665,18 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         """
 
         @jax.custom_jvp
-        def f(
-            x: ArrayLike, alpha: ArrayLike, low: ArrayLike, high: ArrayLike
-        ) -> ArrayLike:
+        def f(x: ArrayLike, alpha: ArrayLike, low: ArrayLike, high: ArrayLike) -> Array:
             neq_neg1_mask = jnp.not_equal(alpha, -1.0)
             neq_neg1_alpha = jnp.where(neq_neg1_mask, alpha, 0.0)
 
-            def cdf_when_alpha_neq_neg1() -> ArrayLike:
+            def cdf_when_alpha_neq_neg1() -> Array:
                 one_more_alpha = 1.0 + neq_neg1_alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 return (jnp.power(x, one_more_alpha) - low_pow_one_more_alpha) / (
                     jnp.power(high, one_more_alpha) - low_pow_one_more_alpha
                 )
 
-            def cdf_when_alpha_eq_neg1() -> ArrayLike:
+            def cdf_when_alpha_eq_neg1() -> Array:
                 return jnp.log(x / low) / jnp.log(high / low)
 
             cdf_val = jnp.where(
@@ -694,7 +690,7 @@ def cdf_when_alpha_eq_neg1() -> ArrayLike:
         def f_jvp(
             primals: tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike],
             tangents: tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike],
-        ) -> tuple[ArrayLike, ArrayLike]:
+        ) -> tuple[Array, Array]:
             x, alpha, low, high = primals
             x_t, alpha_t, low_t, high_t = tangents
 
@@ -710,13 +706,13 @@ def f_jvp(
             primal_out = f(*primals)
 
             # Tangents for alpha not equals -1
-            def x_neq_neg1(alpha: ArrayLike) -> ArrayLike:
+            def x_neq_neg1(alpha: ArrayLike) -> Array:
                 one_more_alpha = 1.0 + alpha
                 return (one_more_alpha * jnp.power(x, alpha)) / (
                     jnp.power(high, one_more_alpha) - jnp.power(low, one_more_alpha)
                 )
 
-            def alpha_neq_neg1(alpha: ArrayLike) -> ArrayLike:
+            def alpha_neq_neg1(alpha: ArrayLike) -> Array:
                 one_more_alpha = 1.0 + alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 high_pow_one_more_alpha = jnp.power(high, one_more_alpha)
@@ -733,7 +729,7 @@ def alpha_neq_neg1(alpha: ArrayLike) -> ArrayLike:
                 ) / jnp.square(high_pow_one_more_alpha - low_pow_one_more_alpha)
                 return term1 - term2
 
-            def low_neq_neg1(alpha: ArrayLike) -> ArrayLike:
+            def low_neq_neg1(alpha: ArrayLike) -> Array:
                 one_more_alpha = 1.0 + alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 high_pow_one_more_alpha = jnp.power(high, one_more_alpha)
@@ -743,7 +739,7 @@ def low_neq_neg1(alpha: ArrayLike) -> ArrayLike:
                 term1 = term2 * (x_pow_one_more_alpha - low_pow_one_more_alpha) / change
                 return term1 - term2
 
-            def high_neq_neg1(alpha: ArrayLike) -> ArrayLike:
+            def high_neq_neg1(alpha: ArrayLike) -> Array:
                 one_more_alpha = 1.0 + alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 high_pow_one_more_alpha = jnp.power(high, one_more_alpha)
@@ -755,15 +751,15 @@ def high_neq_neg1(alpha: ArrayLike) -> ArrayLike:
                 ) / jnp.square(high_pow_one_more_alpha - low_pow_one_more_alpha)
 
             # Tangents for alpha equals -1
-            def x_eq_neg1() -> ArrayLike:
+            def x_eq_neg1() -> Array:
                 return jnp.reciprocal(x * (log_high - log_low))
 
-            def low_eq_neg1() -> ArrayLike:
+            def low_eq_neg1() -> Array:
                 return (log_x - log_low) / (
                     jnp.square(log_high - log_low) * low
                 ) - jnp.reciprocal((log_high - log_low) * low)
 
-            def high_eq_neg1() -> ArrayLike:
+            def high_eq_neg1() -> Array:
                 return (log_x - log_low) / (jnp.square(log_high - log_low) * high)
 
             # Including approximation for alpha = -1
@@ -791,7 +787,7 @@ def high_eq_neg1() -> ArrayLike:
 
         return f(value, self.alpha, self.low, self.high)
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         r"""Inverse cumulated probability distribution:
         Z inequal minus one:
 
@@ -807,13 +803,11 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
         """
 
         @jax.custom_jvp
-        def f(
-            q: ArrayLike, alpha: ArrayLike, low: ArrayLike, high: ArrayLike
-        ) -> ArrayLike:
+        def f(q: ArrayLike, alpha: ArrayLike, low: ArrayLike, high: ArrayLike) -> Array:
             neq_neg1_mask = jnp.not_equal(alpha, -1.0)
             neq_neg1_alpha = jnp.where(neq_neg1_mask, alpha, 0.0)
 
-            def icdf_alpha_neq_neg1() -> ArrayLike:
+            def icdf_alpha_neq_neg1() -> Array:
                 one_more_alpha = 1.0 + neq_neg1_alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 high_pow_one_more_alpha = jnp.power(high, one_more_alpha)
@@ -823,7 +817,7 @@ def icdf_alpha_neq_neg1() -> ArrayLike:
                     jnp.reciprocal(one_more_alpha),
                 )
 
-            def icdf_alpha_eq_neg1() -> ArrayLike:
+            def icdf_alpha_eq_neg1() -> Array:
                 return jnp.power(high / low, q) * low
 
             icdf_val = jnp.where(
@@ -837,7 +831,7 @@ def icdf_alpha_eq_neg1() -> ArrayLike:
         def f_jvp(
             primals: tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike],
             tangents: tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike],
-        ) -> tuple[ArrayLike, ArrayLike]:
+        ) -> tuple[Array, Array]:
             x, alpha, low, high = primals
             x_t, alpha_t, low_t, high_t = tangents
 
@@ -852,7 +846,7 @@ def f_jvp(
             primal_out = f(*primals)
 
             # Tangents for alpha not equal -1
-            def x_neq_neg1(alpha: ArrayLike) -> ArrayLike:
+            def x_neq_neg1(alpha: ArrayLike) -> Array:
                 one_more_alpha = 1.0 + alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 high_pow_one_more_alpha = jnp.power(high, one_more_alpha)
@@ -865,7 +859,7 @@ def x_neq_neg1(alpha: ArrayLike) -> ArrayLike:
                     )
                 ) / one_more_alpha
 
-            def alpha_neq_neg1(alpha: ArrayLike) -> ArrayLike:
+            def alpha_neq_neg1(alpha: ArrayLike) -> Array:
                 one_more_alpha = 1.0 + alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 high_pow_one_more_alpha = jnp.power(high, one_more_alpha)
@@ -884,7 +878,7 @@ def alpha_neq_neg1(alpha: ArrayLike) -> ArrayLike:
                 term3 = jnp.log(factor0) / jnp.square(one_more_alpha)
                 return term1 * (term2 - term3)
 
-            def low_neq_neg1(alpha: ArrayLike) -> ArrayLike:
+            def low_neq_neg1(alpha: ArrayLike) -> Array:
                 one_more_alpha = 1.0 + alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 high_pow_one_more_alpha = jnp.power(high, one_more_alpha)
@@ -898,7 +892,7 @@ def low_neq_neg1(alpha: ArrayLike) -> ArrayLike:
                     )
                 )
 
-            def high_neq_neg1(alpha: ArrayLike) -> ArrayLike:
+            def high_neq_neg1(alpha: ArrayLike) -> Array:
                 one_more_alpha = 1.0 + alpha
                 low_pow_one_more_alpha = jnp.power(low, one_more_alpha)
                 high_pow_one_more_alpha = jnp.power(high, one_more_alpha)
@@ -913,16 +907,16 @@ def high_neq_neg1(alpha: ArrayLike) -> ArrayLike:
                 )
 
             # Tangents for alpha equals -1
-            def dx_eq_neg1() -> ArrayLike:
+            def dx_eq_neg1() -> Array:
                 return low * jnp.power(high_over_low, x) * (log_high - log_low)
 
-            def low_eq_neg1() -> ArrayLike:
+            def low_eq_neg1() -> Array:
                 return (
                     jnp.power(high_over_low, x)
                     - (high * x * jnp.power(high_over_low, x - 1)) / low
                 )
 
-            def high_eq_neg1() -> ArrayLike:
+            def high_eq_neg1() -> Array:
                 return x * jnp.power(high_over_low, x - 1)
 
             # Including approximation for alpha = -1 \
@@ -950,7 +944,7 @@ def high_eq_neg1() -> ArrayLike:
 
         return f(q, self.alpha, self.low, self.high)
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         u = random.uniform(key, sample_shape + self.batch_shape)
         samples = self.icdf(u)
@@ -1004,7 +998,7 @@ def support(self) -> Constraint:
         return self._support
 
     @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
+    def log_prob(self, value: ArrayLike) -> Array:
         one_more_alpha = 1.0 + self.alpha
         return (
             self.alpha * jnp.log(value)
@@ -1012,7 +1006,7 @@ def log_prob(self, value: ArrayLike) -> ArrayLike:
             - one_more_alpha * jnp.log(self.low)
         )
 
-    def cdf(self, value: ArrayLike) -> ArrayLike:
+    def cdf(self, value: ArrayLike) -> Array:
         cdf_val = jnp.where(
             jnp.less_equal(value, self.low),
             jnp.zeros_like(value),
@@ -1020,7 +1014,7 @@ def cdf(self, value: ArrayLike) -> ArrayLike:
         )
         return cdf_val
 
-    def icdf(self, q: ArrayLike) -> ArrayLike:
+    def icdf(self, q: ArrayLike) -> Array:
         nan_mask = jnp.logical_or(jnp.isnan(q), jnp.less(q, 0.0))
         nan_mask = jnp.logical_or(nan_mask, jnp.greater(q, 1.0))
         return jnp.where(
@@ -1029,7 +1023,7 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
             self.low * jnp.power(1.0 - q, jnp.reciprocal(1.0 + self.alpha)),
         )
 
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> Array:
         assert is_prng_key(key)
         u = random.uniform(key, sample_shape + self.batch_shape)
         samples = self.icdf(u)
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
index 52e07aeb0..1791c6900 100644
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -11,7 +11,7 @@
 import numpy as np
 
 import jax
-from jax import jit, lax, random, vmap
+from jax import Array, jit, lax, random, vmap
 import jax.numpy as jnp
 from jax.scipy.linalg import solve_triangular
 from jax.scipy.special import digamma
@@ -455,7 +455,7 @@ def logmatmulexp(x, y):
 
 
 @jax.custom_jvp
-def log1mexp(x: ArrayLike) -> ArrayLike:
+def log1mexp(x: ArrayLike) -> Array:
     """
     Numerically stable calculation of the quantity
     :math:`\\log(1 - \\exp(x))`, following the algorithm
@@ -485,7 +485,7 @@ def log1mexp(x: ArrayLike) -> ArrayLike:
 log1mexp.defjvps(lambda t, ans, x: -t / jnp.expm1(-x))
 
 
-def logdiffexp(a: ArrayLike, b: ArrayLike) -> ArrayLike:
+def logdiffexp(a: ArrayLike, b: ArrayLike) -> Array:
     """
     Numerically stable calculation of the
     quantity :math:`\\log(\\exp(a) - \\exp(b))`,
@@ -511,7 +511,7 @@ def logdiffexp(a: ArrayLike, b: ArrayLike) -> ArrayLike:
     )
 
 
-def clamp_probs(probs: ArrayLike) -> ArrayLike:
+def clamp_probs(probs: ArrayLike) -> Array:
     finfo = jnp.finfo(jnp.result_type(probs, float))
     return jnp.clip(probs, finfo.tiny, 1.0 - finfo.eps)
 
diff --git a/test/test_distributions.py b/test/test_distributions.py
index 9604ddde8..b1628cbe3 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -8,7 +8,7 @@
 from itertools import product
 import math
 import os
-from typing import Callable
+from typing import Callable, get_type_hints
 
 import numpy as np
 from numpy.testing import assert_allclose, assert_array_equal
@@ -18,7 +18,7 @@
 import scipy.stats as osp
 
 import jax
-from jax import grad, lax, vmap
+from jax import Array, grad, lax, vmap
 import jax.numpy as jnp
 import jax.random as random
 from jax.scipy.special import expit, logsumexp
@@ -385,7 +385,7 @@ def _vmap_over_general_2d_mixture(self: _General2DMixture, locs=None):
 
 
 class _ImproperWrapper(dist.ImproperUniform):
-    def sample(self, key, sample_shape=()):
+    def sample(self, key, sample_shape=()) -> Array:
         transform = biject_to(self.support)
         prototype_value = jnp.zeros(self.event_shape)
         unconstrained_event_shape = jnp.shape(transform.inv(prototype_value))
@@ -2221,6 +2221,76 @@ def fn(*args):
         assert_allclose(jnp.sum(actual_grad), expected_grad, rtol=rtol, atol=atol)
 
 
+@pytest.mark.parametrize("jit", [False, True])
+@pytest.mark.parametrize(
+    "method",
+    [
+        "sample",
+        "log_prob",
+        "mean",
+        "variance",
+        "entropy",
+        "cdf",
+        "icdf",
+        "rsample",
+        "mode",
+        "enumerate_support",
+    ],
+)
+@pytest.mark.parametrize(
+    "jax_dist, sp_dist, params", CONTINUOUS + DISCRETE + DIRECTIONAL
+)
+def test_output_is_array(jax_dist, sp_dist, params, method, jit, request):
+    # Distribution methods should return jax arrays so that consumers can index,
+    # reshape, and compare the result without tripping a type checker (``ArrayLike``
+    # admits python/numpy scalars that are not indexable and have no ``.shape``).
+    # Validation is disabled so the test is independent of global validation state.
+
+    with dist.distribution.validation_enabled(False):
+        d = jax_dist(*params)
+
+        cls = d.__class__
+        # Verify the correct annotation after unpacking.
+        impl = getattr(cls, method)
+        if isinstance(impl, property):
+            impl = impl.fget
+        while isinstance(impl, partial):
+            impl = impl.func
+        return_annotation = get_type_hints(impl).get("return")
+        assert return_annotation is jax.Array, (
+            f"{cls.__name__}.{method} ({impl.__code__.co_filename}:"
+            f"{impl.__code__.co_firstlineno}) is annotated with {return_annotation}"
+        )
+
+        if method in {"sample", "rsample"}:
+            fn = lambda: getattr(d, method)(random.PRNGKey(0))  # noqa: E731
+        elif method in {"log_prob", "cdf"}:
+            # IntervalCensoredDistribution takes an interval (lo, hi) as log_prob
+            # input but returns univariate samples, so log_prob(sample) is not
+            # well-formed for it (mirrors the special-casing in test_dist_shape).
+            if isinstance(d, dist.IntervalCensoredDistribution):
+                pytest.skip(
+                    "log_prob(sample) is ill-formed for interval-censored dists"
+                )
+            value = d.sample(random.PRNGKey(0))
+            fn = lambda: getattr(d, method)(value)  # noqa: E731
+        elif method == "icdf":
+            fn = lambda: d.icdf(0.3)  # noqa: E731
+        elif method == "entropy":
+            fn = lambda: d.entropy()  # noqa: E731
+        elif method in {"mean", "mode", "variance"}:
+            fn = lambda: getattr(d, method)  # noqa: E731
+        elif method == "enumerate_support":
+            fn = lambda: d.enumerate_support()  # noqa: E731
+        else:
+            raise RuntimeError(method)
+        try:
+            out = jax.jit(fn)() if jit else fn()
+        except NotImplementedError:
+            pytest.skip(f"{cls.__name__}.{method} is not implemented")
+    assert isinstance(out, jax.Array), f"{cls.__name__}.{method} returned {type(out)}"
+
+
 @pytest.mark.parametrize(
     "jax_dist, sp_dist, params", CONTINUOUS + DISCRETE + DIRECTIONAL
 )
@@ -2546,15 +2616,22 @@ def test_distribution_constraints(jax_dist, sp_dist, params, prepend_shape):
         with pytest.raises(ValueError):
             jax_dist(*oob_params, validate_args=True)
 
-        with pytest.raises(ValueError):
-            # test error raised under jit omnistaging
-            oob_params = jax.device_get(oob_params)
-
-            def dist_gen_fn():
-                d = jax_dist(*oob_params, validate_args=True)
-                return d
-
-            jax.jit(dist_gen_fn)()
+        # Trace-time value validation only fires when the constraint result
+        # constant-folds to a concrete value. BetaProportion validates its mean
+        # parameter via the inherited Beta.mean property, which now returns a
+        # jax.Array (a tracer under jit), so validation is correctly skipped
+        # there (validate_args is best-effort under jit; see Distribution.
+        # validate_args with strict=False). The eager check above still covers it.
+        if jax_dist is not dist.BetaProportion:
+            with pytest.raises(ValueError):
+                # test error raised under jit omnistaging
+                oob_params = jax.device_get(oob_params)
+
+                def dist_gen_fn():
+                    d = jax_dist(*oob_params, validate_args=True)
+                    return d
+
+                jax.jit(dist_gen_fn)()
 
     d = jax_dist(*valid_params, validate_args=True)