Merge branch 'aesara-devs:main' into ifelse-mixtures

Author: larryshamalama

README.rst

@@ -28,10 +28,11 @@ containing Aesara ``RandomVariable``\s:
 
   from aeppl import joint_logprob, pprint
 
+  srng = at.random.RandomStream()
 
   # A simple scale mixture model
-  S_rv = at.random.invgamma(0.5, 0.5)
-  Y_rv = at.random.normal(0.0, at.sqrt(S_rv))
+  S_rv = srng.invgamma(0.5, 0.5)
+  Y_rv = srng.normal(0.0, at.sqrt(S_rv))
 
   # Compute the joint log-probability
   logprob, (y, s) = joint_logprob(Y_rv, S_rv)
@@ -94,8 +95,8 @@ Joint log-probabilities can be computed for some terms that are *derived* from
 .. code-block:: python
 
   # Create a switching model from a Bernoulli distributed index
-  Z_rv = at.random.normal([-100, 100], 1.0, name="Z")
-  I_rv = at.random.bernoulli(0.5, name="I")
+  Z_rv = srng.normal([-100, 100], 1.0, name="Z")
+  I_rv = srng.bernoulli(0.5, name="I")
 
   M_rv = Z_rv[I_rv]
   M_rv.name = "M"

aeppl/joint_logprob.py

@@ -31,7 +31,8 @@ def conditional_logprob(
 
         import aesara.tensor as at
 
-        srng = at.random.RandomStream(0)
+        srng = at.random.RandomStream()
+
         sigma2_rv = srng.invgamma(0.5, 0.5)
         Y_rv = srng.normal(0, at.sqrt(sigma2_rv))
 
@@ -267,7 +268,7 @@ def joint_logprob(
 
         import aesara.tensor as at
 
-        srng = at.random.RandomStream(0)
+        srng = at.random.RandomStream()
         sigma2_rv = srng.invgamma(0.5, 0.5)
         Y_rv = srng.normal(0, at.sqrt(sigma2_rv))
 

aeppl/logprob.py

@@ -569,7 +569,7 @@ def categorical_logprob(op, values, *inputs, **kwargs):
             )
         )
         # FIXME: `take_along_axis` drops a broadcastable dimension
-        # when `value.broadcastable == p.broadcastable == (True, True, False)`.
+        # when `value.type.shape == p.type.shape == (1, 1, None)`.
     else:
         res = at.log(p[value])
 

aeppl/mixture.py

@@ -180,17 +180,17 @@ def rv_pull_down(x: TensorVariable, dont_touch_vars=None) -> TensorVariable:
 class MixtureRV(Op):
     """A placeholder used to specify a log-likelihood for a mixture sub-graph."""
 
-    __props__ = ("indices_end_idx", "out_dtype", "out_broadcastable")
+    __props__ = ("indices_end_idx", "out_dtype", "out_shape")
 
-    def __init__(self, indices_end_idx, out_dtype, out_broadcastable):
+    def __init__(self, indices_end_idx, out_dtype, out_shape):
         super().__init__()
         self.indices_end_idx = indices_end_idx
         self.out_dtype = out_dtype
-        self.out_broadcastable = out_broadcastable
+        self.out_shape = out_shape
 
     def make_node(self, *inputs):
         return Apply(
-            self, list(inputs), [TensorType(self.out_dtype, self.out_broadcastable)()]
+            self, list(inputs), [TensorType(self.out_dtype, shape=self.out_shape)()]
         )
 
     def perform(self, node, inputs, outputs):
@@ -284,8 +284,8 @@ def mixture_replace(fgraph, node):
     # Replace this sub-graph with a `MixtureRV`
     mix_op = MixtureRV(
         1 + len(mixing_indices),
-        old_mixture_rv.dtype,
-        old_mixture_rv.broadcastable,
+        old_mixture_rv.type.dtype,
+        old_mixture_rv.type.shape,
     )
     new_node = mix_op.make_node(*([join_axis] + mixing_indices + mixture_rvs))
 
@@ -364,8 +364,8 @@ def ifelse_mixture_replace(fgraph, node):
     """
     mix_op = MixtureRV(
         2,
-        old_mixture_rv.dtype,
-        old_mixture_rv.broadcastable,
+        old_mixture_rv.type.dtype,
+        old_mixture_rv.type.shape,
     )
 
     if node.inputs[0].ndim == 0:

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

docs/source/api/distributions.rst

@@ -11,7 +11,7 @@ The :py:func:`aeppl.logprob.logprob` function can be called on any random variab
    import aesara.tensor as at
    from aeppl.logprob import _logprob
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.scalar("mu")
    sigma = at.scalar("sigma")
@@ -29,7 +29,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    p = at.scalar("p")
    x_rv = snrg.bernoulli(p)
@@ -43,7 +43,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    a = at.scalar("a")
    b = at.scalar("b")
@@ -59,7 +59,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    n = at.iscalar("n")
    a = at.scalar("a")
@@ -76,7 +76,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    n = at.iscalar("n")
    p = at.scalar("p")
@@ -92,7 +92,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    loc = at.scalar("loc")
    scale = at.scalar("scale")
@@ -107,7 +107,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    p = at.vector("p")
    x_rv = snrg.categorical(p)
@@ -121,7 +121,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    df = at.scalar("df")
    x_rv = snrg.chisquare(df)
@@ -148,7 +148,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    alpha = at.vector("alpha")
    x_rv = snrg.dirichlet(alpha)
@@ -167,7 +167,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    beta = at.scalar("beta")
    x_rv = snrg.exponential(beta)
@@ -181,7 +181,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    alpha = at.scalar('alpha')
    beta = at.scalar('beta')
@@ -196,7 +196,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    p = at.scalar("p")
    x_rv = snrg.geometric(p)
@@ -210,7 +210,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.scalar('mu')
    beta = at.scalar('beta')
@@ -225,7 +225,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    x0 = at.scalar('x0')
    gamma = at.scalar('gamma')
@@ -240,7 +240,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.scalar('mu')
    sigma = at.scalar('sigma')
@@ -255,7 +255,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    ngood = at.scalar("ngood")
    nbad = at.scalar("nbad")
@@ -271,7 +271,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    alpha = at.scalar('alpha')
    beta = at.scalar('beta')
@@ -286,7 +286,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.scalar("mu")
    lmbda = at.scalar("lambda")
@@ -301,7 +301,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.scalar("mu")
    s = at.scalar("s")
@@ -316,7 +316,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.scalar("mu")
    sigma = at.scalar("sigma")
@@ -331,7 +331,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    n = at.iscalar("n")
    p = at.vector("p")
@@ -346,7 +346,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.vector('mu')
    Sigma = at.matrix('sigma')
@@ -362,7 +362,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    n = at.iscalar("n")
    p = at.scalar("p")
@@ -377,7 +377,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.scalar('mu')
    sigma = at.scalar('sigma')
@@ -392,7 +392,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    b = at.scalar("b")
    scale = at.scalar("scale")
@@ -407,7 +407,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    lmbda = at.scalar("lambda")
    x_rv = snrg.poisson(lmbda)
@@ -421,7 +421,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    df = at.scalar('df')
    loc = at.scalar('loc')
@@ -437,7 +437,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    left = at.scalar('left')
    mode = at.scalar('mode')
@@ -453,7 +453,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    low = at.scalar('low')
    high = at.scalar('high')
@@ -468,7 +468,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.scalar('mu')
    kappa = at.scalar('kappa')
@@ -483,7 +483,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    mu = at.scalar('mu')
    lmbda = at.scalar('lambda')
@@ -499,7 +499,7 @@ Documentation for the Aesara implementation can be found here: :external:py:clas
 
    import aesara.tensor as at
 
-   srng = at.random.RandomStream(0)
+   srng = at.random.RandomStream()
 
    k = at.scalar('k')
    x_rv = srng.weibull(k)