Introduce beartype & fix types

examples/barycentres.pct.py

@@ -24,7 +24,9 @@
 import tensorflow_probability.substrates.jax.distributions as tfd
 from jax.config import config
 
-import gpjax as gpx
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
 
 # Enable Float64 for more stable matrix inversions.
 config.update("jax_enable_x64", True)

examples/classification.pct.py

@@ -19,7 +19,9 @@
 from jax.config import config
 from jaxtyping import Array, Float
 
-import gpjax as gpx
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
 
 # Enable Float64 for more stable matrix inversions.
 config.update("jax_enable_x64", True)

examples/collapsed_vi.pct.py

@@ -17,7 +17,9 @@
 from jax import jit
 from jax.config import config
 
-import gpjax as gpx
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
 
 # Enable Float64 for more stable matrix inversions.
 config.update("jax_enable_x64", True)

examples/deep_kernels.pct.py

@@ -9,9 +9,8 @@
 # Gaussian process model's kernel through a neural network can offer a solution to this.
 
 # %%
-import typing as tp
 from dataclasses import dataclass, field
-from typing import Dict, Any
+from typing import Any
 
 import jax
 import jax.numpy as jnp
@@ -25,12 +24,14 @@
 from simple_pytree import static_field
 import flax
 
-import gpjax as gpx
-import gpjax.kernels as jk
-from gpjax.kernels import DenseKernelComputation
-from gpjax.kernels.base import AbstractKernel
-from gpjax.kernels.computations import AbstractKernelComputation
-from gpjax.base import param_field
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
+    import gpjax.kernels as jk
+    from gpjax.kernels import DenseKernelComputation
+    from gpjax.kernels.base import AbstractKernel
+    from gpjax.kernels.computations import AbstractKernelComputation
+    from gpjax.base import param_field
 
 # Enable Float64 for more stable matrix inversions.
 config.update("jax_enable_x64", True)

examples/graph_kernels.pct.py

@@ -18,7 +18,9 @@
 from jax import jit
 from jax.config import config
 
-import gpjax as gpx
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
 
 # Enable Float64 for more stable matrix inversions.
 config.update("jax_enable_x64", True)
@@ -85,9 +87,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)
 

examples/kernels.pct.py

@@ -19,8 +19,10 @@
 from simple_pytree import static_field
 import numpy as np
 
-import gpjax as gpx
-from gpjax.base.param import param_field
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
+    from gpjax.base.param import param_field
 
 # Enable Float64 for more stable matrix inversions.
 config.update("jax_enable_x64", True)

examples/regression.pct.py

@@ -13,7 +13,9 @@
 from jax import jit
 from jax.config import config
 
-import gpjax as gpx
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
 
 # Enable Float64 for more stable matrix inversions.
 config.update("jax_enable_x64", True)

examples/spatial.pct.py

@@ -22,7 +22,6 @@
 
 import fsspec
 import geopandas as gpd
-import gpjax as gpx
 import jax
 import jax.numpy as jnp
 import jax.random as jr
@@ -33,11 +32,15 @@
 import pystac_client
 import rioxarray as rio
 import xarray as xr
-from gpjax.base import param_field
-from gpjax.dataset import Dataset
 from jaxtyping import Array, Float
 from rioxarray.merge import merge_arrays
 
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
+    from gpjax.base import param_field
+    from gpjax.dataset import Dataset
+
 jax.config.update("jax_enable_x64", True)
 
 key = jr.PRNGKey(123)

examples/uncollapsed_vi.pct.py

@@ -21,8 +21,10 @@
 from jax import jit
 from jax.config import config
 
-import gpjax as gpx
-import gpjax.kernels as jk
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
+    import gpjax.kernels as jk
 
 # Enable Float64 for more stable matrix inversions.
 config.update("jax_enable_x64", True)

examples/yacht.pct.py

@@ -10,7 +10,9 @@
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import StandardScaler
 
-import gpjax as gpx
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):
+    import gpjax as gpx
 
 # Enable Float64 for more stable matrix inversions.
 config.update("jax_enable_x64", True)

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[Num[Array, "N D"]] = None
+    y: Optional[Num[Array, "N Q"]] = 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[Num[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]: