Repository: pymc-devs/pymc
Pull request: Fix dependence of Uniform logp on bound method
This PR fixes #4499  

The Uniform and Discrete Uniform logp methods depended on the dist_math::bound method for correct evaluation of multiple values. This is problematic because the bound can be disabled by setting `check_bounds = False` in `pm.Model()`.

Other changes:
1. I tried to improve the documentation of the bound method to emphasize this. I also tried to make it a bit more clear what the "broadcast_conditions" secret keyword argument means and when it should be used. I would suggest this to be made an optional keyword argument in the function call of `bound`, and perhaps even rename it to `multivariate=False` to make it more obvious when should it be used. Let me know if you agree and I can add it to this PR as well.

2. I tweaked the docstring of `check_bounds` to emphasize this should not be used for models containing discrete variables as some of these depend on bound checks when sampled (unless they have specific samplers that cannot propose invalid values or the logp happens to behave nicely when given out-of-bound points)

3. Removed rendudant `tt.all()` calls in the bound conditions of several multivariate distributions.

*** 

Do we need a release note?

Depending on what your PR does, here are a few things you might want to address in the description:
+ [x] ~~what are the (breaking) changes that this PR makes?~~ None
+ [x] ~~important background, or details about the implementation~~ None
+ [x] are the changes—especially new features—covered by tests and docstrings? Tests cover known bugs. Docstrings improved
+ [x] [linting/style checks have been run](https://github.com/pymc-devs/pymc3/wiki/PyMC3-Python-Code-Style)
+ [x] ~~[consider adding/updating relevant example notebooks](https://github.com/pymc-devs/pymc-examples)~~
+ [x] right before it's ready to merge, mention the PR in the RELEASE-NOTES.md


Problem statement:
ValueError when using check_bounds=False in a model with DiscreteUniform
I was poking around to see if we should change the docstrings of check_bounds in pm.Model to mention this might not be safe in the presence of discrete variables and I found this strange error:

```python
with pm.Model(check_bounds=False) as m:
    y = pm.DiscreteUniform('y', 0, 5)
    trace = pm.sample()
```
```
Traceback (most recent call last):
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3418, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-60-ca8c3478b1dc>", line 3, in <module>
    trace = pm.sample()
  File "/home/ricardo/Documents/Projects/pymc3/pymc3/sampling.py", line 515, in sample
    step = assign_step_methods(model, step, step_kwargs=kwargs)
  File "/home/ricardo/Documents/Projects/pymc3/pymc3/sampling.py", line 217, in assign_step_methods
    return instantiate_steppers(model, steps, selected_steps, step_kwargs)
  File "/home/ricardo/Documents/Projects/pymc3/pymc3/sampling.py", line 138, in instantiate_steppers
    step = step_class(vars=vars, **args)
  File "/home/ricardo/Documents/Projects/pymc3/pymc3/step_methods/metropolis.py", line 187, in __init__
    self.delta_logp = delta_logp(model.logpt, vars, shared)
  File "/home/ricardo/Documents/Projects/pymc3/pymc3/step_methods/metropolis.py", line 888, in delta_logp
    logp1 = pm.CallableTensor(logp0)(inarray1)
  File "/home/ricardo/Documents/Projects/pymc3/pymc3/aesaraf.py", line 308, in __call__
    (oldinput,) = inputvars(self.tensor)
ValueError: not enough values to unpack (expected 1, got 0)
```

I don't see why changing the check_bounds flag would lead to any issues in sampling. I can call the logp method for any value and I get a finite (and wrong) evaluation as would be expected from removing the bounds:

```python
np.array([m.logp({'y': y}) for y in np.arange(-10, 10)])
array([-1.79175947, -1.79175947, -1.79175947, -1.79175947, -1.79175947,
       -1.79175947, -1.79175947, -1.79175947, -1.79175947, -1.79175947,
       -1.79175947, -1.79175947, -1.79175947, -1.79175947, -1.79175947,
       -1.79175947, -1.79175947, -1.79175947, -1.79175947, -1.79175947])
```

I tried with a couple of other Discrete Variables and I had no issues either. Obviously it works as expected with the flag set to True

```python
with pm.Model(check_bounds=True) as m:
    y = pm.DiscreteUniform('y', 0, 5)
    trace = pm.sample()
```

Hints from issue discussion:

Commit message: Add release note

Patch to review:
diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index b8f7537b89..ddf9db4f67 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -13,6 +13,8 @@
 - The `pymc3.memoize` module was removed and replaced with `cachetools`.  The `hashable` function and `WithMemoization` class were moved to `pymc3.util` (see [#4509](https://github.com/pymc-devs/pymc3/pull/4509)).
 - Remove float128 dtype support (see [#4514](https://github.com/pymc-devs/pymc3/pull/4514)).
 - `pm.make_shared_replacements` now retains broadcasting information which fixes issues with Metropolis samplers (see [#4492](https://github.com/pymc-devs/pymc3/pull/4492)).
+- Logp method of `Uniform` and `DiscreteUniform` no longer depends on `pymc3.distributions.dist_math.bound` for proper evaluation (see [#4541](https://github.com/pymc-devs/pymc3/pull/4541)).
+
 + ...
 
 ## PyMC3 3.11.1 (12 February 2021)
diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py
index 25310b4781..d073ceef49 100644
--- a/pymc3/distributions/continuous.py
+++ b/pymc3/distributions/continuous.py
@@ -268,7 +268,11 @@ def logp(self, value):
         """
         lower = self.lower
         upper = self.upper
-        return bound(-aet.log(upper - lower), value >= lower, value <= upper)
+        return bound(
+            aet.fill(value, -aet.log(upper - lower)),
+            value >= lower,
+            value <= upper,
+        )
 
     def logcdf(self, value):
         """
diff --git a/pymc3/distributions/discrete.py b/pymc3/distributions/discrete.py
index 06cd504f40..ab1aa2a198 100644
--- a/pymc3/distributions/discrete.py
+++ b/pymc3/distributions/discrete.py
@@ -1278,7 +1278,11 @@ def logp(self, value):
         """
         upper = self.upper
         lower = self.lower
-        return bound(-aet.log(upper - lower + 1), lower <= value, value <= upper)
+        return bound(
+            aet.fill(value, -aet.log(upper - lower + 1)),
+            lower <= value,
+            value <= upper,
+        )
 
     def logcdf(self, value):
         """
diff --git a/pymc3/distributions/dist_math.py b/pymc3/distributions/dist_math.py
index 90ff0c1908..771d1802b2 100644
--- a/pymc3/distributions/dist_math.py
+++ b/pymc3/distributions/dist_math.py
@@ -48,17 +48,23 @@
 def bound(logp, *conditions, **kwargs):
     """
     Bounds a log probability density with several conditions.
+    When conditions are not met, the logp values are replaced by -inf.
+
+    Note that bound should not be used to enforce the logic of the logp under the normal
+    support as it can be disabled by the user via check_bounds = False in pm.Model()
 
     Parameters
     ----------
     logp: float
     *conditions: booleans
     broadcast_conditions: bool (optional, default=True)
-        If True, broadcasts logp to match the largest shape of the conditions.
-        This is used e.g. in DiscreteUniform where logp is a scalar constant and the shape
-        is specified via the conditions.
-        If False, will return the same shape as logp.
-        This is used e.g. in Multinomial where broadcasting can lead to differences in the logp.
+        If True, conditions are broadcasted and applied element-wise to each value in logp.
+        If False, conditions are collapsed via aet.all(). As a consequence the entire logp
+        array is either replaced by -inf or unchanged.
+
+        Setting broadcasts_conditions to False is necessary for most (all?) multivariate
+        distributions where the dimensions of the conditions do not unambigously match
+        that of the logp.
 
     Returns
     -------
diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 18751e32c7..94eef7bdc1 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -527,9 +527,9 @@ def logp(self, value):
         # only defined for sum(value) == 1
         return bound(
             aet.sum(logpow(value, a - 1) - gammaln(a), axis=-1) + gammaln(aet.sum(a, axis=-1)),
-            aet.all(value >= 0),
-            aet.all(value <= 1),
-            aet.all(a > 0),
+            value >= 0,
+            value <= 1,
+            a > 0,
             broadcast_conditions=False,
         )
 
@@ -671,11 +671,11 @@ def logp(self, x):
 
         return bound(
             factln(n) + aet.sum(-factln(x) + logpow(p, x), axis=-1, keepdims=True),
-            aet.all(x >= 0),
-            aet.all(aet.eq(aet.sum(x, axis=-1, keepdims=True), n)),
-            aet.all(p <= 1),
-            aet.all(aet.eq(aet.sum(p, axis=-1), 1)),
-            aet.all(aet.ge(n, 0)),
+            x >= 0,
+            aet.eq(aet.sum(x, axis=-1, keepdims=True), n),
+            p <= 1,
+            aet.eq(aet.sum(p, axis=-1), 1),
+            n >= 0,
             broadcast_conditions=False,
         )
 
@@ -823,10 +823,10 @@ def logp(self, value):
         # and that each observation value_i sums to n_i.
         return bound(
             result,
-            aet.all(aet.ge(value, 0)),
-            aet.all(aet.gt(a, 0)),
-            aet.all(aet.ge(n, 0)),
-            aet.all(aet.eq(value.sum(axis=-1, keepdims=True), n)),
+            value >= 0,
+            a > 0,
+            n >= 0,
+            aet.eq(value.sum(axis=-1, keepdims=True), n),
             broadcast_conditions=False,
         )
 
@@ -1575,8 +1575,8 @@ def logp(self, x):
         result += (eta - 1.0) * aet.log(det(X))
         return bound(
             result,
-            aet.all(X <= 1),
-            aet.all(X >= -1),
+            X >= -1,
+            X <= 1,
             matrix_pos_def(X),
             eta > 0,
             broadcast_conditions=False,
@@ -2204,9 +2204,10 @@ def logp(self, value):
         logquad = (self.tau * delta * tau_dot_delta).sum(axis=-1)
         return bound(
             0.5 * (logtau + logdet - logquad),
-            aet.all(self.alpha <= 1),
-            aet.all(self.alpha >= -1),
+            self.alpha >= -1,
+            self.alpha <= 1,
             self.tau > 0,
+            broadcast_conditions=False,
         )
 
     def random(self, point=None, size=None):
diff --git a/pymc3/model.py b/pymc3/model.py
index 76800e7960..656eb9f078 100644
--- a/pymc3/model.py
+++ b/pymc3/model.py
@@ -817,7 +817,8 @@ class Model(Factor, WithMemoization, metaclass=ContextMeta):
         Ensure that input parameters to distributions are in a valid
         range. If your model is built in a way where you know your
         parameters can only take on valid values you can set this to
-        False for increased speed.
+        False for increased speed. This should not be used if your model
+        contains discrete variables.
 
     Examples
     --------
diff --git a/pymc3/tests/test_distributions.py b/pymc3/tests/test_distributions.py
index 7662e0ed58..d6a98ff233 100644
--- a/pymc3/tests/test_distributions.py
+++ b/pymc3/tests/test_distributions.py
@@ -2667,6 +2667,17 @@ def test_issue_3051(self, dims, dist_cls, kwargs):
         assert actual_a.shape == (X.shape[0],)
         pass
 
+    def test_issue_4499(self):
+        # Test for bug in Uniform and DiscreteUniform logp when setting check_bounds = False
+        # https://github.com/pymc-devs/pymc3/issues/4499
+        with pm.Model(check_bounds=False) as m:
+            x = pm.Uniform("x", 0, 2, shape=10, transform=None)
+        assert_almost_equal(m.logp_array(np.ones(10)), -np.log(2) * 10)
+
+        with pm.Model(check_bounds=False) as m:
+            x = pm.DiscreteUniform("x", 0, 1, shape=10)
+        assert_almost_equal(m.logp_array(np.ones(10)), -np.log(2) * 10)
+
 
 def test_serialize_density_dist():
     def func(x):


Code review comments:
diff_hunk: @@ -13,6 +13,8 @@
 - The `pymc3.memoize` module was removed and replaced with `cachetools`.  The `hashable` function and `WithMemoization` class were moved to `pymc3.util` (see [#4509](https://github.com/pymc-devs/pymc3/pull/4509)).
 - Remove float128 dtype support (see [#4514](https://github.com/pymc-devs/pymc3/pull/4514)).
 - `pm.make_shared_replacements` now retains broadcasting information which fixes issues with Metropolis samplers (see [#4492](https://github.com/pymc-devs/pymc3/pull/4492)).
+- Logp method of `Uniform` and `DiscreteUniform` no longer depends on `pymc3.distributions.dist_math.bound` for proper evaluation (see [#4541](https://github.com/pymc-devs/pymc3/pull/4541)).
path: RELEASE-NOTES.md
text: @user1:
Something is wrong here as we already release `3.11.2` this should go into `vNext`.

@author:
Right, it seems the release notes were only changed in V3. Should I rebase the release notes from here? https://github.com/pymc-devs/pymc3/blob/v3/RELEASE-NOTES.md

@user2:
Good question. The Release notes are the only document that spans branches and will continue to even after `v4` was released.
I don't think we'll be able to keep up with rebasing it going into the future.

Can you bundle these items in a `vNext (4.0.0)` section, but still point out those that were already released in `3.11.2`? You can link to the Release itself, where the release notes are duplicated.

@author:
Sure, let me give it a try