pdd ¶

def build_pdd_sampling_plan(
    step_indices: Sequence[int],
    schedules: Mapping[str, PDDModalitySchedule],
    *,
    device: torch.device | str | None = None,
) -> PDDSamplingPlan:
    """Partition the fine grid into fused blocks for every modality.

    *step_indices* is the checkpoint's trained partition: fine-grid nodes that
    increase strictly from 0 to the grid size, one fused block per consecutive
    pair. ``MiniMaxH3PipelineConfig.resolve_checkpoint_settings`` validates it.
    """
    max_ts = {schedule.max_t for schedule in schedules.values()}
    if len(max_ts) != 1:
        raise ValueError(f"Every PDD modality schedule must share max_t, got {sorted(max_ts)}")
    fine_grid = pdd_fine_grid(int(step_indices[-1]), max_t=next(iter(max_ts)), device=device)
    node_indices = torch.as_tensor(list(step_indices), dtype=torch.long, device=fine_grid.device)
    weights = {}
    for name, schedule in schedules.items():
        weights[name] = schedule.integration_weights(fine_grid)
        validate_finite_nonzero(weights[name], f"PDD {name} integration weight")
    return PDDSamplingPlan(
        fine_grid=fine_grid,
        indices=node_indices,
        node_sigmas={
            name: schedule.sigma(fine_grid[node_indices])
            for name, schedule in schedules.items()
        },
        integration_weights=weights,
    )

@contextlib.contextmanager
def fuse_pdd_heads(
    linears: Mapping[str, PDDReplicatedLinear],
    start: int,
    end: int,
    integration_weights: Mapping[str, torch.Tensor],
    precision_decoding: torch.dtype,
) -> Iterator[None]:
    """Fuse block ``[start, end)`` on every modality head at once."""
    with contextlib.ExitStack() as stack:
        for name, linear in linears.items():
            stack.enter_context(linear.fuse(start, end, integration_weights[name], precision_decoding))
        yield

def pdd_fine_grid(
    pdd_steps: int,
    *,
    max_t: float = PDD_GRID_MAX_T,
    device: torch.device | str | None = None,
) -> torch.Tensor:
    """The descending base-clock grid ``u_0 = max_t > ... > u_N = 0`` in float64."""
    grid = torch.linspace(max_t, 0.0, pdd_steps + 1, dtype=torch.float64, device=device)
    return grid.clamp(max=max_t)

def shifted_noise_amount(base: torch.Tensor, shift: float, *, max_t: float = PDD_GRID_MAX_T) -> torch.Tensor:
    """Rational time shift ``f_s(u) = s * u * M / (u * (s - 1) + M)`` on ``[0, M]``.

    ``M = max_t`` is a fixed point of every shift, so the first grid node has
    the same noise level in every modality.
    """
    if shift == 1.0:
        return base
    return shift * base * max_t / (base * (shift - 1.0) + max_t)

def shifted_noise_delta(
    base_start: torch.Tensor,
    base_end: torch.Tensor,
    shift: float,
    *,
    max_t: float = PDD_GRID_MAX_T,
) -> torch.Tensor:
    """Return ``f_s(base_end) - f_s(base_start)`` without subtractive cancellation.

    The rational difference simplifies to
    ``s * M^2 * (b - a) / (D(a) * D(b))`` with ``D(u) = M + (s - 1) * u``; the
    identity shift keeps the direct difference bit-exact.
    """
    if shift == 1.0:
        return base_end - base_start
    denominator_start = max_t + (shift - 1.0) * base_start
    denominator_end = max_t + (shift - 1.0) * base_end
    return shift * max_t**2 * (base_end - base_start) / (denominator_start * denominator_end)

def validate_pdd_steps(value: Any, name: str = "pdd_steps") -> int:
    """A PDD grid needs at least two intervals to have a block structure."""
    if isinstance(value, bool) or not isinstance(value, int) or value < 2:
        raise ValueError(f"{name} must be an int >= 2, got {value!r}")
    return value