fix: reconcile posterior device after mapped load

[BHARATH0153]

Summary

Fix stale posterior device state after loading with torch.load(..., map_location=...).

Changes

• Reconcile loaded posterior state with the device of its tensors.
• Move potential functions, distributions, estimators, stored observations, and vector-field state when needed.
• Route VIPosterior restoration through the base reconciliation and its own to() method.
• Add mapped-load coverage for direct, potential-based, VI, and vector-field posterior paths, including post-load runtime operations.

Closes #1954 Supersedes #1959

Ran them locally on MPS since CI skips gpu-marked tests.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 318744fb-f4d5-4033-aa00-42f1e591dde6

📥 Commits

Reviewing files that changed from the base of the PR and between bb0cc2e and 9189ba6.

📒 Files selected for processing (1)

• tests/inference_on_device_test.py

🚧 Files skipped from review as they are similar to previous changes (1)

• tests/inference_on_device_test.py

---

📝 Walkthrough

Walkthrough

This change fixes stale device state after torch.load with map_location. It adds recursive device inference and reconciles posterior components after unpickling. Tests cover CPU restoration for multiple posterior types.

Changes

Device reconciliation on posterior load

Layer / File(s)	Summary
Recursive tensor device inference sbi/utils/torchutils.py	Adds infer_tensor_device, which traverses object graphs and detects tensor or module devices with cycle detection.
Posterior device reconciliation after loading sbi/inference/posteriors/base_posterior.py, sbi/inference/posteriors/vi_posterior.py	NeuralPosterior.__setstate__ detects the restored device and moves posterior components when needed. VIPosterior.__setstate__ delegates to the base implementation.
Device reconciliation regression tests tests/inference_on_device_test.py	Adds GPU-gated serialization tests for Direct, rejection, importance-sampling, MCMC, VI, and vector-field posteriors loaded onto CPU.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• sbi-dev/sbi#1959: Modifies related posterior deserialization and device-inference logic.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the primary change: reconciling posterior device state after mapped loading.
Description check	✅ Passed	The description directly explains the device reconciliation fix, affected posterior types, tests, and linked issues.
Linked Issues check	✅ Passed	The changes address issue #1954 by reconciling devices during unpickling and testing mapped loads across posterior types.
Out of Scope Changes check	✅ Passed	The code and test changes remain focused on mapped-load device reconciliation and related runtime validation.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 49.12281% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.83%. Comparing base (500483d) to head (9189ba6).

Files with missing lines	Patch %	Lines
sbi/inference/posteriors/base_posterior.py	30.43%	16 Missing ⚠️
sbi/utils/torchutils.py	59.37%	13 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1972      +/-   ##
==========================================
- Coverage   88.98%   88.83%   -0.15%     
==========================================
  Files         139      139              
  Lines       13604    13658      +54     
==========================================
+ Hits        12105    12133      +28     
- Misses       1499     1525      +26     

Flag	Coverage Δ
fast	83.36% <49.12%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
sbi/inference/posteriors/vi_posterior.py	85.97% <100.00%> (+0.72%)	⬆️
sbi/utils/torchutils.py	76.59% <59.37%> (-2.21%)	⬇️
sbi/inference/posteriors/base_posterior.py	72.54% <30.43%> (-12.64%)	⬇️

[BHARATH0153]

@janfb created new pr sorry for closing that Thanks Jan for the detailed review. I addressed the remaining VI issue by calling super().__setstate__(state_dict) followed by self.to(self._device), and added post-load sample() and log_prob() assertions for VI.

[BHARATH0153]

@coderabbitai is it covering all the issue and jan review in #1959

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo '=== PR `#1959` review comments and conversation ==='
gh api --paginate repos/sbi-dev/sbi/pulls/1959/comments \
  --jq '.[] | {user: .user.login, path, line, original_line, body, html_url}'
gh api --paginate repos/sbi-dev/sbi/issues/1959/comments \
  --jq '.[] | {user: .user.login, body, html_url}'

echo '=== Current PR changed files ==='
gh pr view 1972 --repo sbi-dev/sbi --json files --jq '.files[] | "\(.path) (+\(.additions)/-\(.deletions))"'

echo '=== Source maps ==='
ast-grep outline sbi/inference/posteriors/base_posterior.py --items all --match '__setstate__|to|NeuralPosterior'
ast-grep outline sbi/inference/posteriors/vi_posterior.py --items all --match '__setstate__|to|VIPosterior'
ast-grep outline sbi/utils/torchutils.py --items all --match 'infer_tensor_device'

echo '=== Relevant implementations ==='
rg -n -A75 -B12 'def __setstate__|def infer_tensor_device' \
  sbi/inference/posteriors/base_posterior.py \
  sbi/inference/posteriors/vi_posterior.py \
  sbi/utils/torchutils.py

echo '=== New mapped-load tests ==='
rg -n -A55 -B8 'map_location|infer_tensor_device|reconcile|loaded.*device' \
  tests/inference_on_device_test.py

Length of output: 59358

---

@BHARATH0153 Yes. The current PR covers the technical issues raised in Jan’s review on #1959 and the device-state problem in #1954.

Covered items:

• VIPosterior now delegates restoration to NeuralPosterior.__setstate__ and then calls VIPosterior.to(self._device). This restores VI-specific state such as distributions and transforms.
• The VI test constructs VIPosterior with a valid potential_fn and prior.
• Device inference reaches potential-based posteriors and VectorFieldPosterior.
• Reconciliation updates both _device and device when the posterior has a device attribute.
• The vector-field test checks loaded.device == "cpu" and executes loaded.sample(...).
• Direct and potential-based posterior tests execute both sample() and potential() after map_location="cpu".
• The VI test executes sample() and log_prob() after loading.
• Tests check that estimators, potential functions, and stored state are on CPU after the mapped load.

One validation limit remains: I can verify the implementation and test coverage from the PR, but I cannot independently confirm local GPU-marked test execution, full pytest, or pyright execution from the available evidence. Please run those checks before merge if they are available.

_{You are interacting with an AI system.}

[BHARATH0153]

@janfb may I know how you would prefer for easy testing my local venv taking more time your suggestion will be helpful thanks!

[janfb]

@janfb may I know how you would prefer for easy testing my local venv taking more time your suggestion will be helpful thanks!

you need to test locally only the specific GPU tests that are affected by your changes, not the entire test suite. I hope this helps.

[BHARATH0153]

thanks for the suggestion

$ pytest tests/inference_on_device_test.py -k "pickle_map_location" -v
================ 6 passed, 239 deselected, 5 warnings in 3.24s ================. 6/6 passed: the test_pickle_map_location_ tests (Direct, rejection, importance, MCMC, VI, vector-field) all ran on MPS and passed after today's VI fix.