Suppress subsample log_prob JIT warning

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
 
 

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__)