From 83eae4fa0fb8a138939b22a4254d2459e2367d2a Mon Sep 17 00:00:00 2001 From: Georg Grab Date: Mon, 10 Aug 2026 16:29:34 +0200 Subject: [PATCH 1/2] fix: make bar distribution reductions device-agnostic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `FullSupportBarDistribution.mean` re-homes its buffer onto the logits' device; `median`, `mode` and `icdf` index `borders` with logit-derived indices and did not. A criterion sitting on a different device than the logits therefore failed every reduction except `mean`, with RuntimeError: indices should be either on cpu or on the same device as the indexed tensor That is reachable whenever a fitted estimator is reassembled and only part of it is placed — restoring one and moving the weights but not the bar distribution, for instance. `mean` silently working made it look like a partial failure rather than a placement problem. Index a same-device view of the borders in `icdf` and `mode`, matching what `mean` already does. `quantile` and `ucb` inherit it via `icdf`. `bucket_widths` derives from the borders, so it follows. The `.to()` is a no-op when the devices already match: results are bit-identical on the co-located path, and match to float32 tolerance across devices. Co-Authored-By: Claude Opus 5 --- .../architectures/shared/bar_distribution.py | 15 ++-- tests/test_bar_distribution_device.py | 69 +++++++++++++++++++ 2 files changed, 80 insertions(+), 4 deletions(-) create mode 100644 tests/test_bar_distribution_device.py diff --git a/src/tabpfn/architectures/shared/bar_distribution.py b/src/tabpfn/architectures/shared/bar_distribution.py index 8411ceea6..74c736bce 100644 --- a/src/tabpfn/architectures/shared/bar_distribution.py +++ b/src/tabpfn/architectures/shared/bar_distribution.py @@ -293,8 +293,11 @@ def icdf(self, logits: torch.Tensor, left_prob: float) -> torch.Tensor: ) rest_prob = left_prob - cumprobs.gather(-1, idx[..., None]).squeeze(-1) - left_border = self.borders[idx] - right_border = self.borders[idx + 1] + # `idx` comes from `logits`, so index a same-device view of the borders. + # No-op when they already match. + borders = self.borders.to(logits.device) + left_border = borders[idx] + right_border = borders[idx + 1] return left_border + (right_border - left_border) * rest_prob / probs.gather( -1, idx[..., None], @@ -346,9 +349,13 @@ def ucb( def mode(self, logits: torch.Tensor) -> torch.Tensor: """Mode of the distribution (center of the highest-density bucket).""" - density = logits.softmax(-1) / self.bucket_widths + # Same-device view of the borders; `bucket_widths` derives from them. + # No-op when they already match. + borders = self.borders.to(logits.device) + widths = borders[1:] - borders[:-1] + density = logits.softmax(-1) / widths mode_inds = density.argmax(-1) - bucket_means = self.borders[:-1] + self.bucket_widths / 2 + bucket_means = borders[:-1] + widths / 2 return bucket_means[mode_inds] def ei( diff --git a/tests/test_bar_distribution_device.py b/tests/test_bar_distribution_device.py new file mode 100644 index 000000000..12253d746 --- /dev/null +++ b/tests/test_bar_distribution_device.py @@ -0,0 +1,69 @@ +# Copyright (c) Prior Labs GmbH 2026. + +"""Reductions must work when the criterion and the logits are on different devices. + +`mean` already re-homed its buffer; `median`, `mode` and `icdf` indexed +`borders` directly and raised. That can happen whenever an estimator is +reassembled and only part of it is moved — e.g. restoring a fitted regressor +and placing the weights but not the bar distribution. +""" + +from __future__ import annotations + +import pytest +import torch + +from tabpfn.architectures.shared.bar_distribution import FullSupportBarDistribution + + +def _second_device() -> torch.device | None: + if torch.cuda.is_available(): + return torch.device("cuda", 0) + if torch.backends.mps.is_available(): + return torch.device("mps") + return None + + +def _reductions(bd: FullSupportBarDistribution, logits: torch.Tensor) -> dict: + return { + "mean": bd.mean(logits), + "median": bd.median(logits), + "mode": bd.mode(logits), + "icdf": bd.icdf(logits, 0.9), + "quantile": bd.quantile(logits), + } + + +@pytest.fixture +def borders() -> torch.Tensor: + return torch.linspace(-3, 3, 101) + + +@pytest.fixture +def logits() -> torch.Tensor: + return torch.randn(64, 100, generator=torch.Generator().manual_seed(0)) + + +def test_reductions_run_with_borders_on_another_device(borders, logits): + device = _second_device() + if device is None: + pytest.skip("needs a second device") + + expected = _reductions(FullSupportBarDistribution(borders.clone()), logits) + # Borders left on CPU, logits on the compute device. + got = _reductions(FullSupportBarDistribution(borders.clone()), logits.to(device)) + + for name, value in got.items(): + assert value.device.type == device.type, name + torch.testing.assert_close( + value.cpu().float(), expected[name].float(), rtol=0, atol=1e-4, msg=name + ) + + +def test_co_located_results_are_unchanged(borders, logits): + """The re-homing must be a no-op when the devices already match.""" + bd = FullSupportBarDistribution(borders) + first = _reductions(bd, logits) + second = _reductions(bd, logits) + for name, value in first.items(): + assert torch.equal(value, second[name]), name From afb1eb2f511d9d8b72bcc4e805b68664ffc8fa0e Mon Sep 17 00:00:00 2001 From: Georg Grab Date: Mon, 10 Aug 2026 16:51:27 +0200 Subject: [PATCH 2/2] docs: add changelog entry for #1170 --- changelog/1170.fixed.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog/1170.fixed.md diff --git a/changelog/1170.fixed.md b/changelog/1170.fixed.md new file mode 100644 index 000000000..be24e1ef4 --- /dev/null +++ b/changelog/1170.fixed.md @@ -0,0 +1 @@ +Regressor output types other than `mean` no longer fail when the bar distribution and the logits are on different devices. \ No newline at end of file