From 2d3ec8f3afcbaa9dcb8401af5fff6932ad7fd158 Mon Sep 17 00:00:00 2001
From: ricardoV94 <28983449+ricardoV94@users.noreply.github.com>
Date: Fri, 29 Jan 2021 17:23:27 +0100
Subject: [PATCH] Add tests for tt.switch related bugs (#4448)

* Add tests for edge cases

* Add release-note
---
 RELEASE-NOTES.md          |  1 +
 pymc3/tests/test_model.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index a2bff57c10b..939b8d54583 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -9,6 +9,7 @@
 
 ### Maintenance
 - We upgraded to `Theano-PyMC v1.1.2` which [includes bugfixes](https://github.com/pymc-devs/aesara/compare/rel-1.1.0...rel-1.1.2) for warning floods and compiledir locking (see [#4444](https://github.com/pymc-devs/pymc3/pull/4444))
+- `Theano-PyMC v1.1.2` also fixed an important issue in `tt.switch` that affected the behavior of several PyMC distributions, including at least the `Bernoulli` and `TruncatedNormal` (see[#4448](https://github.com/pymc-devs/pymc3/pull/4448))
 - `math.log1mexp_numpy` no longer raises RuntimeWarning when given very small inputs. These were commonly observed during NUTS sampling (see [#4428](https://github.com/pymc-devs/pymc3/pull/4428)).
 
 ## PyMC3 3.11.0 (21 January 2021)
diff --git a/pymc3/tests/test_model.py b/pymc3/tests/test_model.py
index c8be76c2d72..2e5a83c1c33 100644
--- a/pymc3/tests/test_model.py
+++ b/pymc3/tests/test_model.py
@@ -366,6 +366,34 @@ def test_tensor_type_conversion(self):
 
         assert m["x2_missing"].type == gf._extra_vars_shared["x2_missing"].type
 
+    def test_theano_switch_broadcast_edge_cases(self):
+        # Tests against two subtle issues related to a previous bug in Theano where tt.switch would not
+        # always broadcast tensors with single values https://github.com/pymc-devs/aesara/issues/270
+
+        # Known issue 1: https://github.com/pymc-devs/pymc3/issues/4389
+        data = np.zeros(10)
+        with pm.Model() as m:
+            p = pm.Beta("p", 1, 1)
+            obs = pm.Bernoulli("obs", p=p, observed=data)
+        # Assert logp is correct
+        npt.assert_allclose(
+            obs.logp(m.test_point),
+            np.log(0.5) * 10,
+        )
+
+        # Known issue 2: https://github.com/pymc-devs/pymc3/issues/4417
+        # fmt: off
+        data = np.array([
+            1.35202174, -0.83690274, 1.11175166, 1.29000367, 0.21282749,
+            0.84430966, 0.24841369, 0.81803141, 0.20550244, -0.45016253,
+        ])
+        # fmt: on
+        with pm.Model() as m:
+            mu = pm.Normal("mu", 0, 5)
+            obs = pm.TruncatedNormal("obs", mu=mu, sigma=1, lower=-1, upper=2, observed=data)
+        # Assert dlogp is correct
+        npt.assert_allclose(m.dlogp([mu])({"mu": 0}), 2.499424682024436, rtol=1e-5)
+
 
 def test_multiple_observed_rv():
     "Test previously buggy MultiObservedRV comparison code."