diff --git a/README.rst b/README.rst
index e838dbc..c9f1f15 100644
--- a/README.rst
+++ b/README.rst
@@ -25,7 +25,7 @@ Example
 
 
     def logprob_fn(y):
-        return joint_logprob({Y_rv: y})
+        return joint_logprob(realized={Y_rv: y})[0]
 
 
     # Build the transition kernel
diff --git a/environment.yml b/environment.yml
index 7220062..0afadb9 100644
--- a/environment.yml
+++ b/environment.yml
@@ -12,7 +12,7 @@ dependencies:
   - numpy>=1.18.1
   - scipy>=1.4.0
   - aesara>=2.8.3
-  - aeppl>=0.0.38
+  - aeppl>=0.0.40
   # Intel BLAS
   - mkl
   - mkl-service
diff --git a/setup.py b/setup.py
index 05c9c0b..d6ace31 100644
--- a/setup.py
+++ b/setup.py
@@ -18,7 +18,7 @@
         "numpy>=1.18.1",
         "scipy>=1.4.0",
         "aesara>=2.8.3",
-        "aeppl>=0.0.38",
+        "aeppl>=0.0.40",
     ],
     tests_require=["pytest"],
     long_description=open("README.md").read() if exists("README.md") else "",
diff --git a/tests/test_hmc.py b/tests/test_hmc.py
index 2758497..61cdc17 100644
--- a/tests/test_hmc.py
+++ b/tests/test_hmc.py
@@ -17,7 +17,7 @@ def test_warmup_scalar():
     Y_rv = srng.normal(1, 2)
 
     def logprob_fn(y: TensorVariable):
-        logprob = joint_logprob({Y_rv: y})
+        logprob, _ = joint_logprob(realized={Y_rv: y})
         return logprob
 
     y_vv = Y_rv.clone()
@@ -57,7 +57,7 @@ def test_warmup_vector():
     Y_rv = srng.multivariate_normal(loc, cov)
 
     def logprob_fn(y: TensorVariable):
-        logprob = joint_logprob({Y_rv: y})
+        logprob, _ = joint_logprob(realized={Y_rv: y})
         return logprob
 
     y_vv = Y_rv.clone()
@@ -104,7 +104,7 @@ def test_univariate_hmc(step_size, diverges):
     Y_rv = srng.normal(1, 2)
 
     def logprob_fn(y):
-        logprob = joint_logprob({Y_rv: y})
+        logprob, _ = joint_logprob(realized={Y_rv: y})
         return logprob
 
     kernel = hmc.new_kernel(srng, logprob_fn)
@@ -162,7 +162,7 @@ def multivariate_normal_model(srng):
     Y_rv = srng.multivariate_normal(loc_tt, cov_tt)
 
     def logprob_fn(y):
-        return joint_logprob({Y_rv: y})
+        return joint_logprob(realized={Y_rv: y})[0]
 
     return (loc, scale, rho), Y_rv, logprob_fn