Add type hints to the primitives module

numpyro/contrib/hsgp/approximation.py

@@ -21,26 +21,22 @@
 import numpyro.distributions as dist
 
 
-def _non_centered_approximation(
-    phi: ArrayLike, spd: ArrayLike, m: int | list[int]
-) -> Array:
+def _non_centered_approximation(phi: Array, spd: Array, m: int) -> Array:
     with numpyro.plate("basis", m):
         beta = numpyro.sample("beta", dist.Normal(loc=0.0, scale=1.0))
 
-    return phi @ (spd * beta)
+    return jnp.asarray(phi @ (spd * beta))
 
 
-def _centered_approximation(
-    phi: ArrayLike, spd: ArrayLike, m: int | list[int]
-) -> Array:
+def _centered_approximation(phi: Array, spd: Array, m: int) -> Array:
     with numpyro.plate("basis", m):
         beta = numpyro.sample("beta", dist.Normal(loc=0.0, scale=spd))
 
-    return phi @ beta
+    return jnp.asarray(phi @ beta)
 
 
 def linear_approximation(
-    phi: ArrayLike, spd: ArrayLike, m: int | list[int], non_centered: bool = True
+    phi: Array, spd: Array, m: int, non_centered: bool = True
 ) -> Array:
     """
     Linear approximation formula of the Hilbert space Gaussian process.
@@ -52,10 +48,10 @@ def linear_approximation(
         1. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
            approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
 
-    :param ArrayLike phi: laplacian eigenfunctions
-    :param ArrayLike spd: square root of the diagonal of the spectral density evaluated at square
+    :param Array phi: laplacian eigenfunctions
+    :param Array spd: square root of the diagonal of the spectral density evaluated at square
         root of the first `m` eigenvalues.
-    :param int | list[int] m: number of eigenfunctions in the approximation
+    :param int m: number of eigenfunctions in the approximation
     :param bool non_centered: whether to use a non-centered parameterization
     :return: The low-rank approximation linear model
     :rtype: Array

numpyro/diagnostics.py

@@ -10,6 +10,7 @@
 from typing import Union
 
 import numpy as np
+import numpy.typing as npt
 
 import jax
 from jax import device_get
@@ -25,7 +26,7 @@
 ]
 
 
-def _compute_chain_variance_stats(x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+def _compute_chain_variance_stats(x: npt.NDArray) -> tuple[npt.NDArray, npt.NDArray]:
     # compute within-chain variance and variance estimator
     # input has shape C x N x sample_shape
     C, N = x.shape[:2]
@@ -41,7 +42,7 @@ def _compute_chain_variance_stats(x: np.ndarray) -> tuple[np.ndarray, np.ndarray
     return var_within, var_estimator
 
 
-def gelman_rubin(x: np.ndarray) -> np.ndarray:
+def gelman_rubin(x: npt.NDArray) -> npt.NDArray:
     """
     Computes R-hat over chains of samples ``x``, where the first dimension of
     ``x`` is chain dimension and the second dimension of ``x`` is draw dimension.
@@ -60,7 +61,7 @@ def gelman_rubin(x: np.ndarray) -> np.ndarray:
     return rhat
 
 
-def split_gelman_rubin(x: np.ndarray) -> np.ndarray:
+def split_gelman_rubin(x: npt.NDArray) -> npt.NDArray:
     """
     Computes split R-hat over chains of samples ``x``, where the first dimension
     of ``x`` is chain dimension and the second dimension of ``x`` is draw dimension.
@@ -97,7 +98,7 @@ def _fft_next_fast_len(target: int) -> int:
         target += 1
 
 
-def autocorrelation(x: np.ndarray, axis: int = 0, bias: bool = True) -> np.ndarray:
+def autocorrelation(x: npt.NDArray, axis: int = 0, bias: bool = True) -> npt.NDArray:
     """
     Computes the autocorrelation of samples at dimension ``axis``.
 
@@ -137,11 +138,11 @@ def autocorrelation(x: np.ndarray, axis: int = 0, bias: bool = True) -> np.ndarr
         autocorr = autocorr / np.arange(N, 0.0, -1)
 
     with np.errstate(invalid="ignore", divide="ignore"):
-        autocorr = autocorr / autocorr[..., :1]
+        autocorr = (autocorr / autocorr[..., :1]).astype(np.float64)
     return np.swapaxes(autocorr, axis, -1)
 
 
-def autocovariance(x: np.ndarray, axis: int = 0, bias: bool = True) -> np.ndarray:
+def autocovariance(x: npt.NDArray, axis: int = 0, bias: bool = True) -> npt.NDArray:
     """
     Computes the autocovariance of samples at dimension ``axis``.
 
@@ -154,7 +155,7 @@ def autocovariance(x: np.ndarray, axis: int = 0, bias: bool = True) -> np.ndarra
     return autocorrelation(x, axis, bias) * x.var(axis=axis, keepdims=True)
 
 
-def effective_sample_size(x: np.ndarray, bias: bool = True) -> np.ndarray:
+def effective_sample_size(x: npt.NDArray, bias: bool = True) -> npt.NDArray:
     """
     Computes effective sample size of input ``x``, where the first dimension of
     ``x`` is chain dimension and the second dimension of ``x`` is draw dimension.
@@ -202,7 +203,7 @@ def effective_sample_size(x: np.ndarray, bias: bool = True) -> np.ndarray:
     return n_eff
 
 
-def hpdi(x: np.ndarray, prob: float = 0.90, axis: int = 0) -> np.ndarray:
+def hpdi(x: npt.NDArray, prob: float = 0.90, axis: int = 0) -> npt.NDArray:
     """
     Computes "highest posterior density interval" (HPDI) which is the narrowest
     interval with probability mass ``prob``.
@@ -285,7 +286,7 @@ def summary(
 
 
 def print_summary(
-    samples: Union[dict, np.ndarray], prob: float = 0.90, group_by_chain: bool = True
+    samples: Union[dict, npt.NDArray], prob: float = 0.90, group_by_chain: bool = True
 ) -> None:
     """
     Prints a summary table displaying diagnostics of ``samples`` from the

numpyro/distributions/distribution.py

@@ -24,19 +24,20 @@
 # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 # POSSIBILITY OF SUCH DAMAGE.
-
 from collections import OrderedDict
 from contextlib import contextmanager
 import functools
 import inspect
+from typing import Protocol, runtime_checkable
 import warnings
 
 import numpy as np
 
 import jax
-from jax import lax, tree_util
+from jax import Array, lax, tree_util
 import jax.numpy as jnp
 from jax.scipy.special import logsumexp
+from jax.typing import ArrayLike
 
 from numpyro.distributions.transforms import AbsTransform, ComposeTransform, Transform
 from numpyro.distributions.util import (
@@ -270,7 +271,7 @@ def validate_args(self, strict: bool = True) -> None:
                 raise RuntimeError("Cannot validate arguments inside jitted code.")
 
     @property
-    def batch_shape(self):
+    def batch_shape(self) -> tuple[int, ...]:
         """
         Returns the shape over which the distribution parameters are batched.
 
@@ -280,7 +281,7 @@ def batch_shape(self):
         return self._batch_shape
 
     @property
-    def event_shape(self):
+    def event_shape(self) -> tuple[int, ...]:
         """
         Returns the shape of a single sample from the distribution without
         batching.
@@ -291,15 +292,15 @@ def event_shape(self):
         return self._event_shape
 
     @property
-    def event_dim(self):
+    def event_dim(self) -> int:
         """
         :return: Number of dimensions of individual events.
         :rtype: int
         """
         return len(self.event_shape)
 
     @property
-    def has_rsample(self):
+    def has_rsample(self) -> bool:
         return set(self.reparametrized_params) == set(self.arg_constraints)
 
     def rsample(self, key, sample_shape=()):
@@ -563,6 +564,38 @@ def is_discrete(self):
         return self.support.is_discrete
 
 
+@runtime_checkable
+class DistributionLike(Protocol):
+    """A protocol for typing distributions.
+
+    Used to type object of type numpyro.distributions.Distribution, funsor.Funsor
+    or tensorflow_probability.distributions.Distribution.
+    """
+
+    @property
+    def batch_shape(self) -> tuple[int, ...]: ...
+
+    @property
+    def event_shape(self) -> tuple[int, ...]: ...
+
+    @property
+    def event_dim(self) -> int: ...
+
+    def sample(self, key: ArrayLike, sample_shape: tuple[int, ...] = ()) -> Array: ...
+
+    def log_prob(self, value: ArrayLike) -> ArrayLike: ...
+
+    @property
+    def mean(self) -> ArrayLike: ...
+
+    @property
+    def variance(self) -> ArrayLike: ...
+
+    def cdf(self, value: ArrayLike) -> ArrayLike: ...
+
+    def icdf(self, q: ArrayLike) -> ArrayLike: ...
+
+
 class ExpandedDistribution(Distribution):
     arg_constraints = {}
     pytree_data_fields = ("base_dist",)