Proof of Concept: Types and MyPy (#1906)

* mypy init got hsgp module

* mypy some modules

* rename step

* fix missing default

* update pre-commit and small typo (caught by codespell)

.github/workflows/ci.yml

@@ -32,7 +32,7 @@ jobs:
         pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip install -r docs/requirements.txt
         pip freeze
-    - name: Lint with ruff
+    - name: Lint with mypy and ruff
       run: |
         make lint
     - name: Build documentation

.pre-commit-config.yaml

@@ -14,6 +14,13 @@ repos:
         language: system
         files: "(.py$)|(.*.ipynb$)"
 
+      - id: mypy
+        name: mypy
+        language: python
+        entry: mypy --install-types --non-interactive
+        files: ^numpyro/
+
+
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0
     hooks:

Makefile

@@ -4,6 +4,7 @@ lint: FORCE
 	ruff check .
 	ruff format . --check
 	python scripts/update_headers.py --check
+	mypy --install-types --non-interactive numpyro
 
 license: FORCE
 	python scripts/update_headers.py

numpyro/contrib/control_flow/cond.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from functools import partial
+from typing import Any, Callable
 
 from jax import device_put, lax
 
@@ -72,7 +73,7 @@ def cond_wrapper(
     return lax.cond(pred, wrapped_true_fun, wrapped_false_fun, wrapped_operand)
 
 
-def cond(pred, true_fun, false_fun, operand):
+def cond(pred: bool, true_fun: Callable, false_fun: Callable, operand: Any) -> Any:
     """
     This primitive conditionally applies ``true_fun`` or ``false_fun``. See
     :func:`jax.lax.cond` for more information.

numpyro/contrib/control_flow/scan.py

@@ -3,6 +3,7 @@
 
 from collections import OrderedDict
 from functools import partial
+from typing import Callable
 
 import jax
 from jax import device_put, lax, random
@@ -278,14 +279,17 @@ def scan_wrapper(
     length,
     reverse,
     rng_key=None,
-    substitute_stack=[],
+    substitute_stack=None,
     enum=False,
     history=1,
     first_available_dim=None,
 ):
     if length is None:
         length = jnp.shape(jax.tree.flatten(xs)[0][0])[0]
 
+    if substitute_stack is None:
+        substitute_stack = []
+
     if enum and history > 0:
         return scan_enum(  # TODO: replay for enum
             f,
@@ -339,7 +343,14 @@ def body_fn(wrapped_carry, x):
     return last_carry, (pytree_trace, ys)
 
 
-def scan(f, init, xs, length=None, reverse=False, history=1):
+def scan(
+    f: Callable,
+    init,
+    xs,
+    length: int | None = None,
+    reverse: bool = False,
+    history: int = 1,
+):
     """
     This primitive scans a function over the leading array axes of
     `xs` while carrying along state. See :func:`jax.lax.scan` for more
@@ -433,7 +444,7 @@ def g(*args, **kwargs):
     :param init: the initial carrying state
     :param xs: the values over which we scan along the leading axis. This can
         be any JAX pytree (e.g. list/dict of arrays).
-    :param length: optional value specifying the length of `xs`
+    :param int | None length: optional value specifying the length of `xs`
         but can be used when `xs` is an empty pytree (e.g. None)
     :param bool reverse: optional boolean specifying whether to run the scan iteration
         forward (the default) or in reverse

numpyro/contrib/hsgp/approximation.py

@@ -21,14 +21,18 @@
 import numpyro.distributions as dist
 
 
-def _non_centered_approximation(phi: ArrayImpl, spd: ArrayImpl, m: int) -> ArrayImpl:
+def _non_centered_approximation(
+    phi: ArrayImpl, spd: ArrayImpl, m: int | list[int]
+) -> ArrayImpl:
     with numpyro.plate("basis", m):
         beta = numpyro.sample("beta", dist.Normal(loc=0.0, scale=1.0))
 
     return phi @ (spd * beta)
 
 
-def _centered_approximation(phi: ArrayImpl, spd: ArrayImpl, m: int) -> ArrayImpl:
+def _centered_approximation(
+    phi: ArrayImpl, spd: ArrayImpl, m: int | list[int]
+) -> ArrayImpl:
     with numpyro.plate("basis", m):
         beta = numpyro.sample("beta", dist.Normal(loc=0.0, scale=spd))
 

numpyro/contrib/hsgp/spectral_densities.py

@@ -22,7 +22,7 @@ def align_param(dim, param):
 
 def spectral_density_squared_exponential(
     dim: int, w: ArrayImpl, alpha: float, length: float | ArrayImpl
-) -> float:
+) -> ArrayImpl:
     """
     Spectral density of the squared exponential kernel.
 
@@ -46,7 +46,7 @@ def spectral_density_squared_exponential(
     :param float alpha: amplitude
     :param float length: length scale
     :return: spectral density value
-    :rtype: float
+    :rtype: ArrayImpl
     """
     length = align_param(dim, length)
     c = alpha * jnp.prod(jnp.sqrt(2 * jnp.pi) * length, axis=-1)

numpyro/distributions/continuous.py

@@ -600,7 +600,7 @@ def __init__(
             transition_matrix.ndim == 2
         ), "`transition_matrix` argument should be a square matrix"
         self.transition_matrix = transition_matrix
-        # Expand the covariance/presicion/scale matrices to the right number of steps.
+        # Expand the covariance/precision/scale matrices to the right number of steps.
         args = {
             "covariance_matrix": covariance_matrix,
             "precision_matrix": precision_matrix,

numpyro/infer/util.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections import namedtuple
+from collections.abc import Sequence
 from contextlib import contextmanager
 from functools import partial
 from typing import Callable, Optional
@@ -931,7 +932,7 @@ def __init__(
         guide: Optional[Callable] = None,
         params: Optional[dict] = None,
         num_samples: Optional[int] = None,
-        return_sites: Optional[list[str]] = None,
+        return_sites: Optional[Sequence[str]] = None,
         infer_discrete: bool = False,
         parallel: bool = False,
         batch_ndims: Optional[int] = None,

pyproject.toml

@@ -102,3 +102,15 @@ doctest_optionflags = [
     "NORMALIZE_WHITESPACE",
     "IGNORE_EXCEPTION_DETAIL",
 ]
+
+[tool.mypy]
+ignore_errors = true
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = [
+    "numpyro.contrib.control_flow.*", # types missing
+    "numpyro.contrib.funsor.*",       # types missing
+    "numpyro.contrib.hsgp.*",
+]
+ignore_errors = false

setup.py

@@ -51,6 +51,7 @@
         "test": [
             "importlib-metadata<5.0",
             "ruff>=0.1.8",
+            "mypy>=1.13",
             "pytest>=4.1",
             "pyro-api>=0.1.1",
             "scikit-learn",