diff --git a/pyro/poutine/subsample_messenger.py b/pyro/poutine/subsample_messenger.py index 5fe9203aea..22d2c9d444 100644 --- a/pyro/poutine/subsample_messenger.py +++ b/pyro/poutine/subsample_messenger.py @@ -109,6 +109,10 @@ def _subsample( """ Helper function for plate. See its docstrings for details. """ + if size is not None and not torch._C._get_tracing_state() and size < 0: + raise ValueError( + "size must be a positive integer, got {}".format(size) + ) if size is None: assert subsample_size is None assert subsample is None diff --git a/tests/poutine/test_poutines.py b/tests/poutine/test_poutines.py index 311837bc87..e09c943d88 100644 --- a/tests/poutine/test_poutines.py +++ b/tests/poutine/test_poutines.py @@ -940,6 +940,18 @@ def model(): assert len(_DIM_ALLOCATOR._stack) == 0, "stack was not cleaned on error" +@pytest.mark.parametrize("size", [-1, -10]) +def test_plate_error_on_negative_size(size): + def model(): + with pyro.plate("foo", size): + pass + + assert len(_DIM_ALLOCATOR._stack) == 0 + with pytest.raises(ValueError, match="size must be a positive integer"): + poutine.trace(model)() + assert len(_DIM_ALLOCATOR._stack) == 0, "stack was not cleaned on error" + + @pytest.mark.parametrize( "graph_type, expected", [("flat", set()), ("dense", {"x", "y"})] )