Merge pull request #103 from thomaspinder/fix_deps

Relax Jax version

Author: daniel-dodd

gpjax/config.py

@@ -5,8 +5,31 @@
 
 __config = None
 
-Identity = dx.Lambda(lambda x: x)
-Softplus = dx.Lambda(lambda x: jnp.log(1.0 + jnp.exp(x)))
+Identity = dx.Lambda(forward=lambda x: x, inverse=lambda x: x)
+Softplus = dx.Lambda(
+    forward=lambda x: jnp.log(1 + jnp.exp(x)),
+    inverse=lambda x: jnp.log(jnp.exp(x) - 1.0),
+)
+
+# class Softplus(dx.Bijector):
+#     def __init__(self):
+#         super().__init__(event_ndims_in=0)
+
+#     def forward_and_log_det(self, x):
+#         softplus = lambda xx: jnp.log(1 + jnp.exp(xx))
+#         y = softplus(x)
+#         logdet = softplus(-x)
+#         return y, logdet
+
+#     def inverse_and_log_det(self, y):
+#         """
+#         Y = Log[1 + exp{X}] ==> X = Log[exp{Y} - 1]
+#         ==> dX/dY = exp{Y} / (exp{Y} - 1)
+#                   = 1 / (1 - exp{-Y})
+#         """
+#         x = jnp.log(jnp.exp(y) - 1.0)
+#         logdet = 1 / (1 - jnp.exp(-y))
+#         return x, logdet
 
 
 def get_defaults() -> ConfigDict:

gpjax/parameters.py

@@ -10,13 +10,12 @@
 import jax.random as jr
 from chex import dataclass
 from jaxtyping import f64
-from tensorflow_probability.substrates.jax import distributions as tfd
 
 from .config import get_defaults
 from .types import PRNGKeyType
 from .utils import merge_dictionaries
 
-Identity = dx.Lambda(lambda x: x)
+Identity = dx.Lambda(forward=lambda x: x, inverse=lambda x: x)
 
 
 ################################
@@ -163,11 +162,13 @@ def inverse(bijector):
 
     bijectors = build_bijectors(params)
 
-    constrainers = jax.tree_map(lambda _: forward, deepcopy(params))
-    unconstrainers = jax.tree_map(lambda _: inverse, deepcopy(params))
+    constrainers = jax.tree_util.tree_map(lambda _: forward, deepcopy(params))
+    unconstrainers = jax.tree_util.tree_map(lambda _: inverse, deepcopy(params))
 
-    constrainers = jax.tree_map(lambda f, b: f(b), constrainers, bijectors)
-    unconstrainers = jax.tree_map(lambda f, b: f(b), unconstrainers, bijectors)
+    constrainers = jax.tree_util.tree_map(lambda f, b: f(b), constrainers, bijectors)
+    unconstrainers = jax.tree_util.tree_map(
+        lambda f, b: f(b), unconstrainers, bijectors
+    )
 
     return constrainers, unconstrainers
 
@@ -182,7 +183,9 @@ def transform(params: tp.Dict, transform_map: tp.Dict) -> tp.Dict:
     Returns:
         tp.Dict: A transformed parameter set.s The dictionary is equal in structure to the input params dictionary.
     """
-    return jax.tree_map(lambda param, trans: trans(param), params, transform_map)
+    return jax.tree_util.tree_map(
+        lambda param, trans: trans(param), params, transform_map
+    )
 
 
 ################################
@@ -200,7 +203,7 @@ def copy_dict_structure(params: dict) -> dict:
     # Copy dictionary structure
     prior_container = deepcopy(params)
     # Set all values to zero
-    prior_container = jax.tree_map(lambda _: None, prior_container)
+    prior_container = jax.tree_util.tree_map(lambda _: None, prior_container)
     return prior_container
 
 
@@ -243,23 +246,16 @@ def prior_checks(priors: dict) -> dict:
     """Run checks on th parameters' prior distributions. This checks that for Gaussian processes that are constructed with non-conjugate likelihoods, the prior distribution on the function's latent values is a unit Gaussian."""
     if "latent" in priors.keys():
         latent_prior = priors["latent"]
-        if isinstance(latent_prior, dx.Distribution) and latent_prior.name != "Normal":
-            warnings.warn(
-                f"A {latent_prior.name} distribution prior has been placed on"
-                " the latent function. It is strongly advised that a"
-                " unit-Gaussian prior is used."
-            )
-        elif (
-            isinstance(latent_prior, tfd.Distribution) and latent_prior.name != "Normal"
-        ):
-            warnings.warn(
-                f"A {latent_prior.name} distribution from Tensorflow Probability has been"
-                "placed on the latent function. We advise using a unit-Gaussian prior from"
-                " Distrax."
-            )
+        if latent_prior is not None:
+            if latent_prior.name != "Normal":
+                warnings.warn(
+                    f"A {latent_prior.name} distribution prior has been placed on"
+                    " the latent function. It is strongly advised that a"
+                    " unit Gaussian prior is used."
+                )
         else:
-            if not latent_prior:
-                priors["latent"] = dx.Normal(loc=0.0, scale=1.0)
+            warnings.warn("Placing unit Gaussian prior on latent function.")
+            priors["latent"] = dx.Normal(loc=0.0, scale=1.0)
     else:
         priors["latent"] = dx.Normal(loc=0.0, scale=1.0)
 
@@ -278,7 +274,7 @@ def build_trainables(params: tp.Dict) -> tp.Dict:
     # Copy dictionary structure
     prior_container = deepcopy(params)
     # Set all values to zero
-    prior_container = jax.tree_map(lambda _: True, prior_container)
+    prior_container = jax.tree_util.tree_map(lambda _: True, prior_container)
     return prior_container
 
 
@@ -289,6 +285,6 @@ def stop_grad(param: tp.Dict, trainable: tp.Dict):
 
 def trainable_params(params: tp.Dict, trainables: tp.Dict) -> tp.Dict:
     """Stop the gradients flowing through parameters whose trainable status is False"""
-    return jax.tree_map(
+    return jax.tree_util.tree_map(
         lambda param, trainable: stop_grad(param, trainable), params, trainables
     )

setup.py

@@ -12,15 +12,14 @@ def parse_requirements_file(filename):
 
 
 REQUIRES = [
-    "jax==0.3.5",
-    "jaxlib==0.3.5",
-    "optax>=0.1.0",
-    "chex==0.1.3",
+    "jax>=0.1.67",
+    "jaxlib>=0.1.47",
+    "optax",
+    "chex",
     "distrax>=0.1.2",
-    "tensorflow-probability==0.16.0",
+    "tensorflow-probability>=0.16.0",
     "tqdm>=4.0.0",
     "ml-collections==0.1.0",
-    "protobuf==3.19.0",
     "jaxtyping",
 ]
 

tests/test_parameters.py

@@ -253,5 +253,5 @@ def test_output(num_datapoints, likelihood):
     a_constrainers, a_unconstrainers = build_transforms(augmented_params)
     assert "test_param" in list(a_constrainers.keys())
     assert "test_param" in list(a_unconstrainers.keys())
-    assert a_constrainers["test_param"](1.0) == 1.0
-    assert a_unconstrainers["test_param"](1.0) == 1.0
+    assert a_constrainers["test_param"](jnp.array([1.0])) == 1.0
+    assert a_unconstrainers["test_param"](jnp.array([1.0])) == 1.0