|
12 | 12 | # See the License for the specific language governing permissions and
|
13 | 13 | # limitations under the License.
|
14 | 14 |
|
| 15 | +import aesara |
15 | 16 | import numpy as np
|
16 | 17 | import pytest
|
17 | 18 |
|
@@ -219,3 +220,32 @@ def test_sample_generate_values(fixture_model, fixture_sizes):
|
219 | 220 | prior = pm.sample_prior_predictive(samples=fixture_sizes)
|
220 | 221 | for rv in RVs:
|
221 | 222 | assert prior[rv.name].shape == size + tuple(rv.distribution.shape)
|
| 223 | + |
| 224 | + |
| 225 | +@pytest.mark.xfail(reason="https://github.com/pymc-devs/aesara/issues/390") |
| 226 | +def test_size32_doesnt_break_broadcasting(): |
| 227 | + size32 = at.constant([1, 10], dtype="int32") |
| 228 | + rv = pm.Normal.dist(0, 1, size=size32) |
| 229 | + assert rv.broadcastable == (True, False) |
| 230 | + |
| 231 | + |
| 232 | +def test_observed_with_column_vector(): |
| 233 | + """This test is related to https://github.com/pymc-devs/aesara/issues/390 which breaks |
| 234 | + broadcastability of column-vector RVs. This unexpected change in type can lead to |
| 235 | + incompatibilities during graph rewriting for model.logp evaluation. |
| 236 | + """ |
| 237 | + with pm.Model() as model: |
| 238 | + # The `observed` is a broadcastable column vector |
| 239 | + obs = at.as_tensor_variable(np.ones((3, 1), dtype=aesara.config.floatX)) |
| 240 | + assert obs.broadcastable == (False, True) |
| 241 | + |
| 242 | + # Both shapes describe broadcastable volumn vectors |
| 243 | + size64 = at.constant([3, 1], dtype="int64") |
| 244 | + # But the second shape is upcasted from an int32 vector |
| 245 | + cast64 = at.cast(at.constant([3, 1], dtype="int32"), dtype="int64") |
| 246 | + |
| 247 | + pm.Normal("x_size64", mu=0, sd=1, size=size64, observed=obs) |
| 248 | + model.logp() |
| 249 | + |
| 250 | + pm.Normal("x_cast64", mu=0, sd=1, size=cast64, observed=obs) |
| 251 | + model.logp() |
0 commit comments