diff --git a/pyro/poutine/subsample_messenger.py b/pyro/poutine/subsample_messenger.py index 5fe9203aea..83d2b05f0f 100644 --- a/pyro/poutine/subsample_messenger.py +++ b/pyro/poutine/subsample_messenger.py @@ -12,6 +12,11 @@ from pyro.util import ignore_jit_warnings +_TORCH_TENSOR_CONSTANT_WARNING = ( + "torch\\.[Tt]ensor results are registered as constants" +) + + class _Subsample(Distribution): """ Randomly select a subsample of a range of indices. @@ -44,7 +49,7 @@ def __init__( use_cuda, device ) ) - with ignore_jit_warnings(["torch.Tensor results are registered as constants"]): + with ignore_jit_warnings([_TORCH_TENSOR_CONSTANT_WARNING]): self.device = device or torch.tensor(tuple()).device @ignore_jit_warnings(["Converting a tensor to a Python boolean"]) @@ -67,7 +72,8 @@ def sample(self, sample_shape: torch.Size = torch.Size()) -> torch.Tensor: def log_prob(self, x: torch.Tensor) -> torch.Tensor: # This is zero so that plate can provide an unbiased estimate of # the non-subsampled log_prob. - result = torch.tensor(0.0, device=self.device) + with ignore_jit_warnings([_TORCH_TENSOR_CONSTANT_WARNING]): + result = torch.tensor(0.0, device=self.device) return result.cuda() if self.use_cuda else result diff --git a/tests/infer/test_jit.py b/tests/infer/test_jit.py index fa013339c6..1871cb5517 100644 --- a/tests/infer/test_jit.py +++ b/tests/infer/test_jit.py @@ -29,6 +29,7 @@ ) from pyro.optim import Adam from pyro.poutine.indep_messenger import CondIndepStackFrame +from pyro.poutine.subsample_messenger import _Subsample from pyro.util import ignore_jit_warnings from tests.common import assert_close, assert_equal @@ -39,6 +40,23 @@ def constant(*args, **kwargs): return torch.tensor(*args, **kwargs) +def test_subsample_log_prob_ignores_jit_constant_warning(): + subsample = _Subsample(3, None) + x = torch.arange(3) + + def f(x): + return subsample.log_prob(x) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", torch.jit.TracerWarning) + torch.jit.trace(f, (x,), check_trace=False) + + assert not any( + "torch.tensor results are registered as constants" in str(w.message).lower() + for w in caught + ) + + logger = logging.getLogger(__name__)