Ensure distribution moments/samples are jax arrays

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,15 +112,15 @@ 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)
     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,15 +228,15 @@ 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)
     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}"

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,24 +399,24 @@ 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::
             \mathrm{Var}[X] = \frac{\alpha}{\lambda^2}(1 + \lambda)
         """
         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:
 

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: