Refactor pm.Simulator (2nd attempt)

docs/source/api/distributions.rst

@@ -8,6 +8,7 @@ Distributions
    distributions/discrete
    distributions/multivariate
    distributions/mixture
+   distributions/simulator
    distributions/timeseries
    distributions/transforms
    distributions/utilities

docs/source/api/distributions/simulator.rst

@@ -0,0 +1,12 @@
+**********
+Simulator
+**********
+
+.. currentmodule:: pymc3.distributions.simulator
+.. autosummary::
+
+   SimulatorRV
+   Simulator
+
+.. automodule:: pymc3.distributions.simulator
+   :members:

pymc3/aesaraf.py

@@ -350,15 +350,26 @@ def rvs_to_value_vars(
 
     """
 
+    # Avoid circular dependency
+    from pymc3.distributions.simulator import SimulatorRV
+
     def transform_replacements(var, replacements):
         rv_var, rv_value_var = extract_rv_and_value_vars(var)
 
         if rv_value_var is None:
-            warnings.warn(
-                f"No value variable found for {rv_var}; "
-                "the random variable will not be replaced."
-            )
-            return []
+            # If RandomVariable does not have a value_var and corresponds to
+            # a SimulatorRV, we allow further replacements in upstream graph
+            if isinstance(rv_var.owner.op, SimulatorRV):
+                # First 3 inputs are just rng, dtype, and size, which don't
+                # need to be replaced.
+                return var.owner.inputs[3:]
+
+            else:
+                warnings.warn(
+                    f"No value variable found for {rv_var}; "
+                    "the random variable will not be replaced."
+                )
+                return []
 
         transform = getattr(rv_value_var.tag, "transform", None)
 

pymc3/distributions/__init__.py

@@ -100,7 +100,7 @@
     Wishart,
     WishartBartlett,
 )
-from pymc3.distributions.simulator import Simulator
+from pymc3.distributions.simulator import Simulator, SimulatorRV
 from pymc3.distributions.timeseries import (
     AR,
     AR1,
@@ -188,6 +188,7 @@
     "Rice",
     "Moyal",
     "Simulator",
+    "SimulatorRV",
     "BART",
     "CAR",
     "PolyaGamma",

pymc3/distributions/distribution.py

@@ -22,7 +22,6 @@
 from typing import Optional
 
 import aesara
-import aesara.tensor as at
 
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.var import RandomStateSharedVariable
@@ -325,47 +324,6 @@ def dist(
         return rv_out
 
 
-class NoDistribution(Distribution):
-    def __init__(
-        self,
-        shape,
-        dtype,
-        initval=None,
-        defaults=(),
-        parent_dist=None,
-        *args,
-        **kwargs,
-    ):
-        super().__init__(
-            shape=shape, dtype=dtype, initval=initval, defaults=defaults, *args, **kwargs
-        )
-        self.parent_dist = parent_dist
-
-    def __getattr__(self, name):
-        # Do not use __getstate__ and __setstate__ from parent_dist
-        # to avoid infinite recursion during unpickling
-        if name.startswith("__"):
-            raise AttributeError("'NoDistribution' has no attribute '%s'" % name)
-        return getattr(self.parent_dist, name)
-
-    def logp(self, x):
-        """Calculate log probability.
-
-        Parameters
-        ----------
-        x: numeric
-            Value for which log-probability is calculated.
-
-        Returns
-        -------
-        TensorVariable
-        """
-        return at.zeros_like(x)
-
-    def _distr_parameters_for_repr(self):
-        return []
-
-
 class Discrete(Distribution):
     """Base class for discrete distributions"""
 
@@ -381,6 +339,13 @@ class Continuous(Distribution):
     """Base class for continuous distributions"""
 
 
+class NoDistribution(Distribution):
+    """Base class for artifical distributions
+
+    RandomVariables that share this type are allowed in logprob graphs
+    """
+
+
 class DensityDist(Distribution):
     """Distribution based on a given log density function.