Fix jaxtyping and flax issues in tests

examples/constructing_new_kernels.py

@@ -92,7 +92,7 @@
 
 meanf = gpx.mean_functions.Zero()
 
-for k, ax, c in zip(kernels, axes.ravel(), cols):
+for k, ax, c in zip(kernels, axes.ravel(), cols, strict=False):
     prior = gpx.gps.Prior(mean_function=meanf, kernel=k)
     rv = prior(x)
     y = rv.sample(seed=key, sample_shape=(10,))

examples/intro_to_gps.py

@@ -214,7 +214,7 @@
 
 cmap = mpl.colors.LinearSegmentedColormap.from_list("custom", ["white", cols[1]], N=256)
 
-for a, t, d in zip([ax0, ax1, ax2], titles, dists):
+for a, t, d in zip([ax0, ax1, ax2], titles, dists, strict=False):
     d_prob = d.prob(jnp.hstack([xx.reshape(-1, 1), yy.reshape(-1, 1)])).reshape(
         xx.shape
     )

examples/intro_to_kernels.py

@@ -37,10 +37,7 @@
 from sklearn.preprocessing import StandardScaler
 
 from examples.utils import use_mpl_style
-from gpjax.parameters import (
-    PositiveReal,
-    Static,
-)
+from gpjax.parameters import Static
 from gpjax.typing import Array
 
 config.update("jax_enable_x64", True)
@@ -204,7 +201,7 @@
 
 meanf = gpx.mean_functions.Zero()
 
-for k, ax in zip(kernels, axes.ravel()):
+for k, ax in zip(kernels, axes.ravel(), strict=False):
     prior = gpx.gps.Prior(mean_function=meanf, kernel=k)
     rv = prior(x)
     y = rv.sample(seed=key, sample_shape=(10,))

examples/utils.py

@@ -70,7 +70,7 @@ def confidence_ellipse(x, y, ax, n_std=3.0, facecolor="none", **kwargs):
 
 def clean_legend(ax):
     handles, labels = ax.get_legend_handles_labels()
-    by_label = dict(zip(labels, handles))
+    by_label = dict(zip(labels, handles, strict=False))
     ax.legend(by_label.values(), by_label.keys())
     return ax
 

gpjax/distributions.py

@@ -162,7 +162,9 @@ def affine_transformation(x):
 
         return vmap(affine_transformation)(Z)
 
-    def sample(self, seed: KeyArray, sample_shape: Tuple[int, ...]):  # pylint: disable=useless-super-delegation
+    def sample(
+        self, seed: KeyArray, sample_shape: Tuple[int, ...]
+    ):  # pylint: disable=useless-super-delegation
         r"""See `Distribution.sample`."""
         return self._sample_n(
             seed, sample_shape[0]

tests/test_dataset.py

@@ -16,9 +16,10 @@
 from dataclasses import is_dataclass
 
 try:
-    import beartype
+    from beartype.roar import BeartypeCallHintParamViolation
+    from jaxtyping import TypeCheckError
 
-    ValidationErrors = (ValueError, beartype.roar.BeartypeCallHintParamViolation)
+    ValidationErrors = (TypeError, BeartypeCallHintParamViolation, TypeCheckError)
 except ImportError:
     ValidationErrors = ValueError
 

tests/test_decision_making/test_search_space.py

@@ -19,6 +19,7 @@
 from jaxtyping import (
     Array,
     Float,
+    TypeCheckError,
 )
 import pytest
 
@@ -64,7 +65,7 @@ def test_continuous_search_space_dtype_consistency(
 def test_continous_search_space_bounds_shape_consistency(
     lower_bounds: Float[Array, " D1"], upper_bounds: Float[Array, " D2"]
 ):
-    with pytest.raises((BeartypeCallHintParamViolation, ValueError)):
+    with pytest.raises((BeartypeCallHintParamViolation, TypeCheckError, ValueError)):
         ContinuousSearchSpace(lower_bounds=lower_bounds, upper_bounds=upper_bounds)
 
 

tests/test_gps.py

@@ -15,8 +15,9 @@
 
 try:
     from beartype.roar import BeartypeCallHintParamViolation
+    from jaxtyping import TypeCheckError
 
-    ValidationErrors = (TypeError, BeartypeCallHintParamViolation)
+    ValidationErrors = (TypeError, BeartypeCallHintParamViolation, TypeCheckError)
 except ImportError:
     ValidationErrors = TypeError
 

tests/test_parameters.py

@@ -39,7 +39,7 @@ def test_transform(param, value):
 
     # Test inverse transformation
     it_params = transform(t_params, DEFAULT_BIJECTION, inverse=True)
-    assert it_params == params
+    assert repr(it_params) == repr(params)
 
 
 @pytest.mark.parametrize(