diff --git a/pyro/poutine/subsample_messenger.py b/pyro/poutine/subsample_messenger.py index 5fe9203aea..bc8a88a562 100644 --- a/pyro/poutine/subsample_messenger.py +++ b/pyro/poutine/subsample_messenger.py @@ -115,6 +115,8 @@ def _subsample( size = -1 # This is PyTorch convention for "arbitrary size" subsample_size = -1 else: + if size < 0: + raise ValueError("size cannot be negative") msg = Message( type="sample", name=name, diff --git a/tests/test_primitives.py b/tests/test_primitives.py index c208a2cfa2..e7f4063218 100644 --- a/tests/test_primitives.py +++ b/tests/test_primitives.py @@ -36,6 +36,11 @@ def test_deterministic_ok(): assert x.shape == () +def test_plate_negative_size_error(): + with pytest.raises(ValueError, match="size cannot be negative"): + pyro.plate("data", -1) + + @pytest.mark.parametrize( "mask", [