Convert RandomVariables to in-place during graph optimization

Author: brandonwillard

pymc3/aesaraf.py

@@ -31,6 +31,7 @@
 import scipy.sparse as sps
 
 from aesara import config, scalar
+from aesara.compile.mode import Mode, get_mode
 from aesara.gradient import grad
 from aesara.graph.basic import (
     Apply,
@@ -861,3 +862,16 @@ def take_along_axis(arr, indices, axis=0):
 
     # use the fancy index
     return arr[_make_along_axis_idx(arr_shape, indices, _axis)]
+
+
+def compile_rv_inplace(inputs, outputs, mode=None, **kwargs):
+    """Use ``aesara.function`` with the random_make_inplace optimization always enabled.
+
+    Using this function ensures that compiled functions containing random
+    variables will produce new samples on each call.
+    """
+    mode = get_mode(mode)
+    opt_qry = mode.provided_optimizer.including("random_make_inplace")
+    mode = Mode(linker=mode.linker, optimizer=opt_qry)
+    aesara_function = aesara.function(inputs, outputs, mode=mode, **kwargs)
+    return aesara_function

pymc3/distributions/distribution.py

@@ -19,12 +19,12 @@
 import warnings
 
 from abc import ABCMeta
-from copy import copy
 from typing import TYPE_CHECKING
 
 import dill
 
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.random.var import RandomStateSharedVariable
 
 from pymc3.distributions import _logcdf, _logp
 
@@ -77,14 +77,6 @@ def _random(*args, **kwargs):
         rv_type = None
 
         if isinstance(rv_op, RandomVariable):
-            if not rv_op.inplace:
-                # TODO: This is a temporary work-around.
-                # Remove this once we know what we want regarding RNG states
-                # and their propagation.
-                rv_op = copy(rv_op)
-                rv_op.inplace = True
-                clsdict["rv_op"] = rv_op
-
             rv_type = type(rv_op)
 
         new_cls = super().__new__(cls, name, bases, clsdict)
@@ -158,15 +150,31 @@ def __new__(cls, name, *args, **kwargs):
         return model.register_rv(rv_out, name, data, total_size, dims=dims, transform=transform)
 
     @classmethod
-    def dist(cls, dist_params, **kwargs):
+    def dist(cls, dist_params, rng=None, **kwargs):
 
         testval = kwargs.pop("testval", None)
 
-        rv_var = cls.rv_op(*dist_params, **kwargs)
+        rv_var = cls.rv_op(*dist_params, rng=rng, **kwargs)
 
         if testval is not None:
             rv_var.tag.test_value = testval
 
+        if (
+            rv_var.owner
+            and isinstance(rv_var.owner.op, RandomVariable)
+            and isinstance(rng, RandomStateSharedVariable)
+            and not getattr(rng, "default_update", None)
+        ):
+            # This tells `aesara.function` that the shared RNG variable
+            # is mutable, which--in turn--tells the `FunctionGraph`
+            # `Supervisor` feature to allow in-place updates on the variable.
+            # Without it, the `RandomVariable`s could not be optimized to allow
+            # in-place RNG updates, forcing all sample results from compiled
+            # functions to be the same on repeated evaluations.
+            new_rng = rv_var.owner.outputs[0]
+            rv_var.update = (rng, new_rng)
+            rng.default_update = new_rng
+
         return rv_var
 
     def _distr_parameters_for_repr(self):

pymc3/model.py

@@ -49,6 +49,7 @@
 
 from pymc3.aesaraf import (
     change_rv_size,
+    compile_rv_inplace,
     gradient,
     hessian,
     inputvars,
@@ -455,7 +456,7 @@ def __init__(
 
         inputs = grad_vars
 
-        self._aesara_function = aesara.function(inputs, outputs, givens=givens, **kwargs)
+        self._aesara_function = compile_rv_inplace(inputs, outputs, givens=givens, **kwargs)
 
     def set_weights(self, values):
         if values.shape != (self._n_costs - 1,):
@@ -1378,7 +1379,7 @@ def makefn(self, outs, mode=None, *args, **kwargs):
         Compiled Aesara function
         """
         with self:
-            return aesara.function(
+            return compile_rv_inplace(
                 self.value_vars,
                 outs,
                 allow_input_downcast=True,

pymc3/sampling.py

@@ -25,19 +25,19 @@
 from copy import copy, deepcopy
 from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast
 
-import aesara
 import aesara.gradient as tg
 import numpy as np
 import packaging
 import xarray
 
+from aesara.compile.mode import Mode
 from aesara.tensor.sharedvar import SharedVariable
 from arviz import InferenceData
 from fastprogress.fastprogress import progress_bar
 
 import pymc3 as pm
 
-from pymc3.aesaraf import change_rv_size, inputvars, walk_model
+from pymc3.aesaraf import change_rv_size, compile_rv_inplace, inputvars, walk_model
 from pymc3.backends.arviz import _DefaultTrace
 from pymc3.backends.base import BaseTrace, MultiTrace
 from pymc3.backends.ndarray import NDArray
@@ -1584,6 +1584,7 @@ def sample_posterior_predictive(
     keep_size: Optional[bool] = False,
     random_seed=None,
     progressbar: bool = True,
+    mode: Optional[Union[str, Mode]] = None,
 ) -> Dict[str, np.ndarray]:
     """Generate posterior predictive samples from a model given a trace.
 
@@ -1617,6 +1618,8 @@ def sample_posterior_predictive(
         Whether or not to display a progress bar in the command line. The bar shows the percentage
         of completion, the sampling speed in samples per second (SPS), and the estimated remaining
         time until completion ("expected time of arrival"; ETA).
+    mode:
+        The mode used by ``aesara.function`` to compile the graph.
 
     Returns
     -------
@@ -1727,12 +1730,13 @@ def sample_posterior_predictive(
     if size is not None:
         vars_to_sample = [change_rv_size(v, size, expand=True) for v in vars_to_sample]
 
-    sampler_fn = aesara.function(
+    sampler_fn = compile_rv_inplace(
         inputs,
         vars_to_sample,
         allow_input_downcast=True,
         accept_inplace=True,
         on_unused_input="ignore",
+        mode=mode,
     )
 
     ppc_trace_t = _DefaultTrace(samples)
@@ -1992,12 +1996,11 @@ def sample_prior_predictive(
 
     vars_to_sample = [model[name] for name in names]
     inputs = [i for i in inputvars(vars_to_sample) if not isinstance(i, SharedVariable)]
-    sampler_fn = aesara.function(
-        inputs,
-        vars_to_sample,
-        allow_input_downcast=True,
-        accept_inplace=True,
+
+    sampler_fn = compile_rv_inplace(
+        inputs, vars_to_sample, allow_input_downcast=True, accept_inplace=True, mode=mode
     )
+
     values = zip(*[sampler_fn() for i in range(samples)])
 
     data = {k: np.stack(v) for k, v in zip(names, values)}

pymc3/sampling_jax.py

@@ -7,7 +7,6 @@
 xla_flags = re.sub(r"xla_force_host_platform_device_count=.+\s", "", xla_flags).split()
 os.environ["XLA_FLAGS"] = " ".join(["--xla_force_host_platform_device_count={}".format(100)])
 
-import aesara.graph.fg
 import aesara.tensor as at
 import arviz as az
 import jax
@@ -23,6 +22,7 @@
 from aesara.tensor.type import TensorType
 
 from pymc3 import modelcontext
+from pymc3.aesaraf import compile_rv_inplace
 
 warnings.warn("This module is experimental.")
 
@@ -209,7 +209,7 @@ def sample_numpyro_nuts(
     print("Compiling...")
 
     tic1 = pd.Timestamp.now()
-    _sample = aesara.function(
+    _sample = compile_rv_inplace(
         [],
         sample_outputs + [numpyro_samples[-1]],
         allow_input_downcast=True,

pymc3/step_methods/metropolis.py

@@ -13,7 +13,6 @@
 #   limitations under the License.
 from typing import Any, Callable, Dict, List, Tuple
 
-import aesara
 import numpy as np
 import numpy.random as nr
 import scipy.linalg
@@ -23,7 +22,7 @@
 
 import pymc3 as pm
 
-from pymc3.aesaraf import floatX, rvs_to_value_vars
+from pymc3.aesaraf import compile_rv_inplace, floatX, rvs_to_value_vars
 from pymc3.blocking import DictToArrayBijection, RaveledVars
 from pymc3.step_methods.arraystep import (
     ArrayStep,
@@ -985,6 +984,6 @@ def delta_logp(point, logp, vars, shared):
 
     logp1 = pm.CallableTensor(logp0)(inarray1)
 
-    f = aesara.function([inarray1, inarray0], logp1 - logp0)
+    f = compile_rv_inplace([inarray1, inarray0], logp1 - logp0)
     f.trust_input = True
     return f

pymc3/tests/test_sampling.py

@@ -31,6 +31,7 @@
 
 import pymc3 as pm
 
+from pymc3.aesaraf import compile_rv_inplace
 from pymc3.backends.ndarray import NDArray
 from pymc3.exceptions import IncorrectArgumentsError, SamplingError
 from pymc3.tests.helpers import SeededTest
@@ -973,7 +974,7 @@ def test_layers(self):
             a = pm.Uniform("a", lower=0, upper=1, size=10)
             b = pm.Binomial("b", n=1, p=a, size=10)
 
-        b_sampler = aesara.function([], b)
+        b_sampler = compile_rv_inplace([], b, mode="FAST_RUN")
         avg = np.stack([b_sampler() for i in range(10000)]).mean(0)
         npt.assert_array_almost_equal(avg, 0.5 * np.ones((10,)), decimal=2)