Use RandomStream in docstring examples

aesara-devs · Dec 16, 2022 · d441294 · d441294
1 parent a94633c
commit d441294
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/aeppl/joint_logprob.py b/aeppl/joint_logprob.py
@@ -35,8 +35,10 @@ def conditional_logprob(
 
         import aesara.tensor as at
 
-        sigma2_rv = at.random.invgamma(0.5, 0.5, name="sigma2")
-        Y_rv = at.random.normal(0, at.sqrt(sigma2_rv), name="Y")
+        srng = at.random.RandomStream()
+
+        sigma2_rv = srng.invgamma(0.5, 0.5, name="sigma2")
+        Y_rv = srng.normal(0, at.sqrt(sigma2_rv), name="Y")
 
     This graph for ``Y_rv`` is equivalent to the following hierarchical model:
 
@@ -60,8 +62,10 @@ def conditional_logprob(
 
         import aesara.tensor as at
 
-        sigma2_rv = at.random.invgamma(0.5, 0.5)
-        Y_rv = at.random.normal(0, at.sqrt(sigma2_rv))
+        srng = at.random.RandomStream()
+
+        sigma2_rv = srng.invgamma(0.5, 0.5)
+        Y_rv = srng.normal(0, at.sqrt(sigma2_rv))
 
         logprobs, value_variables = conditional_logprob(Y_rv, sigma2_rv)
         # print(logprobs.keys())

diff --git a/aeppl/printing.py b/aeppl/printing.py
@@ -417,7 +417,8 @@ class PreamblePPrinter(PPrinter):
     -------
     >>> import aesara.tensor as at
     >>> from aeppl.printing import pprint
-    >>> X_rv = at.random.normal(at.scalar('\\mu'), at.scalar('\\sigma'), name='X')
+    >>> srng = at.random.RandomStream()
+    >>> X_rv = srng.normal(at.scalar('\\mu'), at.scalar('\\sigma'), name='X')
     >>> print(pprint(X_rv))
     \\mu in R
     \\sigma in R