Posterior `.to()` methods fail on plain torch priors and drop cached state
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 865
- Forks
- 260
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 5
Description
🐛 Bug Description
The .to() methods of the posterior classes have two gaps. #1972 worked around them in NeuralPosterior.__setstate__ by updating the device strings directly instead of calling self.to(). The gaps remain for users who call .to() themselves.
DirectPosterior.to()andVectorFieldPosterior.to()raiseValueError("Prior has no attribute to(device)")for plain torch priors such asMultivariateNormal.BasePotential.to()handles the same case withmove_distribution_to_device().- Every posterior
.to()re-runssuper().__init__(). That drops the cached_mapand resets_purpose. A user who callsmap()and then.to()computes the MAP again.
🔄 Steps to Reproduce
import torch
from torch.distributions import MultivariateNormal
from sbi.inference import NPE
prior = MultivariateNormal(torch.zeros(2), torch.eye(2))
theta = prior.sample((200,))
x = theta + 0.1 * torch.randn_like(theta)
trainer = NPE(prior=prior, show_progress_bars=False)
trainer.append_simulations(theta, x).train(max_num_epochs=1)
posterior = trainer.build_posterior()
posterior.to("cpu")
# ValueError: Prior has no attribute to(device).
✅ Expected Behavior
.to()accepts every prior thatprocess_prior()accepts.move_distribution_to_device()already covers plain torch distributions..to()changes device state only. The cached MAP and the other attributes stay.
Once both hold, NeuralPosterior.__setstate__ can call self.to(actual_device) and the attribute-by-attribute reconciliation from #1972 can go.
📌 Additional Context
- Found while reviewing #1972. The
.to()methods live in
sbi/inference/posteriors/*_posterior.py;move_distribution_to_device()is in
sbi/utils/user_input_checks_utils.py. VIPosterior.to()andBasePotential.to()already do it right: they move state in place
and do not re-run__init__.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the .to() implementations in sbi/inference/posteriors/*_posterior.py, then compare them with VIPosterior.to() and BasePotential.to(). Read move_distribution_to_device() in sbi/utils/user_input_checks_utils.py and verify that plain torch priors are supported while cached _map, _purpose, and other posterior state remain unchanged after moving devices.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100