sbi-dev / sbi-dev/sbi

Posterior `.to()` methods fail on plain torch priors and drop cached state

Open
#2,004 3 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

bug
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.

  1. DirectPosterior.to() and VectorFieldPosterior.to() raise ValueError("Prior has no attribute to(device)") for plain torch priors such as MultivariateNormal. BasePotential.to() handles the same case with move_distribution_to_device().
  2. Every posterior .to() re-runs super().__init__(). That drops the cached _map and resets _purpose. A user who calls map() 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 that process_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() and BasePotential.to() already do it right: they move state in place
    and do not re-run __init__.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.