Repository: jaxgaussianprocesses/gpjax
Issue: bug: Installation with GPU support

# Bug Report

Moving this from `Discussions` as not seeing much traffic there. 
**GPJax version: 0.4.10**

**Current behavior:**
I can install GPJax with CPU only support but not with GPU support. If I try to install gpjax with cuda support, either based on the current gpjax readthedocs or by first installing Jax according to the latest install with cuda directions:

- installation page for GPJax: https://gpjax.readthedocs.io/en/latest/installation.html
- also tried to follow the Jax install guide: https://github.com/google/jax#installation

I've tried to find a matching CUDA-compatible version of `jax==0.3.5` but haven't found a version that is also supported by my machine.

I've also tried installing it on a Google Colab notebook with `!pip install gpjax` and this disables GPU support based on the package versions specified in setup.py.

I'm thinking the setup.py file needs to be updated to a higher jax version and potentially make corresponding updates to changes in jax since 0.3.5, or I'm misunderstanding how to configure GPU support with GPJax.

**Expected behavior:**
I would expect that GPJax does have GPU support. 

**Steps to reproduce:**

Ubuntu 20.04.3 and have tried python 3.8-3.10 
run the regression example for instance:
```
from pprint import PrettyPrinter

import gpjax as gpx
import jax
import jax.numpy as jnp
import jax.random as jr
from jax.example_libraries import optimizers
from jax import jit
import optax as ox

pp = PrettyPrinter(indent=4)
key = jr.PRNGKey(123)


def pred_dist(xtest, y, training, final_params):
    latent_distribution = posterior(training, final_params)(xtest)
    predictive_distribution = likelihood(latent_distribution, final_params)
    return jnp.sum(predictive_distribution.log_prob(y))

if __name__ == '__main__':
    import matplotlib.pyplot as plt

    n = 100
    noise = 0.3

    x = jr.uniform(key=key, minval=-3.0, maxval=3.0, shape=(n,)).sort().reshape(-1, 1)
    f = lambda x: jnp.sin(4 * x) + jnp.cos(2 * x)
    signal = f(x)
    y = signal + jr.normal(key, shape=signal.shape) * noise

    D = gpx.Dataset(X=x, y=y)

    xtest = jnp.linspace(-3.5, 3.5, 500).reshape(-1, 1)
    ytest = f(xtest)

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(xtest, ytest, label="Latent function")
    ax.plot(x, y, "o", label="Observations")
    ax.legend(loc="best")

    kernel = gpx.RBF()
    prior = gpx.Prior(kernel=kernel)
    likelihood = gpx.Gaussian(num_datapoints=D.n)
    posterior = prior * likelihood
    params, trainable, constrainer, unconstrainer = gpx.initialise(posterior)
    pp.pprint(params)
    params = gpx.transform(params, unconstrainer)
    mll = jit(posterior.marginal_log_likelihood(D, constrainer, negative=True))
    mll(params)

    opt = ox.adam(learning_rate=0.01)
    final_params = gpx.fit(
        mll,
        params,
        trainable,
        opt,
        n_iters=500,
    )

    final_params = gpx.transform(final_params, constrainer)
    pp.pprint(final_params)

    latent_dist = posterior(D, final_params)(xtest)
    predictive_dist = likelihood(latent_dist, final_params)

    predictive_mean = predictive_dist.mean()
    predictive_std = predictive_dist.stddev()
    # test gradient behavior with respect to posterior
    gtest = jax.grad(pred_dist)
    grad = gtest(xtest, y, D, final_params)
```
Running this returns the following warning and error:
```
/home/paperspace/anaconda3/envs/ml310/lib/python3.10/site-packages/chex-0.1.3-py3.10.egg/chex/_src/pytypes.py:37: FutureWarning: jax.tree_structure is deprecated, and will be removed in a future release. Use jax.tree_util.tree_structure instead.
  PyTreeDef = type(jax.tree_structure(None))
{   'kernel': {   'lengthscale': DeviceArray([1.], dtype=float64),
                  'variance': DeviceArray([1.], dtype=float64)},
    'likelihood': {'obs_noise': DeviceArray([1.], dtype=float64)},
    'mean_function': {}}
Traceback (most recent call last):
  File "/home/paperspace/pycharm_projects/exp_design/gpjax_demo.py", line 47, in <module>
    params = gpx.transform(params, unconstrainer)
  File "/home/paperspace/GPJax/gpjax/parameters.py", line 147, in transform
    return jax.tree_util.tree_map(lambda param, trans: trans(param), params, transform_map)
  File "/home/paperspace/anaconda3/envs/ml310/lib/python3.10/site-packages/jax/_src/tree_util.py", line 205, in tree_map
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/home/paperspace/anaconda3/envs/ml310/lib/python3.10/site-packages/jax/_src/tree_util.py", line 205, in <genexpr>
    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
  File "/home/paperspace/GPJax/gpjax/parameters.py", line 147, in <lambda>
    return jax.tree_util.tree_map(lambda param, trans: trans(param), params, transform_map)
  File "/home/paperspace/anaconda3/envs/ml310/lib/python3.10/site-packages/distrax-0.1.2-py3.10.egg/distrax/_src/bijectors/lambda_bijector.py", line 112, in inverse
    return self._inverse(y)
  File "/home/paperspace/anaconda3/envs/ml310/lib/python3.10/site-packages/distrax-0.1.2-py3.10.egg/distrax/_src/utils/transformations.py", line 119, in wrapped
    out = _interpret_inverse(jaxpr, consts, *args)
  File "/home/paperspace/anaconda3/envs/ml310/lib/python3.10/site-packages/distrax-0.1.2-py3.10.egg/distrax/_src/utils/transformations.py", line 250, in _interpret_inverse
    write(jax.core.unitvar, jax.core.unit)
AttributeError: module 'jax.core' has no attribute 'unitvar'
```
I can try to provide something more minimal, but I'm guessing I'm just missing something more fundamental on install.

In Colab, run 
`!pip install gpjax`
Then reload as per the prompt
`import gpjax as gpx`
At that point it should indicate their is no GPU/TPU found

Thanks!


Changed files:
gpjax/parameters.py, tests/test_parameters.py, setup.py, gpjax/config.py

diff --git a/gpjax/config.py b/gpjax/config.py
index a953dd3..13695d9 100644
--- a/gpjax/config.py
+++ b/gpjax/config.py
@@ -5,8 +5,31 @@ from ml_collections import ConfigDict
 
 __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:
diff --git a/gpjax/parameters.py b/gpjax/parameters.py
index bda0590..7daf6c6 100644
--- a/gpjax/parameters.py
+++ b/gpjax/parameters.py
@@ -10,13 +10,12 @@ import jax.numpy as jnp
 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 build_transforms(params: tp.Dict) -> tp.Tuple[tp.Dict, tp.Dict]:
 
     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
     )
diff --git a/setup.py b/setup.py
index 34c36cc..6cae770 100644
--- a/setup.py
+++ b/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",
 ]
 
diff --git a/tests/test_parameters.py b/tests/test_parameters.py
index 52a26c3..c0ce0c0 100644
--- a/tests/test_parameters.py
+++ b/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