fix: make bar distribution reductions device-agnostic

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.

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(

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