## Shared Distribution Utilities
# We will represent each distribution by two callables:
# - `pdf(x)` returning the probability density evaluated on a 1D tensor `x`.
# - `sample(n, base_noise)` producing differentiable samples (re-using base noise when we need gradients).
# To keep gradients tractable, we lean on reparameterizations for Gaussians and deterministic mixtures of Gaussians. For numerical integrals we discretize the domain on a dense grid.
# Dense grid for numerical integration (wide enough that a model Gaussian with |mu| up to 15 stays on it);
# density plots are clipped to PLOT_XLIM.
GRID_X = torch.linspace(-20.0, 20.0, 8192, dtype=DEFAULT_DTYPE, device=device)
PLOT_XLIM = (-6.0, 6.0)
DX = GRID_X[1] - GRID_X[0]
EPS = torch.tensor(1e-12, dtype=DEFAULT_DTYPE, device=device)
def gaussian_pdf_torch(x: torch.Tensor, mean: torch.Tensor | float, std: torch.Tensor | float) -> torch.Tensor:
"""Evaluate the Gaussian pdf on `x` using broadcasted parameters."""
mean_t = torch.as_tensor(mean, dtype=DEFAULT_DTYPE, device=device)
std_t = torch.as_tensor(std, dtype=DEFAULT_DTYPE, device=device)
var = std_t ** 2
norm_const = 1.0 / torch.sqrt(2.0 * math.pi * var)
z = (x - mean_t) ** 2 / (2.0 * var)
return norm_const * torch.exp(-z)
def gaussian_log_pdf_torch(x: torch.Tensor, mean: torch.Tensor | float, std: torch.Tensor | float) -> torch.Tensor:
"""Log of the Gaussian pdf, evaluated directly so that far tails do not underflow to zero."""
mean_t = torch.as_tensor(mean, dtype=DEFAULT_DTYPE, device=device)
std_t = torch.as_tensor(std, dtype=DEFAULT_DTYPE, device=device)
return -0.5 * ((x - mean_t) / std_t) ** 2 - torch.log(std_t) - 0.5 * math.log(2.0 * math.pi)
def uniform_pdf_torch(x: torch.Tensor, low: float, high: float) -> torch.Tensor:
low_t = torch.as_tensor(low, dtype=DEFAULT_DTYPE, device=device)
high_t = torch.as_tensor(high, dtype=DEFAULT_DTYPE, device=device)
inside = (x >= low_t) & (x <= high_t)
return inside.to(DEFAULT_DTYPE) / (high_t - low_t + EPS)
def gaussian_mixture_pdf_torch(
x: torch.Tensor,
means: torch.Tensor,
stds: torch.Tensor,
weights: torch.Tensor,
) -> torch.Tensor:
comps = [w * gaussian_pdf_torch(x, m, s) for w, m, s in zip(weights, means, stds)]
return torch.stack(comps, dim=0).sum(dim=0)
@dataclass(frozen=True)
class DistributionSpec:
name: str
pdf: Callable[[torch.Tensor], torch.Tensor]
sampler: Callable[[int, torch.Tensor | None], torch.Tensor]
metadata: Dict[str, float]
def make_gaussian(mean: float = 0.0, std: float = 1.0) -> DistributionSpec:
mean_t = torch.tensor(mean, dtype=DEFAULT_DTYPE, device=device)
std_t = torch.tensor(std, dtype=DEFAULT_DTYPE, device=device)
def _pdf(x: torch.Tensor) -> torch.Tensor:
return gaussian_pdf_torch(x, mean_t, std_t)
def _sample(n: int, base_noise: torch.Tensor | None = None) -> torch.Tensor:
eps = torch.randn(n, dtype=DEFAULT_DTYPE, device=device) if base_noise is None else base_noise
return mean_t + std_t * eps
meta = {"mean": float(mean), "std": float(std)}
return DistributionSpec(name=f"Gaussian(μ={mean:.2f}, σ={std:.2f})", pdf=_pdf, sampler=_sample, metadata=meta)
def make_uniform(low: float = -2.0, high: float = 2.0) -> DistributionSpec:
low_t = torch.tensor(low, dtype=DEFAULT_DTYPE, device=device)
high_t = torch.tensor(high, dtype=DEFAULT_DTYPE, device=device)
def _pdf(x: torch.Tensor) -> torch.Tensor:
return uniform_pdf_torch(x, float(low_t), float(high_t))
def _sample(n: int, base_noise: torch.Tensor | None = None) -> torch.Tensor:
if base_noise is None:
base_noise = torch.rand(n, dtype=DEFAULT_DTYPE, device=device)
return low_t + (high_t - low_t) * base_noise
meta = {"low": float(low), "high": float(high)}
return DistributionSpec(name=f"Uniform[{low:.1f}, {high:.1f}]", pdf=_pdf, sampler=_sample, metadata=meta)
def make_gaussian_mixture(
means: Tuple[float, float] = (-2.0, 2.0),
stds: Tuple[float, float] = (0.6, 0.6),
weights: Tuple[float, float] = (0.5, 0.5),
) -> DistributionSpec:
means_t = torch.tensor(means, dtype=DEFAULT_DTYPE, device=device)
stds_t = torch.tensor(stds, dtype=DEFAULT_DTYPE, device=device)
weights_t = torch.tensor(weights, dtype=DEFAULT_DTYPE, device=device)
weights_t = weights_t / weights_t.sum()
def _pdf(x: torch.Tensor) -> torch.Tensor:
return gaussian_mixture_pdf_torch(x, means_t, stds_t, weights_t)
def _sample(n: int, base_noise: torch.Tensor | None = None) -> torch.Tensor:
# `base_noise` is standard-normal noise shared with the model (see sample_with_matching_noise);
# the component of each sample is assigned deterministically in the mixture proportions.
eps = torch.randn(n, dtype=DEFAULT_DTYPE, device=device) if base_noise is None else base_noise
cumulative = torch.cumsum(weights_t, dim=0)
u = torch.linspace(0.0, 1.0, n, dtype=DEFAULT_DTYPE, device=device)
choices = torch.clamp(torch.bucketize(u, cumulative), max=len(means_t) - 1)
return means_t[choices] + stds_t[choices] * eps
meta = {
"means": tuple(float(m) for m in means),
"stds": tuple(float(s) for s in stds),
"weights": tuple(float(w) for w in weights),
}
desc = ", ".join([f"μ={m:.1f}" for m in means])
return DistributionSpec(name=f"Gaussian Mixture ({desc})", pdf=_pdf, sampler=_sample, metadata=meta)
