Update joint_logprob usage for newer AePPL

aesara-devs · Jan 3, 2023 · b8b9641 · b8b9641
1 parent 96ea86d
commit b8b9641
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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