Fix tests with jax 0.5.1 (#494)

* Bump tolerance on expected improvement utility

* Relax get_batch uniqueness test

* Surpress docstrings

* Bound Flax

* Deprecate QuadContourSet collections

* Bump version

* Bound mkdocstrings

---------

Co-authored-by: Thomas Pinder <[email protected]>

Author: stephen-huan

examples/intro_to_gps.py

@@ -218,9 +218,7 @@
     d_prob = d.prob(jnp.hstack([xx.reshape(-1, 1), yy.reshape(-1, 1)])).reshape(
         xx.shape
     )
-    cntf = a.contourf(xx, yy, jnp.exp(d_prob), levels=20, antialiased=True, cmap=cmap)
-    for c in cntf.collections:
-        c.set_edgecolor("face")
+    cntf = a.contourf(xx, yy, jnp.exp(d_prob), levels=20, antialiased=True, cmap=cmap, edgecolor="face")
     a.set_xlim(-2.75, 2.75)
     a.set_ylim(-2.75, 2.75)
     samples = d.sample(seed=key, sample_shape=(5000,))

gpjax/__init__.py

@@ -40,7 +40,7 @@
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.9.3"
+__version__ = "0.9.4"
 
 __all__ = [
     "base",

pyproject.toml

@@ -30,7 +30,7 @@ dependencies = [
   "beartype>0.16.1",
   "cola-ml==0.0.5",
   "jaxopt==0.8.2",
-  "flax>=0.8.5",
+  "flax<0.10.0",
   "numpy<2.0.0",
 ]
 
@@ -42,7 +42,7 @@ python = "3.10"
 dependencies = [
   "mkdocs>=1.5.3",
   "mkdocs-material>=9.5.12",
-  "mkdocstrings[python]>=0.25.1",
+  "mkdocstrings[python]<0.28.0",
   "mkdocs-jupyter>=0.24.3",
   "mkdocs-gen-files>=0.5.0",
   "mkdocs-literate-nav>=0.6.0",

tests/test_decision_making/test_utility_functions/test_expected_improvement.py

@@ -64,4 +64,4 @@ def test_expected_improvement_utility_function_correct_values(
     eta = get_best_latent_observation_val(posterior, dataset)
     mc_ei = jnp.expand_dims(jnp.mean(jnp.maximum(eta - samples, 0), 0), -1)
     assert jnp.all(ei >= 0)
-    assert jnp.allclose(ei, mc_ei, rtol=0.01)
+    assert jnp.allclose(ei, mc_ei, rtol=0.08, atol=1e-6)

tests/test_fit.py

@@ -322,5 +322,5 @@ def test_get_batch(n_data: int, n_dim: int, batch_size: int):
     assert New.n == batch_size
     assert New.X.shape[1:] == x.shape[1:]
     assert New.y.shape[1:] == y.shape[1:]
-    assert (New.X != B.X).all()
-    assert (New.y != B.y).all()
+    assert jnp.sum(New.X == B.X) <= n_dim * batch_size / n_data
+    assert jnp.sum(New.y == B.y) <= n_dim * batch_size / n_data