"""Objective-preserving primitives for InfoFlow training. This module intentionally depends only on PyTorch. It contains no dataset or model-specific policy, so its contracts can be tested before sampler wiring. """ from __future__ import annotations import math from dataclasses import dataclass import torch @dataclass(frozen=True) class SafeMixture: """A normalized safe proposal and its exact importance correction.""" base_probability: torch.Tensor proposal_probability: torch.Tensor sampling_probability: torch.Tensor importance_weight: torch.Tensor max_importance_weight: float def _normalized_probability(probability: torch.Tensor, *, name: str) -> torch.Tensor: if not isinstance(probability, torch.Tensor): raise TypeError(f"{name} must be a torch.Tensor") if probability.ndim != 1 or probability.numel() == 0: raise ValueError(f"{name} must be a non-empty one-dimensional tensor") if probability.dtype == torch.bool or not probability.dtype.is_floating_point: raise TypeError(f"{name} must have a floating-point dtype") probability64 = probability.detach().to(dtype=torch.float64) if not torch.isfinite(probability64).all(): raise ValueError(f"{name} must contain only finite values") if (probability64 < 0).any(): raise ValueError(f"{name} must be non-negative") total = probability64.sum() if total <= 0: raise ValueError(f"{name} must have a positive sum") return probability64 / total def build_safe_mixture( base_probability: torch.Tensor, proposal_probability: torch.Tensor, *, alpha: float, ) -> SafeMixture: """Build ``q = (1 - alpha) p + alpha r`` and exact ``p / q`` weights.""" if base_probability.device != proposal_probability.device: raise ValueError("base_probability and proposal_probability must be on the same device") if base_probability.shape != proposal_probability.shape: raise ValueError("base_probability and proposal_probability must have the same shape") if not isinstance(alpha, (int, float)) or isinstance(alpha, bool) or not math.isfinite(float(alpha)): raise TypeError("alpha must be a finite real number") alpha = float(alpha) if not 0.0 <= alpha < 1.0: raise ValueError("alpha must satisfy 0 <= alpha < 1") p = _normalized_probability(base_probability, name="base_probability") r = _normalized_probability(proposal_probability, name="proposal_probability") q = (1.0 - alpha) * p + alpha * r positive_support = p > 0 if not (q[positive_support] > 0).all(): raise RuntimeError("safe mixture lost support for a positive-probability baseline item") weight = torch.zeros_like(p) weight[positive_support] = p[positive_support] / q[positive_support] max_weight = 1.0 / (1.0 - alpha) if (weight > max_weight + 1e-12).any(): raise RuntimeError("safe-mixture importance bound was violated") return SafeMixture( base_probability=p, proposal_probability=r, sampling_probability=q, importance_weight=weight.detach(), max_importance_weight=max_weight, ) def reduce_flow_matching_loss(elementwise_loss: torch.Tensor, *, reduction: str = "mean") -> torch.Tensor: """Reduce flow loss while retaining the original batch dimension.""" if not isinstance(elementwise_loss, torch.Tensor): raise TypeError("elementwise_loss must be a torch.Tensor") if elementwise_loss.ndim == 0 or elementwise_loss.shape[0] == 0: raise ValueError("elementwise_loss must have a non-empty batch dimension") if reduction not in {"mean", "none"}: raise ValueError("reduction must be either 'mean' or 'none'") # Preserve the legacy scalar path's reduction order exactly. Computing # per-sample means first can introduce a tiny floating-point difference. if reduction == "mean": return elementwise_loss.mean() if elementwise_loss.ndim == 1: per_sample = elementwise_loss else: per_sample = elementwise_loss.mean(dim=tuple(range(1, elementwise_loss.ndim))) return per_sample def aggregate_repeated_losses( repeated_per_sample_loss: torch.Tensor, *, repeat_count: int, batch_size: int | None = None, ) -> torch.Tensor: """Average repeat-major stochastic flow draws back to shape ``[B]``.""" if not isinstance(repeated_per_sample_loss, torch.Tensor): raise TypeError("repeated_per_sample_loss must be a torch.Tensor") if repeated_per_sample_loss.ndim != 1 or repeated_per_sample_loss.numel() == 0: raise ValueError("repeated_per_sample_loss must be a non-empty one-dimensional tensor") if not isinstance(repeat_count, int) or isinstance(repeat_count, bool) or repeat_count <= 0: raise ValueError("repeat_count must be a positive integer") if repeated_per_sample_loss.numel() % repeat_count != 0: raise ValueError("loss count must be divisible by repeat_count") inferred_batch_size = repeated_per_sample_loss.numel() // repeat_count if batch_size is not None and batch_size != inferred_batch_size: raise ValueError("batch_size does not match the repeated loss shape") return repeated_per_sample_loss.reshape(repeat_count, inferred_batch_size).mean(dim=0) def ordinary_importance_corrected_mean( per_sample_loss: torch.Tensor, importance_weight: torch.Tensor, ) -> torch.Tensor: """Return the ordinary (not self-normalized) importance estimator.""" if not isinstance(per_sample_loss, torch.Tensor) or not isinstance(importance_weight, torch.Tensor): raise TypeError("per_sample_loss and importance_weight must be torch.Tensor instances") if per_sample_loss.ndim != 1 or per_sample_loss.numel() == 0: raise ValueError("per_sample_loss must be a non-empty one-dimensional tensor") if importance_weight.shape != per_sample_loss.shape: raise ValueError("importance_weight must have the same shape as per_sample_loss") detached_weight = importance_weight.detach() if not torch.isfinite(detached_weight).all() or (detached_weight < 0).any(): raise ValueError("importance_weight must be finite and non-negative") work_dtype = torch.float64 if per_sample_loss.dtype == torch.float64 else torch.float32 loss = per_sample_loss.to(dtype=work_dtype) weight = detached_weight.to(device=per_sample_loss.device, dtype=work_dtype) return (weight * loss).mean()