Fix typo

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:

setup.py

@@ -17,7 +17,7 @@ def parse_requirements_file(filename):
     "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",
     "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