Merge remote-tracking branch 'upstream/v4'

Author: Larry Dong

.github/workflows/pytest.yml

@@ -62,6 +62,7 @@ jobs:
             --ignore=pymc3/tests/test_shape_handling.py
             --ignore=pymc3/tests/test_distributions.py
             --ignore=pymc3/tests/test_distributions_random.py
+            --ignore=pymc3/tests/test_idata_conversion.py
 
           - |
             pymc3/tests/test_modelcontext.py
@@ -73,6 +74,7 @@ jobs:
             pymc3/tests/test_updates.py
 
           - |
+            pymc3/tests/test_idata_conversion.py
             pymc3/tests/test_distributions.py
             pymc3/tests/test_distributions_random.py
             pymc3/tests/test_examples.py

pymc3/aesaraf.py

@@ -11,6 +11,8 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+import warnings
+
 from typing import (
     Callable,
     Dict,
@@ -169,19 +171,17 @@ def change_rv_size(
 def extract_rv_and_value_vars(
     var: TensorVariable,
 ) -> Tuple[TensorVariable, TensorVariable]:
-    """Extract a random variable and its corresponding value variable from a generic
-    `TensorVariable`.
+    """Return a random variable and it's observations or value variable, or ``None``.
 
     Parameters
     ==========
     var
-        A variable corresponding to a `RandomVariable`.
+        A variable corresponding to a ``RandomVariable``.
 
     Returns
     =======
-    The first value in the tuple is the `RandomVariable`, and the second is the
-    measure-space variable that corresponds with the latter (i.e. the "value"
-    variable).
+    The first value in the tuple is the ``RandomVariable``, and the second is the
+    measure/log-likelihood value variable that corresponds with the latter.
 
     """
     if not var.owner:
@@ -195,7 +195,7 @@ def extract_rv_and_value_vars(
 
 
 def extract_obs_data(x: TensorVariable) -> np.ndarray:
-    """Extract data observed symbolic variables.
+    """Extract data from observed symbolic variables.
 
     Raises
     ------
@@ -331,17 +331,24 @@ 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 []
 
         transform = getattr(rv_value_var.tag, "transform", None)
 
         if transform is None or not apply_transforms:
             replacements[var] = rv_value_var
-            return []
+            # In case the value variable is itself a graph, we walk it for
+            # potential replacements
+            return [rv_value_var]
 
         trans_rv_value = transform.backward(rv_var, rv_value_var)
         replacements[var] = trans_rv_value
 
+        # Walk the transformed variable and make replacements
         return [trans_rv_value]
 
     return replace_rvs_in_graphs(graphs, transform_replacements, initial_replacements, **kwargs)

pymc3/blocking.py

@@ -19,12 +19,17 @@
 """
 import collections
 
-from typing import Dict, List, Optional, Union
+from functools import partial
+from typing import Callable, Dict, Optional, TypeVar
 
 import numpy as np
 
 __all__ = ["DictToArrayBijection"]
 
+
+T = TypeVar("T")
+PointType = Dict[str, np.ndarray]
+
 # `point_map_info` is a tuple of tuples containing `(name, shape, dtype)` for
 # each of the raveled variables.
 RaveledVars = collections.namedtuple("RaveledVars", "data, point_map_info")
@@ -38,7 +43,7 @@ class DictToArrayBijection:
     """
 
     @staticmethod
-    def map(var_dict: Dict[str, np.ndarray]) -> RaveledVars:
+    def map(var_dict: PointType) -> RaveledVars:
         """Map a dictionary of names and variables to a concatenated 1D array space."""
         vars_info = tuple((v, k, v.shape, v.dtype) for k, v in var_dict.items())
         raveled_vars = [v[0].ravel() for v in vars_info]
@@ -50,42 +55,41 @@ def map(var_dict: Dict[str, np.ndarray]) -> RaveledVars:
 
     @staticmethod
     def rmap(
-        array: RaveledVars, as_list: Optional[bool] = False
-    ) -> Union[Dict[str, np.ndarray], List[np.ndarray]]:
+        array: RaveledVars,
+        start_point: Optional[PointType] = None,
+    ) -> PointType:
         """Map 1D concatenated array to a dictionary of variables in their original spaces.
 
         Parameters
         ==========
         array
             The array to map.
-        as_list
-            When ``True``, return a list of the original variables instead of a
-            ``dict`` keyed each variable's name.
+        start_point
+            An optional dictionary of initial values.
+
         """
-        if as_list:
-            res = []
+        if start_point:
+            res = dict(start_point)
         else:
             res = {}
 
         if not isinstance(array, RaveledVars):
-            raise TypeError("`apt` must be a `RaveledVars` type")
+            raise TypeError("`array` must be a `RaveledVars` type")
 
         last_idx = 0
         for name, shape, dtype in array.point_map_info:
             arr_len = np.prod(shape, dtype=int)
             var = array.data[last_idx : last_idx + arr_len].reshape(shape).astype(dtype)
-            if as_list:
-                res.append(var)
-            else:
-                res[name] = var
+            res[name] = var
             last_idx += arr_len
 
         return res
 
     @classmethod
-    def mapf(cls, f):
-        """
-         function f: DictSpace -> T to ArraySpace -> T
+    def mapf(cls, f: Callable[[PointType], T], start_point: Optional[PointType] = None) -> T:
+        """Create a callable that first maps back to ``dict`` inputs and then applies a function.
+
+        function f: DictSpace -> T to ArraySpace -> T
 
         Parameters
         ----------
@@ -95,7 +99,7 @@ def mapf(cls, f):
         -------
         f: array -> T
         """
-        return Compose(f, cls.rmap)
+        return Compose(f, partial(cls.rmap, start_point=start_point))
 
 
 class Compose:

pymc3/distributions/discrete.py

@@ -16,7 +16,14 @@
 import aesara.tensor as at
 import numpy as np
 
-from aesara.tensor.random.basic import bernoulli, binomial, categorical, nbinom, poisson
+from aesara.tensor.random.basic import (
+    RandomVariable,
+    bernoulli,
+    binomial,
+    categorical,
+    nbinom,
+    poisson,
+)
 from scipy import stats
 
 from pymc3.aesaraf import floatX, intX, take_along_axis
@@ -434,6 +441,22 @@ def _distr_parameters_for_repr(self):
         return ["p"]
 
 
+class DiscreteWeibullRV(RandomVariable):
+    name = "discrete_weibull"
+    ndim_supp = 0
+    ndims_params = [0, 0]
+    dtype = "int64"
+    _print_name = ("dWeibull", "\\operatorname{dWeibull}")
+
+    @classmethod
+    def rng_fn(cls, rng, q, beta, size):
+        p = rng.uniform(size=size)
+        return np.ceil(np.power(np.log(1 - p) / np.log(q), 1.0 / beta)) - 1
+
+
+discrete_weibull = DiscreteWeibullRV()
+
+
 class DiscreteWeibull(Discrete):
     R"""Discrete Weibull log-likelihood
 
@@ -473,51 +496,15 @@ def DiscreteWeibull(q, b, x):
     Variance  :math:`2 \sum_{x = 1}^{\infty} x q^{x^{\beta}} - \mu - \mu^2`
     ========  ======================
     """
+    rv_op = discrete_weibull
 
-    def __init__(self, q, beta, *args, **kwargs):
-        super().__init__(*args, defaults=("median",), **kwargs)
-
-        self.q = at.as_tensor_variable(floatX(q))
-        self.beta = at.as_tensor_variable(floatX(beta))
-
-        self.median = self._ppf(0.5)
-
-    def _ppf(self, p):
-        r"""
-        The percentile point function (the inverse of the cumulative
-        distribution function) of the discrete Weibull distribution.
-        """
-        q = self.q
-        beta = self.beta
-
-        return (at.ceil(at.power(at.log(1 - p) / at.log(q), 1.0 / beta)) - 1).astype("int64")
-
-    def _random(self, q, beta, size=None):
-        p = np.random.uniform(size=size)
-
-        return np.ceil(np.power(np.log(1 - p) / np.log(q), 1.0 / beta)) - 1
-
-    def random(self, point=None, size=None):
-        r"""
-        Draw random values from DiscreteWeibull distribution.
-
-        Parameters
-        ----------
-        point: dict, optional
-            Dict of variable values on which random values are to be
-            conditioned (uses default point if not specified).
-        size: int, optional
-            Desired size of random sample (returns one sample if not
-            specified).
-
-        Returns
-        -------
-        array
-        """
-        # q, beta = draw_values([self.q, self.beta], point=point, size=size)
-        # return generate_samples(self._random, q, beta, dist_shape=self.shape, size=size)
+    @classmethod
+    def dist(cls, q, beta, *args, **kwargs):
+        q = at.as_tensor_variable(floatX(q))
+        beta = at.as_tensor_variable(floatX(beta))
+        return super().dist([q, beta], **kwargs)
 
-    def logp(self, value):
+    def logp(value, q, beta):
         r"""
         Calculate log-probability of DiscreteWeibull distribution at specified value.
 
@@ -531,8 +518,6 @@ def logp(self, value):
         -------
         TensorVariable
         """
-        q = self.q
-        beta = self.beta
         return bound(
             at.log(at.power(q, at.power(value, beta)) - at.power(q, at.power(value + 1, beta))),
             0 <= value,
@@ -541,7 +526,7 @@ def logp(self, value):
             0 < beta,
         )
 
-    def logcdf(self, value):
+    def logcdf(value, q, beta):
         """
         Compute the log of the cumulative distribution function for Discrete Weibull distribution
         at the specified value.
@@ -556,9 +541,6 @@ def logcdf(self, value):
         -------
         TensorVariable
         """
-        q = self.q
-        beta = self.beta
-
         return bound(
             at.log1p(-at.power(q, at.power(value + 1, beta))),
             0 <= value,