Fix FullSupportBarDistribution tail sampling and quantiles

changelog/1215.fixed.md

@@ -0,0 +1 @@
+Fix `FullSupportBarDistribution` CDF, quantiles, and sampling to match its half-normal tails while preserving batch shape, device, and dtype.

src/tabpfn/architectures/shared/bar_distribution.py

@@ -501,15 +501,88 @@ def assert_support(self, *, allow_zero_bucket_left: bool = False) -> None:
 
     @staticmethod
     def halfnormal_with_p_weight_before(
-        range_max: float,
+        range_max: float | torch.Tensor,
         p: float = 0.5,
     ) -> torch.distributions.HalfNormal:
         """Build a half-normal placing ``p`` of its mass below ``range_max``."""
-        s = range_max / torch.distributions.HalfNormal(torch.tensor(1.0)).icdf(
-            torch.tensor(p),
-        )
+        range_max = torch.as_tensor(range_max)
+        unit_halfnormal = torch.distributions.HalfNormal(torch.ones_like(range_max))
+        s = range_max / unit_halfnormal.icdf(torch.full_like(range_max, p))
         return torch.distributions.HalfNormal(s)
 
+    @override
+    def cdf(self, logits: torch.Tensor, ys: torch.Tensor) -> torch.Tensor:
+        """Calculate the CDF, including the two half-normal tails."""
+        if len(ys.shape) < len(logits.shape) and len(ys.shape) == 1:
+            ys = ys.repeat((*logits.shape[:-1], 1))
+        else:
+            assert ys.shape[:-1] == logits.shape[:-1], (
+                f"ys.shape: {ys.shape} logits.shape: {logits.shape}"
+            )
+
+        prob_left_of_ys = super().cdf(logits, ys)
+        probs = torch.softmax(logits, dim=-1)
+        side_normals = (
+            self.halfnormal_with_p_weight_before(self.bucket_widths[0]),
+            self.halfnormal_with_p_weight_before(self.bucket_widths[-1]),
+        )
+
+        left_tail_cdf = probs[..., 0, None] * (
+            1.0 - side_normals[0].cdf((self.borders[1] - ys).clamp_min(0.0))
+        )
+        right_tail_cdf = 1.0 - probs[..., -1, None] * (
+            1.0 - side_normals[1].cdf((ys - self.borders[-2]).clamp_min(0.0))
+        )
+
+        prob_left_of_ys = torch.where(
+            ys < self.borders[1],
+            left_tail_cdf,
+            prob_left_of_ys,
+        )
+        prob_left_of_ys = torch.where(
+            ys >= self.borders[-2],
+            right_tail_cdf,
+            prob_left_of_ys,
+        )
+        return prob_left_of_ys.clip(0.0, 1.0)
+
+    @override
+    def icdf(self, logits: torch.Tensor, left_prob: float) -> torch.Tensor:
+        """Calculate quantiles using half-normal tails in the outer buckets."""
+        probs = logits.softmax(-1)
+        cumprobs = torch.cumsum(probs, -1)
+        left_prob_tensor = torch.full(
+            (*cumprobs.shape[:-1], 1),
+            left_prob,
+            dtype=logits.dtype,
+            device=logits.device,
+        )
+        idx = torch.searchsorted(cumprobs, left_prob_tensor).squeeze(-1)
+        idx = idx.clamp(0, self.num_bars - 1)
+
+        cumprobs_before = torch.cat(
+            (torch.zeros_like(cumprobs[..., :1]), cumprobs[..., :-1]),
+            dim=-1,
+        )
+        selected_probs = probs.gather(-1, idx[..., None]).squeeze(-1)
+        conditional_prob = (
+            left_prob - cumprobs_before.gather(-1, idx[..., None]).squeeze(-1)
+        ) / selected_probs
+
+        values = self.borders[idx] + self.bucket_widths[idx] * conditional_prob
+        side_normals = (
+            self.halfnormal_with_p_weight_before(self.bucket_widths[0]),
+            self.halfnormal_with_p_weight_before(self.bucket_widths[-1]),
+        )
+        left_tail_values = self.borders[1] - side_normals[0].icdf(
+            1.0 - conditional_prob,
+        )
+        right_tail_values = self.borders[-2] + side_normals[1].icdf(
+            conditional_prob,
+        )
+        values = torch.where(idx == 0, left_tail_values, values)
+        return torch.where(idx == self.num_bars - 1, right_tail_values, values)
+
     @override
     def forward(
         self,
@@ -606,9 +679,32 @@ def sample(self, logits: torch.Tensor, t: float = 1.0) -> torch.Tensor:
 
         Temperature t.
         """
-        p_cdf = torch.rand(*logits.shape[:-1])
-        return torch.tensor(
-            [self.icdf(logits[i, :] / t, p) for i, p in enumerate(p_cdf.tolist())],
+        bucket_indices = torch.distributions.Categorical(logits=logits / t).sample()
+        uniform_samples = torch.rand(
+            bucket_indices.shape,
+            dtype=logits.dtype,
+            device=logits.device,
+        )
+        samples = (
+            self.borders[bucket_indices]
+            + self.bucket_widths[bucket_indices] * uniform_samples
+        )
+
+        side_normals = (
+            self.halfnormal_with_p_weight_before(self.bucket_widths[0]),
+            self.halfnormal_with_p_weight_before(self.bucket_widths[-1]),
+        )
+        left_tail_samples = self.borders[1] - side_normals[0].sample(
+            bucket_indices.shape,
+        )
+        right_tail_samples = self.borders[-2] + side_normals[1].sample(
+            bucket_indices.shape,
+        )
+        samples = torch.where(bucket_indices == 0, left_tail_samples, samples)
+        return torch.where(
+            bucket_indices == self.num_bars - 1,
+            right_tail_samples,
+            samples,
         )
 
     @override

tests/test_architectures/test_bar_distribution.py

@@ -6,6 +6,18 @@
 import torch
 
 from tabpfn.architectures.shared import bar_distribution
+from tests.utils import get_pytest_devices_with_mps_marked_slow
+
+
+def _make_full_support_distribution(
+    *,
+    dtype: torch.dtype = torch.float32,
+    device: str = "cpu",
+) -> tuple[bar_distribution.FullSupportBarDistribution, torch.Tensor]:
+    borders = torch.tensor([-2.0, -1.0, 1.0, 2.0], dtype=dtype, device=device)
+    dist = bar_distribution.FullSupportBarDistribution(borders)
+    logits = torch.tensor([0.25, 0.5, 0.25], dtype=dtype, device=device).log()
+    return dist, logits
 
 
 def test_cdf_out_of_bounds():
@@ -39,6 +51,107 @@ def test_move_to_larger():
     )
 
 
+def test_full_support_cdf_and_icdf_checkpoints():
+    dist, logits = _make_full_support_distribution()
+
+    assert dist.icdf(logits, 0.125).item() == pytest.approx(-2.0)
+    assert dist.icdf(logits, 0.875).item() == pytest.approx(2.0)
+
+    ys = torch.tensor([float("-inf"), -2.0, -1.0, 0.0, 1.0, 2.0, float("inf")])
+    expected = torch.tensor([0.0, 0.125, 0.25, 0.5, 0.75, 0.875, 1.0])
+    assert torch.allclose(dist.cdf(logits, ys), expected)
+
+
+@pytest.mark.parametrize(
+    "left_prob",
+    [0.0, 0.001, 0.125, 0.249, 0.25, 0.5, 0.75, 0.875, 0.999, 1.0],
+)
+def test_full_support_cdf_icdf_round_trip(left_prob: float):
+    dist, logits = _make_full_support_distribution(dtype=torch.float64)
+    batch_logits = torch.stack(
+        (logits, torch.tensor([0.1, 0.3, 0.6], dtype=logits.dtype).log())
+    )
+
+    values = dist.icdf(batch_logits, left_prob)
+    actual = dist.cdf(batch_logits, values.unsqueeze(-1)).squeeze(-1)
+
+    assert torch.allclose(
+        actual,
+        torch.full_like(actual, left_prob),
+        atol=1e-12,
+        rtol=1e-12,
+    )
+
+
+def test_full_support_inherited_quantiles_and_border_translation():
+    dist, logits = _make_full_support_distribution()
+    batch_logits = logits.expand(2, -1)
+
+    assert torch.equal(dist.median(batch_logits), dist.icdf(batch_logits, 0.5))
+    assert torch.equal(
+        dist.quantile(batch_logits, center_prob=0.75),
+        torch.stack(
+            (dist.icdf(batch_logits, 0.125), dist.icdf(batch_logits, 0.875)),
+            dim=-1,
+        ),
+    )
+    assert torch.equal(
+        dist.ucb(batch_logits, best_f=0.0, rest_prob=0.125),
+        dist.icdf(batch_logits, 0.875),
+    )
+    assert torch.equal(
+        dist.ucb(batch_logits, best_f=0.0, rest_prob=0.125, maximize=False),
+        dist.icdf(batch_logits, 0.125),
+    )
+
+    new_borders = torch.tensor([-3.0, -2.0, 0.0, 2.0, 3.0])
+    translated = dist.get_probs_for_different_borders(batch_logits, new_borders)
+    expected = torch.tensor([0.125, 0.375, 0.375, 0.125]).expand(2, -1)
+    assert torch.allclose(translated, expected)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+@pytest.mark.parametrize("device", get_pytest_devices_with_mps_marked_slow())
+def test_full_support_sample_preserves_shape_device_and_dtype(
+    device: str,
+    dtype: torch.dtype,
+):
+    if device == "mps" and dtype == torch.float64:
+        pytest.skip("MPS does not support float64 tensors")
+
+    dist, logits = _make_full_support_distribution(dtype=dtype, device=device)
+    batch_logits = logits.expand(2, 3, -1).contiguous()
+
+    scalar_sample = dist.sample(logits)
+    samples = dist.sample(batch_logits)
+
+    assert scalar_sample.shape == logits.shape[:-1]
+    assert scalar_sample.device.type == torch.device(device).type
+    assert scalar_sample.dtype == dtype
+    assert torch.isfinite(scalar_sample)
+    assert samples.shape == batch_logits.shape[:-1]
+    assert samples.device.type == torch.device(device).type
+    assert samples.dtype == dtype
+    assert torch.isfinite(samples).all()
+
+
+def test_full_support_sample_matches_tail_mass_and_cdf():
+    torch.manual_seed(7)
+    dist, logits = _make_full_support_distribution(dtype=torch.float64)
+    samples = dist.sample(logits.repeat(20_000, 1))
+
+    outside = ((samples < dist.borders[0]) | (samples > dist.borders[-1])).double()
+    assert outside.mean().item() == pytest.approx(0.25, abs=0.015)
+
+    ys = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0], dtype=samples.dtype)
+    empirical_cdf = torch.stack([(samples <= y).double().mean() for y in ys])
+    expected_cdf = torch.tensor(
+        [0.125, 0.25, 0.5, 0.75, 0.875],
+        dtype=samples.dtype,
+    )
+    assert torch.allclose(empirical_cdf, expected_cdf, atol=0.015, rtol=0.0)
+
+
 def test_average_bar_distributions_into_different_one():
     num_bars = [100, 80, 10, 5]
     logits = [torch.arange(nb - 1).float() for nb in num_bars]