Skip to content

[bug] TracerWarning from subsample_messenger when running NUTS with jit_compile=True #3449

Description

@richardmkit

Issue Description

When running pyro.infer.MCMC with NUTS(..., jit_compile=True), Pyro emits the following warning from pyro/ poutine/subsample_messenger.py:

/opt/miniconda3/envs/pyro/lib/python3.10/site-packages/pyro/poutine/subsample_messenger.py:70:
TracerWarning: torch.tensor results are registered as constants in the trace.
You can safely ignore this warning if you use this function to create tensors out of constant variables that would be
the same every time you call this function.
In any other case, this might cause the trace to be incorrect.
  result = torch.tensor(0.0, device=self.device)

This appears even when the model runs successfully and the warning seems to come from Pyro internals rather than user
model code.

I understand this may be harmless because the tensor is a constant, but it is confusing for users because:

1. the warning looks like a potential tracing bug,
2. it points to Pyro internal code instead of user code,
3. the only obvious way to suppress it is to disable jit_compile=True, which may reduce performance.

I am opening this issue to ask whether this warning is expected under JIT-compiled NUTS, and whether Pyro should
suppress or avoid triggering it internally.

### Environment

- OS: macOS 26.3.1 (arm64)
- Python: 3.10.19
- PyTorch: 2.5.1
- Pyro: 1.9.1+ab0491a

### Code Snippet

A minimal pattern that triggers the warning is:

import pyro
from pyro.infer import MCMC, NUTS

nuts = NUTS(model, jit_compile=True)
mcmc = MCMC(nuts, warmup_steps=1000, num_samples=1000, num_chains=4)
mcmc.run(obs=data)

In my case the warning is raised from:

result = torch.tensor(0.0, device=self.device)

inside pyro/poutine/subsample_messenger.py.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions