Introduce beartype & fix types

examples/graph_kernels.pct.py

@@ -85,9 +85,9 @@
 
 true_kernel = gpx.GraphKernel(
     laplacian=L,
-    lengthscale=jnp.array([2.3]),
-    variance=jnp.array([3.2]),
-    smoothness=jnp.array([6.1]),
+    lengthscale=2.3,
+    variance=3.2,
+    smoothness=6.1,
 )
 prior = gpx.Prior(mean_function=gpx.Zero(), kernel=true_kernel)
 

gpjax/__init__.py

@@ -13,19 +13,21 @@
 # limitations under the License.
 # ==============================================================================
 
-from .dataset import Dataset
-from .fit import fit
-from .gps import Prior, construct_posterior
-from .kernels import *
-from .likelihoods import Bernoulli, Gaussian
-from .mean_functions import Constant, Zero
-from .objectives import (ELBO, CollapsedELBO, ConjugateMLL,
-                         LogPosteriorDensity, NonConjugateMLL)
-from .variational_families import (CollapsedVariationalGaussian,
-                                   ExpectationVariationalGaussian,
-                                   NaturalVariationalGaussian,
-                                   VariationalGaussian,
-                                   WhitenedVariationalGaussian)
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    from .dataset import Dataset
+    from .fit import fit
+    from .gps import Prior, construct_posterior
+    from .kernels import *
+    from .likelihoods import Bernoulli, Gaussian
+    from .mean_functions import Constant, Zero
+    from .objectives import (ELBO, CollapsedELBO, ConjugateMLL,
+                             LogPosteriorDensity, NonConjugateMLL)
+    from .variational_families import (CollapsedVariationalGaussian,
+                                       ExpectationVariationalGaussian,
+                                       NaturalVariationalGaussian,
+                                       VariationalGaussian,
+                                       WhitenedVariationalGaussian)
 
 __license__ = "MIT"
 __description__ = "Didactic Gaussian processes in JAX"

gpjax/base/module.py

@@ -13,14 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 __all__ = ["Module", "meta_leaves", "meta_flatten", "meta_map", "meta"]
 
 import dataclasses
 import os
 from copy import copy, deepcopy
-from typing import Any, Callable, Dict, Iterable, List, Tuple
+from beartype.typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, TypeVar, Union
 
 import jax
 import jax.tree_util as jtu
@@ -31,7 +30,9 @@
                               PyTreeCheckpointer, PyTreeCheckpointHandler,
                               RestoreArgs, SaveArgs)
 from simple_pytree import Pytree, static_field
-from typing_extensions import Self
+
+
+Self = TypeVar('Self')
 
 
 class Module(Pytree):
@@ -49,7 +50,7 @@ def __init_subclass__(cls, mutable: bool = False):
             ):
                 cls._pytree__meta[field] = {**value.metadata}
 
-    def replace(self, **kwargs: Any) -> Self:
+    def replace(self: Self, **kwargs: Any) -> Self:
         """
         Replace the values of the fields of the object.
 
@@ -68,7 +69,7 @@ def replace(self, **kwargs: Any) -> Self:
         pytree.__dict__.update(kwargs)
         return pytree
 
-    def replace_meta(self, **kwargs: Any) -> Self:
+    def replace_meta(self: Self, **kwargs: Any) -> Self:
         """
         Replace the metadata of the fields.
 
@@ -87,7 +88,7 @@ def replace_meta(self, **kwargs: Any) -> Self:
         pytree.__dict__.update(_pytree__meta={**pytree._pytree__meta, **kwargs})
         return pytree
 
-    def update_meta(self, **kwargs: Any) -> Self:
+    def update_meta(self: Self, **kwargs: Any) -> Self:
         """
         Update the metadata of the fields. The metadata must already exist.
 
@@ -112,15 +113,15 @@ def update_meta(self, **kwargs: Any) -> Self:
         pytree.__dict__.update(_pytree__meta=new)
         return pytree
 
-    def replace_trainable(self: Module, **kwargs: Dict[str, bool]) -> Self:
+    def replace_trainable(self: Self, **kwargs: Dict[str, bool]) -> Self:
         """Replace the trainability status of local nodes of the Module."""
         return self.update_meta(**{k: {"trainable": v} for k, v in kwargs.items()})
 
-    def replace_bijector(self: Module, **kwargs: Dict[str, tfb.Bijector]) -> Self:
+    def replace_bijector(self: Self, **kwargs: Dict[str, tfb.Bijector]) -> Self:
         """Replace the bijectors of local nodes of the Module."""
         return self.update_meta(**{k: {"bijector": v} for k, v in kwargs.items()})
 
-    def constrain(self) -> Self:
+    def constrain(self: Self) -> Self:
         """Transform model parameters to the constrained space according to their defined bijectors.
 
         Returns:
@@ -137,7 +138,7 @@ def _apply_constrain(meta_leaf):
 
         return meta_map(_apply_constrain, self)
 
-    def unconstrain(self) -> Self:
+    def unconstrain(self: Self) -> Self:
         """Transform model parameters to the unconstrained space according to their defined bijectors.
 
         Returns:
@@ -154,7 +155,7 @@ def _apply_unconstrain(meta_leaf):
 
         return meta_map(_apply_unconstrain, self)
 
-    def stop_gradient(self) -> Self:
+    def stop_gradient(self: Self) -> Self:
         """Stop gradients flowing through the Module.
 
         Returns:
@@ -176,7 +177,7 @@ def _apply_stop_grad(meta_leaf):
         return meta_map(_apply_stop_grad, self)
 
 
-def _toplevel_meta(pytree: Any) -> List[Dict[str, Any]]:
+def _toplevel_meta(pytree: Any) -> List[Optional[Dict[str, Any]]]:
     """Unpacks a list of meta corresponding to the top-level nodes of the pytree.
 
     Args:
@@ -197,8 +198,8 @@ def _toplevel_meta(pytree: Any) -> List[Dict[str, Any]]:
 def meta_leaves(
     pytree: Module,
     *,
-    is_leaf: Callable[[Any], bool] | None = None,
-) -> List[Tuple[Dict[str, Any], Any]]:
+    is_leaf: Optional[Callable[[Any], bool]] = None,
+) -> List[Tuple[Optional[Dict[str, Any]], Any]]:
     """
     Returns the meta of the leaves of the pytree.
 
@@ -212,8 +213,8 @@ def meta_leaves(
 
     def _unpack_metadata(
         meta_leaf: Any,
-        pytree: Module,
-        is_leaf: Callable[[Any], bool] | None,
+        pytree: Union[Module, Any],
+        is_leaf: Optional[Callable[[Any], bool]],
     ):
         """Recursively unpack leaf metadata."""
         if is_leaf and is_leaf(pytree):
@@ -235,8 +236,8 @@ def _unpack_metadata(
 
 
 def meta_flatten(
-    pytree: Module, *, is_leaf: Callable[[Any], bool] | None = None
-) -> Module:
+    pytree: Union[Module, Any], *, is_leaf: Optional[Callable[[Any], bool]] = None
+) -> Union[Module, Any]:
     """
     Returns the meta of the Module.
 
@@ -254,10 +255,10 @@ def meta_flatten(
 
 def meta_map(
     f: Callable[[Any, Dict[str, Any]], Any],
-    pytree: Module,
+    pytree: Union[Module, Any],
     *rest: Any,
-    is_leaf: Callable[[Any], bool] | None = None,
-) -> Module:
+    is_leaf: Optional[Callable[[Any], bool]] = None,
+) -> Union[Module, Any]:
     """Apply a function to a Module where the first argument are the pytree leaves, and the second argument are the Module metadata leaves.
     Args:
         f (Callable[[Any, Dict[str, Any]], Any]): The function to apply to the pytree.
@@ -273,7 +274,7 @@ def meta_map(
     return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
 
 
-def meta(pytree: Module, *, is_leaf: Callable[[Any], bool] | None = None) -> Module:
+def meta(pytree: Module, *, is_leaf: Optional[Callable[[Any], bool]] = None) -> Module:
     """Returns the metadata of the Module as a pytree.
 
     Args:

gpjax/base/param.py

@@ -13,12 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 __all__ = ["param_field"]
 
 import dataclasses
-from typing import Any, Mapping, Optional
+from beartype.typing import Any, Mapping, Optional
 
 import tensorflow_probability.substrates.jax.bijectors as tfb
 

gpjax/dataset.py

@@ -12,15 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Optional
+from beartype.typing import Optional, Union
 
 import jax.numpy as jnp
-from jaxtyping import Array, Float
+from jaxtyping import Float, Num
 from simple_pytree import Pytree
-
+from gpjax.typing import Array
 
 @dataclass
 class Dataset(Pytree):
@@ -31,8 +30,8 @@ class Dataset(Pytree):
         y (Optional[Float[Array, "N Q"]]): Output data.
     """
 
-    X: Optional[Float[Array, "N D"]] = None
-    y: Optional[Float[Array, "N Q"]] = None
+    X: Optional[Union[Num[Array, "N D"], Num[Array, "..."]]] = None
+    y: Optional[Union[Num[Array, "N Q"], Num[Array, "..."]]] = None
 
     def __post_init__(self) -> None:
         """Checks that the shapes of X and y are compatible."""
@@ -54,7 +53,7 @@ def is_unsupervised(self) -> bool:
         """Returns `True` if the dataset is unsupervised."""
         return self.X is None and self.y is not None
 
-    def __add__(self, other: Dataset) -> Dataset:
+    def __add__(self, other: "Dataset") -> "Dataset":
         """Combine two datasets. Right hand dataset is stacked beneath the left."""
 
         X = None
@@ -84,7 +83,7 @@ def out_dim(self) -> int:
         return self.y.shape[1]
 
 
-def _check_shape(X: Float[Array, "N D"], y: Float[Array, "N Q"]) -> None:
+def _check_shape(X: Optional[Num[Array, "..."]], y: Optional[Float[Array, "..."]]) -> None:
     """Checks that the shapes of X and y are compatible."""
     if X is not None and y is not None:
         if X.shape[0] != y.shape[0]:

gpjax/fit.py

@@ -13,17 +13,18 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Any, Optional, Tuple
+from beartype.typing import Any, Optional, Tuple, Union, Callable
 
 import jax
 import jax.random as jr
 import optax as ox
 from jax._src.random import _check_prng_key
-from jax.random import KeyArray
-from jaxtyping import Array, Float
+from gpjax.typing import Array
+from jaxtyping import Float
 from jaxlib.xla_extension import PjitFunction
 from warnings import warn
 
+from gpjax.typing import ScalarFloat, KeyArray
 from .base import Module
 from .dataset import Dataset
 from .objectives import AbstractObjective
@@ -33,7 +34,7 @@
 def fit(
     *,
     model: Module,
-    objective: AbstractObjective,
+    objective: Union[AbstractObjective, Callable[[Module, Dataset], ScalarFloat]],
     train_data: Dataset,
     optim: ox.GradientTransformation,
     num_iters: Optional[int] = 100,
@@ -69,7 +70,7 @@ def fit(
         >>> model = LinearModel(weight=1.0, bias=1.0)
         >>>
         >>> # (3) Define your loss function:
-        >>> class MeanSqaureError(gpx.AbstractObjective):
+        >>> class MeanSquareError(gpx.AbstractObjective):
         ...     def evaluate(self, model: LinearModel, train_data: gpx.Dataset) -> float:
         ...         return jnp.mean((train_data.y - model(train_data.X)) ** 2)
         ...
@@ -117,7 +118,7 @@ def fit(
         _check_verbose(verbose)
 
     # Unconstrained space loss function with stop-gradient rule for non-trainable params.
-    def loss(model: Module, batch: Dataset) -> Float[Array, "1"]:
+    def loss(model: Module, batch: Dataset) -> ScalarFloat:
         model = model.stop_gradient()
         return objective(model.constrain(), batch)