FullSupportBarDistribution.sample()/cdf()/icdf() disagree with forward() in outer buckets

FullSupportBarDistribution.sample()/cdf()/icdf() disagree with forward() in outer buckets

Summary

FullSupportBarDistribution.forward() evaluates the first and last components as half-normal tails with support on all of R, but cdf(), icdf(), and sample() inherit the finite, within-bin-uniform implementation from BarDistribution.

Consequently, values returned by sample() are not distributed according to the density returned by -forward(). The same inconsistency affects tail quantiles, median(), and methods such as ucb() that call the inherited icdf().

This is still present in PriorLabs/TabPFN commit 930ea728b40f7f321d965250838bca8507bd70b9 (2026-08-20):

• FullSupportBarDistribution.forward() adds the half-normal tail density.
• FullSupportBarDistribution.sample() calls self.icdf(...).
• icdf() is inherited from BarDistribution and linearly interpolates between the two finite borders of every selected bucket, including both outer buckets.

Minimal reproduction

import torch

try:
    from tabpfn.architectures.shared.bar_distribution import FullSupportBarDistribution
except ImportError:
    from tabpfn.model.bar_distribution import FullSupportBarDistribution

torch.manual_seed(7)
dist = FullSupportBarDistribution(torch.tensor([-2.0, -1.0, 1.0, 2.0]))
probs = torch.tensor([0.25, 0.50, 0.25])
logits = probs.log().repeat(20_000, 1)

samples = dist.sample(logits)
buckets = dist.map_to_bucket_idx(samples).clamp(0, dist.num_bars - 1)
reported_logq = -dist(logits, samples)
sample_logq = probs[buckets].log() - dist.bucket_widths[buckets].log()
edge = (buckets == 0) | (buckets == dist.num_bars - 1)

print(samples.min().item(), samples.max().item())
print(((samples < -2.0) | (samples > 2.0)).sum().item())
print((reported_logq[edge] - sample_logq[edge]).min().item())
print((reported_logq[edge] - sample_logq[edge]).max().item())

Observed:

sample range: approximately [-2, 2]
outside [-2, 2]: 0 / 20,000
reported_logq - sample_logq in edge buckets: [-0.847, -0.620]

The documented full-support construction puts half of each outer component's mass beyond its finite outer border. With probabilities [0.25, 0.50, 0.25], the expected outside fraction is therefore

0.5 * (0.25 + 0.25) = 0.25,

or about 5,000 of 20,000 samples, rather than zero.

The complete reproducible audit also compares empirical and analytic CDFs. It obtains KS distance 0.00585 to the inherited finite ICDF distribution versus 0.125 to the density implemented by forward().

Expected behavior

For borders b[0] < ... < b[K] and component probabilities pi, the two outer components should be sampled from the same half-normal densities used by forward():

left:  Y = b[1]   - HalfNormal(scale_left)
right: Y = b[K-1] + HalfNormal(scale_right)

where each scale keeps half of the component mass inside its nominal outer bin:

scale_left  = (b[1] - b[0])     / HalfNormal(1).icdf(0.5)
scale_right = (b[K] - b[K - 1]) / HalfNormal(1).icdf(0.5)

For a uniform variate u that selects the left component, with conditional rank r = u / pi[0], the inverse CDF should be

b[1] - HalfNormal(scale_left).icdf(1 - r).

For the right component it should be

b[K-1] + HalfNormal(scale_right).icdf(r).

Interior buckets should retain their current within-bin-uniform interpolation.

Requested fix

1. Override cdf() and icdf() in FullSupportBarDistribution with the half-normal formulas already used by forward().
2. Make sample() use that corrected inverse CDF, or sample the selected outer components directly from their half-normal distributions.
3. Preserve batch shape, device, and dtype; avoid rebuilding the result through torch.tensor(list_of_tensors), which can move CUDA results back to CPU.
4. Keep the existing borders, component probabilities, and half-normal scale convention so checkpoint likelihoods remain unchanged.
5. Audit inherited methods that depend on cdf()/icdf(), including quantile(), median(), ucb(), and border-probability translation.
6. State the affected releases and include the correction in the release notes.

Proposed regression tests

1. Deterministic inverse-CDF checkpoints. For borders [-2,-1,1,2] and probabilities [0.25,0.50,0.25], verify icdf(0.125) == -2 and icdf(0.875) == 2 within numerical tolerance. The current inherited implementation returns -1.5 and 1.5.
2. CDF/ICDF round trip. Check cdf(logits, icdf(logits, u)) == u for ranks inside both tails and all interior buckets.
3. Sampling tail mass. A seeded large sample should place approximately 25% of samples outside [-2,2] in the reproduction above.
4. Density/sampling consistency. A goodness-of-fit test should accept the distribution defined by forward() and reject the finite inherited-ICDF law.
5. Device/dtype coverage. Test CPU and CUDA when available, including batched logits and float32/float64.

Why this matters

The mismatch invalidates any use that assumes samples and density values refer to the same distribution, including importance sampling, Monte Carlo expectations, density estimation, generative sampling, and calibrated tail quantiles. For importance sampling it can additionally violate the proposal support condition because the implemented sampler has finite support while the reported proposal density is full-support.