Merge remote-tracking branch 'upstream/master' into less-strict-test

Author: MarcoGorelli

RELEASE-NOTES.md

@@ -26,6 +26,7 @@ pip install theano-pymc
 This new version of `Theano-PyMC` comes with an experimental JAX backend which, when combined with the new and experimental JAX samplers in PyMC3, can greatly speed up sampling in your model. As this is still very new, please do not use it in production yet but do test it out and let us know if anything breaks and what results you are seeing, especially speed-wise.
 
 ### New features
+- New experimental JAX samplers in `pymc3.sample_jax` (see [notebook](https://docs.pymc.io/notebooks/GLM-hierarchical-jax.html) and [#4247](https://github.com/pymc-devs/pymc3/pull/4247)). Requires JAX and either TFP or numpyro.
 - Add MLDA, a new stepper for multilevel sampling. MLDA can be used when a hierarchy of approximate posteriors of varying accuracy is available, offering improved sampling efficiency especially in high-dimensional problems and/or where gradients are not available (see [#3926](https://github.com/pymc-devs/pymc3/pull/3926))
 - Add Bayesian Additive Regression Trees (BARTs) [#4183](https://github.com/pymc-devs/pymc3/pull/4183))
 - Added `pymc3.gp.cov.Circular` kernel for Gaussian Processes on circular domains, e.g. the unit circle (see [#4082](https://github.com/pymc-devs/pymc3/pull/4082)).
@@ -42,6 +43,8 @@ This new version of `Theano-PyMC` comes with an experimental JAX backend which,
 - Fixed numerical instability in ExGaussian's logp by preventing `logpow` from returning `-inf` (see [#4050](https://github.com/pymc-devs/pymc3/pull/4050)).
 - Numerically improved stickbreaking transformation - e.g. for the `Dirichlet` distribution. [#4129](https://github.com/pymc-devs/pymc3/pull/4129)
 - Enabled the `Multinomial` distribution to handle batch sizes that have more than 2 dimensions. [#4169](https://github.com/pymc-devs/pymc3/pull/4169)
+- Test model logp before starting any MCMC chains (see [#4116](https://github.com/pymc-devs/pymc3/issues/4116))
+- Fix bug in `model.check_test_point` that caused the `test_point` argument to be ignored. (see [PR #4211](https://github.com/pymc-devs/pymc3/pull/4211#issuecomment-727142721))
 
 ### Documentation
 - Added a new notebook demonstrating how to incorporate sampling from a conjugate Dirichlet-multinomial posterior density in conjunction with other step methods (see [#4199](https://github.com/pymc-devs/pymc3/pull/4199)).

docs/source/notebooks/GLM-hierarchical-jax.ipynb

@@ -64,5 +64,6 @@ Gallery.contents = {
     "MLDA_introduction": "MCMC",
     "MLDA_simple_linear_regression": "MCMC",
     "MLDA_gravity_surveying": "MCMC",
-    "MLDA_variance_reduction_linear_regression": "MCMC"
+    "MLDA_variance_reduction_linear_regression": "MCMC",
+    "GLM-hierarchical-jax": "MCMC"
 }

docs/source/notebooks/table_of_contents_examples.js

@@ -246,7 +246,7 @@ def _str_repr(self, name=None, dist=None, formatting="plain"):
         alpha = self.alpha
         m = self.m
 
-        if formatting == "latex":
+        if "latex" in formatting:
             return f"$\\text{{{name}}} \\sim  \\text{{BART}}(\\text{{alpha = }}\\text{{{alpha}}}, \\text{{m = }}\\text{{{m}}})$"
         else:
             return f"{name} ~ BART(alpha = {alpha}, m = {m})"

pymc3/distributions/bart.py

@@ -157,13 +157,13 @@ def _distr_name_for_repr(self):
 
     def _str_repr(self, **kwargs):
         distr_repr = self._wrapped._str_repr(**{**kwargs, "dist": self._wrapped})
-        if "formatting" in kwargs and kwargs["formatting"] == "latex":
+        if "formatting" in kwargs and "latex" in kwargs["formatting"]:
             distr_repr = distr_repr[distr_repr.index(r" \sim") + 6 :]
         else:
             distr_repr = distr_repr[distr_repr.index(" ~") + 3 :]
         self_repr = super()._str_repr(**kwargs)
 
-        if "formatting" in kwargs and kwargs["formatting"] == "latex":
+        if "formatting" in kwargs and "latex" in kwargs["formatting"]:
             return self_repr + " -- " + distr_repr
         else:
             return self_repr + "-" + distr_repr

pymc3/distributions/bound.py

@@ -164,44 +164,61 @@ def _distr_name_for_repr(self):
         return self.__class__.__name__
 
     def _str_repr(self, name=None, dist=None, formatting="plain"):
-        """Generate string representation for this distribution, optionally
+        """
+        Generate string representation for this distribution, optionally
         including LaTeX markup (formatting='latex').
+
+        Parameters
+        ----------
+        name : str
+            name of the distribution
+        dist : Distribution
+            the distribution object
+        formatting : str
+            one of { "latex", "plain", "latex_with_params", "plain_with_params" }
         """
         if dist is None:
             dist = self
         if name is None:
             name = "[unnamed]"
+        supported_formattings = {"latex", "plain", "latex_with_params", "plain_with_params"}
+        if not formatting in supported_formattings:
+            raise ValueError(f"Unsupported formatting ''. Choose one of {supported_formattings}.")
 
         param_names = self._distr_parameters_for_repr()
         param_values = [
             get_repr_for_variable(getattr(dist, x), formatting=formatting) for x in param_names
         ]
 
-        if formatting == "latex":
+        if "latex" in formatting:
             param_string = ",~".join(
                 [fr"\mathit{{{name}}}={value}" for name, value in zip(param_names, param_values)]
             )
-            return r"$\text{{{var_name}}} \sim \text{{{distr_name}}}({params})$".format(
-                var_name=name, distr_name=dist._distr_name_for_repr(), params=param_string
+            if formatting == "latex_with_params":
+                return r"$\text{{{var_name}}} \sim \text{{{distr_name}}}({params})$".format(
+                    var_name=name, distr_name=dist._distr_name_for_repr(), params=param_string
+                )
+            return r"$\text{{{var_name}}} \sim \text{{{distr_name}}}$".format(
+                var_name=name, distr_name=dist._distr_name_for_repr()
             )
         else:
-            # 'plain' is default option
+            # one of the plain formattings
             param_string = ", ".join(
                 [f"{name}={value}" for name, value in zip(param_names, param_values)]
             )
-            return "{var_name} ~ {distr_name}({params})".format(
-                var_name=name, distr_name=dist._distr_name_for_repr(), params=param_string
-            )
+            if formatting == "plain_with_params":
+                return f"{name} ~ {dist._distr_name_for_repr()}({param_string})"
+            return f"{name} ~ {dist._distr_name_for_repr()}"
 
     def __str__(self, **kwargs):
         try:
             return self._str_repr(formatting="plain", **kwargs)
         except:
             return super().__str__()
 
-    def _repr_latex_(self, **kwargs):
+    def _repr_latex_(self, *, formatting="latex_with_params", **kwargs):
         """Magic method name for IPython to use for LaTeX formatting."""
-        return self._str_repr(formatting="latex", **kwargs)
+        return self._str_repr(formatting=formatting, **kwargs)
 
     def logp_nojac(self, *args, **kwargs):
         """Return the logp, but do not include a jacobian term for transforms.

pymc3/distributions/distribution.py

@@ -126,7 +126,7 @@ def _str_repr(self, name=None, dist=None, formatting="plain"):
         sum_stat = self.sum_stat.__name__ if hasattr(self.sum_stat, "__call__") else self.sum_stat
         distance = getattr(self.distance, "__name__", self.distance.__class__.__name__)
 
-        if formatting == "latex":
+        if "latex" in formatting:
             return f"$\\text{{{name}}} \\sim  \\text{{Simulator}}(\\text{{{function}}}({params}), \\text{{{distance}}}, \\text{{{sum_stat}}})$"
         else:
             return f"{name} ~ Simulator({function}({params}), {distance}, {sum_stat})"

pymc3/distributions/simulator.py

@@ -65,7 +65,7 @@ def __rmatmul__(self, other):
 
     def _str_repr(self, name=None, dist=None, formatting="plain"):
         if getattr(self, "distribution", None) is None:
-            if formatting == "latex":
+            if "latex" in formatting:
                 return None
             else:
                 return super().__str__()
@@ -76,8 +76,8 @@ def _str_repr(self, name=None, dist=None, formatting="plain"):
             dist = self.distribution
         return self.distribution._str_repr(name=name, dist=dist, formatting=formatting)
 
-    def _repr_latex_(self, **kwargs):
-        return self._str_repr(formatting="latex", **kwargs)
+    def _repr_latex_(self, *, formatting="latex_with_params", **kwargs):
+        return self._str_repr(formatting=formatting, **kwargs)
 
     def __str__(self, **kwargs):
         try:
@@ -1368,15 +1368,15 @@ def check_test_point(self, test_point=None, round_vals=2):
             test_point = self.test_point
 
         return Series(
-            {RV.name: np.round(RV.logp(self.test_point), round_vals) for RV in self.basic_RVs},
+            {RV.name: np.round(RV.logp(test_point), round_vals) for RV in self.basic_RVs},
             name="Log-probability of test_point",
         )
 
     def _str_repr(self, formatting="plain", **kwargs):
         all_rv = itertools.chain(self.unobserved_RVs, self.observed_RVs)
 
-        if formatting == "latex":
-            rv_reprs = [rv.__latex__() for rv in all_rv]
+        if "latex" in formatting:
+            rv_reprs = [rv.__latex__(formatting=formatting) for rv in all_rv]
             rv_reprs = [
                 rv_repr.replace(r"\sim", r"&\sim &").strip("$")
                 for rv_repr in rv_reprs
@@ -1407,8 +1407,8 @@ def _str_repr(self, formatting="plain", **kwargs):
     def __str__(self, **kwargs):
         return self._str_repr(formatting="plain", **kwargs)
 
-    def _repr_latex_(self, **kwargs):
-        return self._str_repr(formatting="latex", **kwargs)
+    def _repr_latex_(self, *, formatting="latex", **kwargs):
+        return self._str_repr(formatting=formatting, **kwargs)
 
     __latex__ = _repr_latex_
 
@@ -1874,24 +1874,27 @@ def _walk_up_rv(rv, formatting="plain"):
             all_rvs.extend(_walk_up_rv(parent, formatting=formatting))
     else:
         name = rv.name if rv.name else "Constant"
-        fmt = r"\text{{{name}}}" if formatting == "latex" else "{name}"
+        fmt = r"\text{{{name}}}" if "latex" in formatting else "{name}"
         all_rvs.append(fmt.format(name=name))
     return all_rvs
 
 
 class DeterministicWrapper(tt.TensorVariable):
     def _str_repr(self, formatting="plain"):
-        if formatting == "latex":
-            return r"$\text{{{name}}} \sim \text{{Deterministic}}({args})$".format(
-                name=self.name, args=r",~".join(_walk_up_rv(self, formatting=formatting))
-            )
+        if "latex" in formatting:
+            if formatting == "latex_with_params":
+                return r"$\text{{{name}}} \sim \text{{Deterministic}}({args})$".format(
+                    name=self.name, args=r",~".join(_walk_up_rv(self, formatting=formatting))
+                )
+            return fr"$\text{{{self.name}}} \sim \text{{Deterministic}}$"
         else:
-            return "{name} ~ Deterministic({args})".format(
-                name=self.name, args=", ".join(_walk_up_rv(self, formatting=formatting))
-            )
+            if formatting == "plain_with_params":
+                args = ", ".join(_walk_up_rv(self, formatting=formatting))
+                return f"{self.name} ~ Deterministic({args})"
+            return f"{self.name} ~ Deterministic"
 
-    def _repr_latex_(self):
-        return self._str_repr(formatting="latex")
+    def _repr_latex_(self, *, formatting="latex_with_params", **kwargs):
+        return self._str_repr(formatting=formatting)
 
     __latex__ = _repr_latex_
 

pymc3/model.py

@@ -121,7 +121,7 @@ def update_input_map(key: str, val: Set[VarName]):
                     pass
         return input_map
 
-    def _make_node(self, var_name, graph):
+    def _make_node(self, var_name, graph, *, formatting: str = "plain"):
         """Attaches the given variable to a graphviz Digraph"""
         v = self.model[var_name]
 
@@ -146,7 +146,7 @@ def _make_node(self, var_name, graph):
         elif isinstance(v, SharedVariable):
             label = f"{var_name}\n~\nData"
         else:
-            label = str(v).replace(" ~ ", "\n~\n")
+            label = v._str_repr(formatting=formatting).replace(" ~ ", "\n~\n")
 
         graph.node(var_name.replace(":", "&"), label, **attrs)
 
@@ -181,7 +181,7 @@ def get_plates(self):
             plates[shape].add(var_name)
         return plates
 
-    def make_graph(self):
+    def make_graph(self, formatting: str = "plain"):
         """Make graphviz Digraph of PyMC3 model
 
         Returns
@@ -205,20 +205,20 @@ def make_graph(self):
                 # must be preceded by 'cluster' to get a box around it
                 with graph.subgraph(name="cluster" + label) as sub:
                     for var_name in var_names:
-                        self._make_node(var_name, sub)
+                        self._make_node(var_name, sub, formatting=formatting)
                     # plate label goes bottom right
                     sub.attr(label=label, labeljust="r", labelloc="b", style="rounded")
             else:
                 for var_name in var_names:
-                    self._make_node(var_name, graph)
+                    self._make_node(var_name, graph, formatting=formatting)
 
         for key, values in self.make_compute_graph().items():
             for value in values:
                 graph.edge(value.replace(":", "&"), key.replace(":", "&"))
         return graph
 
 
-def model_to_graphviz(model=None):
+def model_to_graphviz(model=None, *, formatting: str = "plain"):
     """Produce a graphviz Digraph from a PyMC3 model.
 
     Requires graphviz, which may be installed most easily with
@@ -228,6 +228,15 @@ def model_to_graphviz(model=None):
     and then `pip install graphviz` to get the python bindings.  See
     http://graphviz.readthedocs.io/en/stable/manual.html
     for more information.
+
+    Parameters
+    ----------
+    model : pm.Model
+        The model to plot. Not required when called from inside a modelcontext.
+    formatting : str
+        one of { "plain", "plain_with_params" }
     """
+    if not "plain" in formatting:
+        raise ValueError(f"Unsupported formatting for graph nodes: '{formatting}'. See docstring.")
     model = pm.modelcontext(model)
-    return ModelGraph(model).make_graph()
+    return ModelGraph(model).make_graph(formatting=formatting)

pymc3/model_graph.py

@@ -54,6 +54,7 @@
     PGBART,
 )
 from .util import (
+    check_start_vals,
     update_start_vals,
     get_untransformed_name,
     is_transformed_name,
@@ -419,7 +420,16 @@ def sample(
 
     """
     model = modelcontext(model)
+    if start is None:
+        start = model.test_point
+    else:
+        if isinstance(start, dict):
+            update_start_vals(start, model.test_point, model)
+        else:
+            for chain_start_vals in start:
+                update_start_vals(chain_start_vals, model.test_point, model)
 
+    check_start_vals(start, model)
     if cores is None:
         cores = min(4, _cpu_count())
 
@@ -487,6 +497,7 @@ def sample(
                 progressbar=progressbar,
                 **kwargs,
             )
+            check_start_vals(start_, model)
             if start is None:
                 start = start_
         except (AttributeError, NotImplementedError, tg.NullTypeGradError):