From bf046a1c1bf40d9777afabfd0d97cc5de01b8088 Mon Sep 17 00:00:00 2001 From: ppguo <17300750016@fudan.edu.cn> Date: Thu, 27 Aug 2026 22:30:19 +0800 Subject: [PATCH 1/4] Fix FullSupportBarDistribution tail sampling and quantiles --- changelog/PR.fixed.md | 1 + .../architectures/shared/bar_distribution.py | 110 ++++++++++++++++-- .../test_bar_distribution.py | 100 ++++++++++++++++ 3 files changed, 204 insertions(+), 7 deletions(-) create mode 100644 changelog/PR.fixed.md diff --git a/changelog/PR.fixed.md b/changelog/PR.fixed.md new file mode 100644 index 000000000..f117d4248 --- /dev/null +++ b/changelog/PR.fixed.md @@ -0,0 +1 @@ +Fix `FullSupportBarDistribution` CDF, quantiles, and sampling to match its half-normal tails while preserving batch shape, device, and dtype. diff --git a/src/tabpfn/architectures/shared/bar_distribution.py b/src/tabpfn/architectures/shared/bar_distribution.py index 8411ceea6..1ddc734ef 100644 --- a/src/tabpfn/architectures/shared/bar_distribution.py +++ b/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 diff --git a/tests/test_architectures/test_bar_distribution.py b/tests/test_architectures/test_bar_distribution.py index f12b09ed4..7ef88d088 100644 --- a/tests/test_architectures/test_bar_distribution.py +++ b/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,94 @@ 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() + + assert torch.equal(dist.median(logits), dist.icdf(logits, 0.5)) + assert torch.equal( + dist.quantile(logits, center_prob=0.75), + torch.stack((dist.icdf(logits, 0.125), dist.icdf(logits, 0.875))), + ) + assert torch.equal( + dist.ucb(logits, best_f=0.0, rest_prob=0.125), + dist.icdf(logits, 0.875), + ) + assert torch.equal( + dist.ucb(logits, best_f=0.0, rest_prob=0.125, maximize=False), + dist.icdf(logits, 0.125), + ) + + new_borders = torch.tensor([-3.0, -2.0, 0.0, 2.0, 3.0]) + translated = dist.get_probs_for_different_borders(logits, new_borders) + assert torch.allclose(translated, torch.tensor([0.125, 0.375, 0.375, 0.125])) + + +@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() + + samples = dist.sample(batch_logits) + + 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 = dist.cdf(logits, ys) + 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] From d62c368fc40e32510303799ba16cc679fb10562a Mon Sep 17 00:00:00 2001 From: ppguo <17300750016@fudan.edu.cn> Date: Thu, 27 Aug 2026 22:33:16 +0800 Subject: [PATCH 2/4] Name changelog fragment for PR 1215 --- changelog/{PR.fixed.md => 1215.fixed.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changelog/{PR.fixed.md => 1215.fixed.md} (100%) diff --git a/changelog/PR.fixed.md b/changelog/1215.fixed.md similarity index 100% rename from changelog/PR.fixed.md rename to changelog/1215.fixed.md From 3b4d7278cc525f4ac9365af011d57224f65d6b3c Mon Sep 17 00:00:00 2001 From: ppguo <17300750016@fudan.edu.cn> Date: Thu, 27 Aug 2026 22:36:46 +0800 Subject: [PATCH 3/4] Strengthen full-support distribution regression coverage --- .../test_bar_distribution.py | 33 +++++++++++++------ 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/tests/test_architectures/test_bar_distribution.py b/tests/test_architectures/test_bar_distribution.py index 7ef88d088..568c8d1ee 100644 --- a/tests/test_architectures/test_bar_distribution.py +++ b/tests/test_architectures/test_bar_distribution.py @@ -85,24 +85,29 @@ def test_full_support_cdf_icdf_round_trip(left_prob: float): 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(logits), dist.icdf(logits, 0.5)) + assert torch.equal(dist.median(batch_logits), dist.icdf(batch_logits, 0.5)) assert torch.equal( - dist.quantile(logits, center_prob=0.75), - torch.stack((dist.icdf(logits, 0.125), dist.icdf(logits, 0.875))), + 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(logits, best_f=0.0, rest_prob=0.125), - dist.icdf(logits, 0.875), + dist.ucb(batch_logits, best_f=0.0, rest_prob=0.125), + dist.icdf(batch_logits, 0.875), ) assert torch.equal( - dist.ucb(logits, best_f=0.0, rest_prob=0.125, maximize=False), - dist.icdf(logits, 0.125), + 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(logits, new_borders) - assert torch.allclose(translated, torch.tensor([0.125, 0.375, 0.375, 0.125])) + 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]) @@ -117,8 +122,13 @@ def test_full_support_sample_preserves_shape_device_and_dtype( 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 @@ -135,7 +145,10 @@ def test_full_support_sample_matches_tail_mass_and_cdf(): 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 = dist.cdf(logits, 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) From e8312b2a2ad1616144e4224da400c9fede892ef6 Mon Sep 17 00:00:00 2001 From: ppguo <17300750016@fudan.edu.cn> Date: Thu, 27 Aug 2026 23:33:24 +0800 Subject: [PATCH 4/4] Clamp full-support conditional quantile probabilities --- src/tabpfn/architectures/shared/bar_distribution.py | 1 + tests/test_architectures/test_bar_distribution.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/src/tabpfn/architectures/shared/bar_distribution.py b/src/tabpfn/architectures/shared/bar_distribution.py index 1ddc734ef..3a8caebed 100644 --- a/src/tabpfn/architectures/shared/bar_distribution.py +++ b/src/tabpfn/architectures/shared/bar_distribution.py @@ -568,6 +568,7 @@ def icdf(self, logits: torch.Tensor, left_prob: float) -> torch.Tensor: conditional_prob = ( left_prob - cumprobs_before.gather(-1, idx[..., None]).squeeze(-1) ) / selected_probs + conditional_prob = conditional_prob.clamp(0.0, 1.0) values = self.borders[idx] + self.bucket_widths[idx] * conditional_prob side_normals = ( diff --git a/tests/test_architectures/test_bar_distribution.py b/tests/test_architectures/test_bar_distribution.py index 568c8d1ee..1de42eeb9 100644 --- a/tests/test_architectures/test_bar_distribution.py +++ b/tests/test_architectures/test_bar_distribution.py @@ -83,6 +83,13 @@ def test_full_support_cdf_icdf_round_trip(left_prob: float): ) +def test_full_support_icdf_clamps_conditional_probability_roundoff(): + dist, _ = _make_full_support_distribution() + logits = torch.tensor([0.40334684, 0.83802634, -0.7192576]) + + assert torch.isposinf(dist.icdf(logits, 1.0)) + + def test_full_support_inherited_quantiles_and_border_translation(): dist, logits = _make_full_support_distribution() batch_logits = logits.expand(2, -1)