Module 12. Diffusion Language Models

import hashlib
import json
import math
import subprocess
import sys
from dataclasses import asdict, dataclass
from numbers import Integral, Real
from typing import Any, Callable, Literal

try:
    import numpy as np
except ModuleNotFoundError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "numpy"])
    import numpy as np

REFERENCE_JSON = '{"block_mask":{"allowed":[[1,0,0,0,0,0,0],[1,1,0,0,0,0,0],[1,1,1,0,0,0,0],[1,1,1,1,1,1,1],[1,1,1,1,1,1,1],[1,1,1,1,1,1,1],[1,1,1,1,1,1,1]],"block_length":4,"prefix_length":3},"corruption":{"clean_tokens":[10,11,12,13,14,15,16,17],"eligible_positions":[false,false,true,true,true,true,true,true],"mask_id":-1,"masked_positions":[2,4,5,7],"noisy_tokens":[10,11,-1,13,-1,-1,16,-1],"time":0.5},"diffusiongemma_passport":{"active_parameters":3800000000,"active_routed_experts":8,"canvas_tokens":256,"committed_tokens_per_pass_max":20,"committed_tokens_per_pass_min":15,"context_tokens":262144,"h100_developer_reported_tokens_per_s_lower_bound":1100,"h100_precision":"FP8","idealized_active_4bit_bytes":1900000000.0,"idealized_active_4bit_decimal_gb":1.9,"idealized_active_4bit_gib":1.7695128917694092,"idealized_full_bank_4bit_bytes":12600000000.0,"idealized_full_bank_4bit_decimal_gb":12.6,"idealized_full_bank_4bit_gib":11.73466444015503,"layers":30,"native_nvfp4_hardware_family":"NVIDIA Blackwell","release_date":"2026-06-10","routed_experts":128,"shared_experts":1,"sliding_window_tokens":1024,"source_snapshot":"2026-10-11","total_parameters":25200000000,"traffic_note":"Neither idealized payload equals measured bytes per diffusion pass; actual traffic depends on expert union, precision, metadata, cache, kernels, and runtime."},"loss":{"correct_probabilities":[0.9,0.7,0.6,0.8],"masked_positions_count":4,"reduction":"sum_over_masked_positions_then_multiply_by_1_over_t","target_tokens":[0,1,2,3],"time":0.5,"value":2.392009269353518},"policy":{"confidence":1.0,"random":0.6854166666666667,"random_with_remask":0.8291666666666666,"target":[0,1,2,3,4,5,1,2,3,4,5,0],"trials":200},"roofline":{"ar_compute_time_fraction_of_step":0.003383838383838384,"autoregressive":{"bottleneck":"memory","compute_time_s":1.616161616161616e-05,"ideal_arithmetic_share_of_step":0.003383838383838384,"memory_time_s":0.004776119402985075,"step_time_s":0.004776119402985075,"tokens_per_s":209.375,"weight_bits":16.0,"weight_payload_bytes":16000000000.0,"weight_payload_decimal_gb":16.0,"weight_payload_gib":14.901161193847656},"batch1_throughput_ratio_at_64_steps":4.617537313432836,"batch_parity_under_linear_ar_scaling":4.617537313432836,"block_4bit":{"bottleneck":"compute","compute_time_s":0.004137373737373737,"memory_time_s":0.0011940298507462687,"step_time_s":0.004137373737373737,"tokens_per_s":3867.1875,"weight_bits":4.0,"weight_payload_bytes":4000000000.0,"weight_payload_decimal_gb":4.0,"weight_payload_gib":3.725290298461914},"block_bf16":{"bottleneck":"memory","compute_time_s":0.004137373737373737,"memory_time_s":0.004776119402985075,"step_time_s":0.004776119402985075,"tokens_per_s":3350.0,"weight_bits":16.0,"weight_payload_bytes":16000000000.0,"weight_payload_decimal_gb":16.0,"weight_payload_gib":14.901161193847656},"block_steps_for_1000_tokens_per_s_4bit":61.875,"diffusion":{"128":{"bottleneck":"compute","compute_time_s":0.01654949494949495,"memory_time_s":0.004776119402985075,"step_time_s":0.01654949494949495,"tokens_per_s":483.3984375,"weight_bits":16.0,"weight_payload_bytes":16000000000.0,"weight_payload_decimal_gb":16.0,"weight_payload_gib":14.901161193847656},"16":{"bottleneck":"compute","compute_time_s":0.01654949494949495,"memory_time_s":0.004776119402985075,"step_time_s":0.01654949494949495,"tokens_per_s":3867.1875,"weight_bits":16.0,"weight_payload_bytes":16000000000.0,"weight_payload_decimal_gb":16.0,"weight_payload_gib":14.901161193847656},"64":{"bottleneck":"compute","compute_time_s":0.01654949494949495,"memory_time_s":0.004776119402985075,"step_time_s":0.01654949494949495,"tokens_per_s":966.796875,"weight_bits":16.0,"weight_payload_bytes":16000000000.0,"weight_payload_decimal_gb":16.0,"weight_payload_gib":14.901161193847656}},"hardware":{"memory_bandwidth_bytes_per_s":3350000000000.0,"peak_flops_per_s":990000000000000.0},"maximum_steps_strictly_faster":295,"model":{"flops_per_parameter_per_token":2.0,"parameter_count":8000000000},"scenario":"generic_dense_8b_h100_lower_bound","scope":"Arithmetic lower-bound scenario only: excludes KV traffic, activations, communication, routing, kernel/runtime overhead, batching, and measured utilization."},"schedule":{"counts":[8,6,4,2,0],"length":8,"steps":4,"weights":{"0.25":4.0,"0.5":2.0,"0.75":1.3333333333333333}},"trace":{"candidates":[3,1,4,0,3,2,4,3],"confidence":[0.447488947046876,0.595067902012219,0.6111163345147275,0.40178371163146176,0.2776006503454896,0.5910785484224659,0.9279309181635038,0.46568452224635726],"correct":6,"final":[3,1,4,0,3,2,4,3],"model_calls":4,"opened_order":[[2,6],[1,5],[0,7],[3,4]],"target":[3,1,4,1,0,2,4,3]}}'
REFERENCE_SHA256 = "aba7d3b44833a1497cd34c8509147ad60c530897cfa280b075721c9a2e36c54d"
assert hashlib.sha256(REFERENCE_JSON.encode("utf-8")).hexdigest() == REFERENCE_SHA256
REFERENCE = json.loads(REFERENCE_JSON)
np.set_printoptions(precision=6, suppress=True)
print('Embedded Module 12 contract verified', REFERENCE_SHA256)

def _positive_int(name: str, value: object) -> int:
    """Return a strictly positive Python integer and reject booleans explicitly."""

    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    result = int(value)
    if result <= 0:
        raise ValueError(f"{name} must be positive")
    return result


def _nonnegative_int(name: str, value: object) -> int:
    """Return a non-negative Python integer and reject booleans explicitly."""

    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    result = int(value)
    if result < 0:
        raise ValueError(f"{name} must be non-negative")
    return result


def _finite_real(name: str, value: object) -> float:
    """Return a finite real scalar and reject booleans and complex values."""

    if isinstance(value, bool) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real scalar")
    result = float(value)
    if not math.isfinite(result):
        raise ValueError(f"{name} must be finite")
    return result


def _bounded_real(
    name: str,
    value: object,
    *,
    lower: float,
    upper: float,
    lower_inclusive: bool = True,
    upper_inclusive: bool = True,
) -> float:
    """Validate one finite scalar against an explicit interval."""

    result = _finite_real(name, value)
    lower_ok = result >= lower if lower_inclusive else result > lower
    upper_ok = result <= upper if upper_inclusive else result < upper
    if not (lower_ok and upper_ok):
        left = "[" if lower_inclusive else "("
        right = "]" if upper_inclusive else ")"
        raise ValueError(f"{name} must lie in {left}{lower}, {upper}{right}")
    return result


def _finite_array(
    name: str,
    value: object,
    *,
    ndim: int | None = None,
    dtype: np.dtype[Any] | type = np.float64,
) -> np.ndarray:
    """Convert to a NumPy array and enforce shape rank and finite values."""

    array = np.asarray(value, dtype=dtype)
    if ndim is not None and array.ndim != ndim:
        raise ValueError(f"{name} must have rank {ndim}, got shape {array.shape}")
    if array.size == 0:
        raise ValueError(f"{name} must be non-empty")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    return array


def _integer_tokens(name: str, value: object, *, ndim: int = 1) -> np.ndarray:
    """Validate an integer token tensor without silently truncating floats."""

    raw = np.asarray(value)
    if raw.ndim != ndim:
        raise ValueError(f"{name} must have rank {ndim}, got shape {raw.shape}")
    if raw.size == 0:
        raise ValueError(f"{name} must be non-empty")
    if raw.dtype.kind not in "iu":
        raise TypeError(f"{name} must contain integer token IDs")
    return raw.astype(np.int64, copy=False)

@dataclass(frozen=True)
class CorruptionResult:
    """One absorbing-mask forward sample for a one-dimensional token sequence."""

    clean_tokens: np.ndarray
    noisy_tokens: np.ndarray
    masked_positions: np.ndarray
    eligible_positions: np.ndarray
    time: float
    mask_id: int


def linear_mask_schedule(length: int, steps: int) -> tuple[int, ...]:
    """Return monotone target mask counts from ``length`` to zero.

    ``ceil`` gives a declared policy for non-divisible lengths. Duplicate counts are allowed
    when ``steps > length``; such a step performs a model evaluation but opens no token.
    """

    n = _positive_int("length", length)
    s = _positive_int("steps", steps)
    counts = tuple((n * remaining + s - 1) // s for remaining in range(s, -1, -1))
    if counts[0] != n or counts[-1] != 0:
        raise AssertionError("schedule endpoints are inconsistent")
    if any(left < right for left, right in zip(counts, counts[1:])):
        raise AssertionError("mask schedule must be non-increasing")
    return counts


def elbo_time_weight(time: Real) -> float:
    """Return the absorbing-diffusion slice weight ``1 / t`` for ``0 < t <= 1``."""

    t = _bounded_real("time", time, lower=0.0, upper=1.0, lower_inclusive=False)
    return 1.0 / t


def corrupt_absorbing(
    tokens: object,
    *,
    time: Real,
    mask_id: int,
    rng: np.random.Generator,
    eligible_positions: object | None = None,
) -> CorruptionResult:
    """Replace eligible tokens independently by ``mask_id`` with probability ``time``.

    The mask token must be outside the clean sequence so that masked and observed positions
    remain unambiguous. ``eligible_positions`` can protect a fixed prefix or special tokens.
    """

    clean = _integer_tokens("tokens", tokens)
    if isinstance(mask_id, bool) or not isinstance(mask_id, Integral):
        raise TypeError("mask_id must be an integer")
    mask = int(mask_id)
    if np.any(clean == mask):
        raise ValueError("mask_id must not occur in the clean sequence")
    if not isinstance(rng, np.random.Generator):
        raise TypeError("rng must be numpy.random.Generator")
    t = _bounded_real("time", time, lower=0.0, upper=1.0)

    if eligible_positions is None:
        eligible = np.ones(clean.shape, dtype=bool)
    else:
        eligible = np.asarray(eligible_positions)
        if eligible.shape != clean.shape or eligible.dtype.kind != "b":
            raise ValueError("eligible_positions must be a boolean vector matching tokens")
        eligible = eligible.astype(bool, copy=False)

    selected = eligible & (rng.random(clean.shape[0]) < t)
    noisy = clean.copy()
    noisy[selected] = mask
    return CorruptionResult(
        clean_tokens=clean.copy(),
        noisy_tokens=noisy,
        masked_positions=selected,
        eligible_positions=eligible.copy(),
        time=t,
        mask_id=mask,
    )


def stable_log_softmax(logits: object) -> np.ndarray:
    """Compute row-wise log-softmax for a finite ``[positions, vocabulary]`` tensor."""

    values = _finite_array("logits", logits, ndim=2)
    maxima = np.max(values, axis=-1, keepdims=True)
    shifted = values - maxima
    log_normalizer = np.log(np.exp(shifted).sum(axis=-1, keepdims=True))
    result = shifted - log_normalizer
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("log-softmax produced a non-finite value")
    return result


def masked_diffusion_loss(
    logits: object,
    target_tokens: object,
    masked_positions: object,
    *,
    time: Real,
) -> float:
    """Return ``(1/t) * sum(-log p_theta(x_i | x_t))`` on masked positions only."""

    log_probs = stable_log_softmax(logits)
    targets = _integer_tokens("target_tokens", target_tokens)
    if targets.shape[0] != log_probs.shape[0]:
        raise ValueError("target_tokens and logits must have the same position axis")
    if np.any(targets < 0) or np.any(targets >= log_probs.shape[1]):
        raise ValueError("target token ID is outside the vocabulary")
    selected = np.asarray(masked_positions)
    if selected.shape != targets.shape or selected.dtype.kind != "b":
        raise ValueError("masked_positions must be a boolean vector matching target_tokens")
    selected = selected.astype(bool, copy=False)
    weight = elbo_time_weight(time)
    nll = -log_probs[np.arange(targets.size), targets]
    loss = weight * float(nll[selected].sum())
    if not math.isfinite(loss):
        raise FloatingPointError("diffusion loss is non-finite")
    return loss


ref = REFERENCE["schedule"]
assert linear_mask_schedule(ref["length"], ref["steps"]) == tuple(ref["counts"])
assert linear_mask_schedule(3, 5) == (3, 3, 2, 2, 1, 0)
for time_text, expected in ref["weights"].items():
    assert math.isclose(elbo_time_weight(float(time_text)), expected, rel_tol=1e-15)

cor_ref = REFERENCE["corruption"]
corrupted = corrupt_absorbing(
    np.array(cor_ref["clean_tokens"], dtype=np.int64),
    time=cor_ref["time"],
    mask_id=cor_ref["mask_id"],
    rng=np.random.default_rng(73),
    eligible_positions=np.array(cor_ref["eligible_positions"], dtype=bool),
)
assert np.flatnonzero(corrupted.masked_positions).tolist() == cor_ref["masked_positions"]
assert corrupted.noisy_tokens.tolist() == cor_ref["noisy_tokens"]
assert not corrupted.masked_positions[:2].any()  # protected prefix

# Empirical mask rate over eligible positions; this checks the Bernoulli contract, not a fixture.
rng = np.random.default_rng(700)
empirical = float((rng.random((4000, 6)) < 0.5).mean())
assert abs(empirical - 0.5) < 0.02
print("schedule", ref["counts"], "| noisy", corrupted.noisy_tokens.tolist(), "| empirical", round(empirical, 4))
print('Schedule, protected prefix, and mask rate verified')

ref = REFERENCE["loss"]
targets = np.array(ref["target_tokens"], dtype=np.int64)
p_correct = np.array(ref["correct_probabilities"], dtype=np.float64)
vocabulary = 5
probabilities = np.empty((targets.size, vocabulary), dtype=np.float64)
for row, (token, p_correct_value) in enumerate(zip(targets, p_correct)):
    probabilities[row] = (1.0 - p_correct_value) / (vocabulary - 1)
    probabilities[row, token] = p_correct_value
logits = np.log(probabilities)
selected = np.ones(targets.shape, dtype=bool)
loss = masked_diffusion_loss(logits, targets, selected, time=ref["time"])
assert math.isclose(loss, ref["value"], rel_tol=1e-14)

# A non-selected row must not affect the loss, even with extreme finite logits.
extra_logits = np.vstack([logits, np.array([1000.0, -1000.0, -1000.0, -1000.0, -1000.0])])
extra_targets = np.append(targets, 4)
extra_mask = np.append(selected, False)
assert math.isclose(
    masked_diffusion_loss(extra_logits, extra_targets, extra_mask, time=ref["time"]),
    loss,
    rel_tol=1e-14,
)
assert np.allclose(np.exp(stable_log_softmax(extra_logits)).sum(axis=1), 1.0)
print(f"L_t={loss:.9f} for t={ref['time']} and four masked positions")
print('Masked loss and numerical stability verified')

@dataclass(frozen=True)
class DenoiseStep:
    """Inspectable state transition produced by one denoising model call."""

    step: int
    target_masked: int
    remasked_positions: tuple[int, ...]
    opened_positions: tuple[int, ...]
    masked_after: int
    sequence: tuple[int, ...]


@dataclass(frozen=True)
class DenoiseResult:
    """Final sequence and full trace of a masked-diffusion sampling run."""

    final_tokens: np.ndarray
    schedule: tuple[int, ...]
    steps: tuple[DenoiseStep, ...]
    model_calls: int


def stable_softmax(logits: object) -> np.ndarray:
    """Compute row-wise softmax for finite logits and verify probability normalization."""

    log_probs = stable_log_softmax(logits)
    probabilities = np.exp(log_probs)
    if not np.allclose(probabilities.sum(axis=-1), 1.0, atol=1e-12, rtol=0.0):
        raise FloatingPointError("softmax rows do not sum to one")
    return probabilities


def sample_masked_diffusion(
    predict: Callable[[np.ndarray, int], object],
    *,
    length: int,
    steps: int,
    mask_id: int,
    rng: np.random.Generator,
    selection_policy: Literal["confidence", "random", "entropy"] = "confidence",
    token_policy: Literal["greedy", "categorical"] = "greedy",
    remask_fraction: Real = 0.0,
) -> DenoiseResult:
    """Run a transparent masked-diffusion sampler with one model call per step.

    ``predict`` receives a copy of the current token vector and the one-based step number,
    and must return finite logits of shape ``[length, vocabulary]``. The mask ID must lie
    outside that vocabulary. Opened tokens remain fixed unless explicitly selected for
    remasking. Confidence ties are resolved by lower position index.
    """

    n = _positive_int("length", length)
    s = _positive_int("steps", steps)
    if isinstance(mask_id, bool) or not isinstance(mask_id, Integral):
        raise TypeError("mask_id must be an integer")
    mask = int(mask_id)
    if not isinstance(rng, np.random.Generator):
        raise TypeError("rng must be numpy.random.Generator")
    if selection_policy not in {"confidence", "random", "entropy"}:
        raise ValueError("selection_policy must be 'confidence', 'random', or 'entropy'")
    if token_policy not in {"greedy", "categorical"}:
        raise ValueError("token_policy must be 'greedy' or 'categorical'")
    remask = _bounded_real("remask_fraction", remask_fraction, lower=0.0, upper=1.0)

    schedule = linear_mask_schedule(n, s)
    sequence = np.full(n, mask, dtype=np.int64)
    trace: list[DenoiseStep] = []
    stored_confidence = np.zeros(n)
    expected_vocabulary = None

    for step_number, target_masked in enumerate(schedule[1:], start=1):
        # Select from scores saved when tokens were last predicted. Mask BEFORE
        # this call, so the network conditions on the actual corrupted input.
        opened_before = np.flatnonzero(sequence != mask)
        remasked = np.empty(0, dtype=np.int64)
        remask_count = math.floor(remask * opened_before.size)
        if remask_count:
            order = np.lexsort((opened_before, stored_confidence[opened_before]))
            remasked = opened_before[order[:remask_count]]
            sequence[remasked] = mask
        raw_logits = predict(sequence.copy(), step_number)
        logits = _finite_array("predict output", raw_logits, ndim=2)
        if logits.shape[0] != n:
            raise ValueError("predict output must have one row per sequence position")
        vocabulary = logits.shape[1]
        if vocabulary <= 0:
            raise ValueError("predict output vocabulary must be non-empty")
        if 0 <= mask < vocabulary:
            raise ValueError("mask_id must lie outside the model vocabulary")
        if expected_vocabulary is None:
            expected_vocabulary = vocabulary
        elif vocabulary != expected_vocabulary:
            raise ValueError("predict changed vocabulary size between calls")
        probabilities = stable_softmax(logits)

        if token_policy == "greedy":
            candidates = np.argmax(probabilities, axis=-1).astype(np.int64)
        else:
            candidates = np.array(
                [rng.choice(vocabulary, p=row) for row in probabilities],
                dtype=np.int64,
            )
        candidate_confidence = probabilities[np.arange(n), candidates]

        masked_now = np.flatnonzero(sequence == mask)
        open_count = masked_now.size - target_masked
        if open_count < 0:
            raise AssertionError("sampler has fewer masks than the schedule requires")
        if selection_policy == "confidence":
            order = np.lexsort((masked_now, -candidate_confidence[masked_now]))
            opened = masked_now[order[:open_count]]
        elif selection_policy == "entropy":
            log_probabilities = stable_log_softmax(logits)
            entropy = -(probabilities * log_probabilities).sum(axis=-1)
            order = np.lexsort((masked_now, entropy[masked_now]))
            opened = masked_now[order[:open_count]]
        else:
            opened = rng.permutation(masked_now)[:open_count]
        sequence[opened] = candidates[opened]
        stored_confidence[opened] = candidate_confidence[opened]

        masked_after = int(np.count_nonzero(sequence == mask))
        if masked_after != target_masked:
            raise AssertionError("sampler did not reach the scheduled mask count")
        trace.append(
            DenoiseStep(
                step=step_number,
                target_masked=target_masked,
                remasked_positions=tuple(sorted(map(int, remasked))),
                opened_positions=tuple(sorted(map(int, opened))),
                masked_after=masked_after,
                sequence=tuple(map(int, sequence)),
            )
        )

    if np.any(sequence == mask):
        raise AssertionError("final sequence still contains mask tokens")
    return DenoiseResult(
        final_tokens=sequence.copy(),
        schedule=schedule,
        steps=tuple(trace),
        model_calls=s,
    )


def make_contextual_toy_predictor(
    target_tokens: object,
    *,
    vocabulary_size: int,
) -> Callable[[np.ndarray, int], np.ndarray]:
    """Create a deterministic toy model whose odd positions need their left anchor.

    Even positions are high-confidence anchors. An odd position initially prefers a wrong
    token; once its left anchor is opened correctly, the target becomes high-confidence.
    This makes reveal order and remasking observable without training a neural network.
    """

    target = _integer_tokens("target_tokens", target_tokens)
    vocab = _positive_int("vocabulary_size", vocabulary_size)
    if np.any(target < 0) or np.any(target >= vocab):
        raise ValueError("target token ID is outside vocabulary_size")

    def predict(sequence: np.ndarray, step: int) -> np.ndarray:
        """Return deterministic contextual logits for one sampler step."""

        current = _integer_tokens("sequence", sequence)
        if current.shape != target.shape:
            raise ValueError("sequence shape differs from the toy target")
        _positive_int("step", step)
        logits = np.full((target.size, vocab), -2.5, dtype=np.float64)
        for position, target_id in enumerate(target):
            wrong_id = (int(target_id) + 1) % vocab
            if position % 2 == 0:
                logits[position, int(target_id)] = 3.5
                logits[position, wrong_id] = 0.0
            elif current[position - 1] == target[position - 1]:
                logits[position, int(target_id)] = 3.3
                logits[position, wrong_id] = 0.0
            else:
                logits[position, wrong_id] = 1.5
                logits[position, int(target_id)] = 1.1
        return logits

    return predict


ref = REFERENCE["trace"]
target = np.array(ref["target"], dtype=np.int64)
trace_rng = np.random.default_rng(7)
fixed_logits = trace_rng.standard_normal((target.size, 5))
fixed_logits[np.arange(target.size), target] += 2.0
model_calls = 0

def fixed_predict(sequence, step):
    """Return the frozen logits and count exactly one call per denoising step."""

    global model_calls
    assert sequence.shape == target.shape
    assert 1 <= step <= 4
    model_calls += 1
    return fixed_logits

result = sample_masked_diffusion(
    fixed_predict,
    length=target.size,
    steps=4,
    mask_id=-1,
    rng=np.random.default_rng(99),
    selection_policy="confidence",
)
assert model_calls == result.model_calls == ref["model_calls"] == 4
assert [list(step.opened_positions) for step in result.steps] == ref["opened_order"]
assert result.final_tokens.tolist() == ref["final"]
assert int(np.count_nonzero(result.final_tokens == target)) == ref["correct"] == 6
assert [step.masked_after for step in result.steps] == list(result.schedule[1:])

# Equal confidence is resolved by lower position index.
def tie_predict(sequence, step):
    """Return equal logits to exercise the deterministic position tie-break."""

    return np.zeros((sequence.size, 3), dtype=np.float64)

tie_result = sample_masked_diffusion(
    tie_predict,
    length=4,
    steps=2,
    mask_id=-1,
    rng=np.random.default_rng(1),
)
assert tie_result.steps[0].opened_positions == (0, 1)
print("opened", ref["opened_order"], "| final", result.final_tokens.tolist(), "| correct", ref["correct"])
print('Demasking trace and one model call per step verified')

ref = REFERENCE["policy"]
target = np.array(ref["target"], dtype=np.int64)
predictor = make_contextual_toy_predictor(target, vocabulary_size=6)
settings = {
    "confidence": ("confidence", 0.0),
    "random": ("random", 0.0),
    "random_with_remask": ("random", 0.25),
}
measured = {}
for name, (selection, remask) in settings.items():
    accuracies = []
    for trial in range(ref["trials"]):
        sampled = sample_masked_diffusion(
            predictor,
            length=target.size,
            steps=4,
            mask_id=-1,
            rng=np.random.default_rng(1000 + trial),
            selection_policy=selection,
            remask_fraction=remask,
        )
        accuracies.append(float(np.mean(sampled.final_tokens == target)))
    measured[name] = float(np.mean(accuracies))
    assert math.isclose(measured[name], ref[name], rel_tol=0.0, abs_tol=1e-15)

assert measured["confidence"] > measured["random"]
assert measured["random_with_remask"] > measured["random"]
print({name: round(value, 4) for name, value in measured.items()})
print('Reveal order and remasking verified')

def block_diffusion_attention_mask(prefix_length: int, block_length: int) -> np.ndarray:
    """Return a full boolean mask for a fixed causal prefix plus one bidirectional block.

    Prefix queries remain causal and cannot read the current block. Queries in the current
    block may read every prefix token and every token of their own block. The fixed prefix is
    therefore cacheable between denoising steps, while the current block is recomputed.
    """

    prefix = _nonnegative_int("prefix_length", prefix_length)
    block = _positive_int("block_length", block_length)
    total = prefix + block
    allowed = np.zeros((total, total), dtype=bool)
    if prefix:
        allowed[:prefix, :prefix] = np.tril(np.ones((prefix, prefix), dtype=bool))
    allowed[prefix:, :total] = True
    return allowed


ref = REFERENCE["block_mask"]
allowed = block_diffusion_attention_mask(ref["prefix_length"], ref["block_length"])
assert allowed.astype(int).tolist() == ref["allowed"]
prefix = ref["prefix_length"]
assert np.array_equal(allowed[:prefix, :prefix], np.tril(np.ones((prefix, prefix), dtype=bool)))
assert not allowed[:prefix, prefix:].any()
assert allowed[prefix:, :].all()
print(allowed.astype(int))
print('Block-diffusion attention geometry verified')

@dataclass(frozen=True)
class HardwareSpec:
    """Decimal-bandwidth and FLOP/s inputs for a transparent roofline scenario."""

    memory_bandwidth_bytes_per_s: float
    peak_flops_per_s: float

    def __post_init__(self) -> None:
        """Validate positive finite hardware rates."""

        object.__setattr__(
            self,
            "memory_bandwidth_bytes_per_s",
            _bounded_real(
                "memory_bandwidth_bytes_per_s",
                self.memory_bandwidth_bytes_per_s,
                lower=0.0,
                upper=float("inf"),
                lower_inclusive=False,
            ),
        )
        object.__setattr__(
            self,
            "peak_flops_per_s",
            _bounded_real(
                "peak_flops_per_s",
                self.peak_flops_per_s,
                lower=0.0,
                upper=float("inf"),
                lower_inclusive=False,
            ),
        )


@dataclass(frozen=True)
class DenseModelSpec:
    """Dense model size and approximate arithmetic per token."""

    parameter_count: int
    flops_per_parameter_per_token: float = 2.0

    def __post_init__(self) -> None:
        """Validate model size and per-token arithmetic coefficient."""

        object.__setattr__(self, "parameter_count", _positive_int("parameter_count", self.parameter_count))
        object.__setattr__(
            self,
            "flops_per_parameter_per_token",
            _bounded_real(
                "flops_per_parameter_per_token",
                self.flops_per_parameter_per_token,
                lower=0.0,
                upper=float("inf"),
                lower_inclusive=False,
            ),
        )


@dataclass(frozen=True)
class RooflineEstimate:
    """Arithmetic floor for one repeated full-sequence or block denoising workload."""

    weight_bits: float
    weight_payload_bytes: float
    weight_payload_decimal_gb: float
    weight_payload_gib: float
    memory_time_s: float
    compute_time_s: float
    step_time_s: float
    tokens_per_s: float
    bottleneck: str


def bytes_to_gib(byte_count: Real) -> float:
    """Convert a non-negative byte count to binary GiB."""

    value = _bounded_real("byte_count", byte_count, lower=0.0, upper=float("inf"))
    return value / 2**30


def bytes_to_decimal_gb(byte_count: Real) -> float:
    """Convert a non-negative byte count to decimal GB."""

    value = _bounded_real("byte_count", byte_count, lower=0.0, upper=float("inf"))
    return value / 1e9


def estimate_roofline(
    model: DenseModelSpec,
    hardware: HardwareSpec,
    *,
    positions_per_step: int,
    denoise_steps: int,
    weight_bits: Real,
) -> RooflineEstimate:
    """Estimate a weight-stream/compute floor for repeated full-width model evaluations.

    The estimate reads the weight payload once per model call and counts dense arithmetic as
    ``flops_per_parameter_per_token * P * positions_per_step``. It excludes KV traffic,
    activations, communication, kernels, runtime overhead, batching, and measured utilization.
    """

    positions = _positive_int("positions_per_step", positions_per_step)
    steps = _positive_int("denoise_steps", denoise_steps)
    bits = _bounded_real("weight_bits", weight_bits, lower=0.0, upper=float("inf"), lower_inclusive=False)
    payload = model.parameter_count * bits / 8.0
    memory_time = payload / hardware.memory_bandwidth_bytes_per_s
    compute_time = (
        model.flops_per_parameter_per_token * model.parameter_count * positions
        / hardware.peak_flops_per_s
    )
    step_time = max(memory_time, compute_time)
    throughput = positions / (steps * step_time)
    bottleneck = "memory" if memory_time >= compute_time else "compute"
    return RooflineEstimate(
        weight_bits=bits,
        weight_payload_bytes=payload,
        weight_payload_decimal_gb=bytes_to_decimal_gb(payload),
        weight_payload_gib=bytes_to_gib(payload),
        memory_time_s=memory_time,
        compute_time_s=compute_time,
        step_time_s=step_time,
        tokens_per_s=throughput,
        bottleneck=bottleneck,
    )


def maximum_steps_strictly_faster(
    *,
    positions_per_step: int,
    diffusion_step_time_s: Real,
    autoregressive_step_time_s: Real,
) -> int:
    """Return the largest integer ``S`` with ``N/(S*t_diff) > 1/t_ar``."""

    positions = _positive_int("positions_per_step", positions_per_step)
    t_diff = _bounded_real(
        "diffusion_step_time_s",
        diffusion_step_time_s,
        lower=0.0,
        upper=float("inf"),
        lower_inclusive=False,
    )
    t_ar = _bounded_real(
        "autoregressive_step_time_s",
        autoregressive_step_time_s,
        lower=0.0,
        upper=float("inf"),
        lower_inclusive=False,
    )
    ratio = positions * t_ar / t_diff
    return max(0, math.ceil(ratio) - 1)


ref = REFERENCE["roofline"]
hardware = HardwareSpec(**ref["hardware"])
model = DenseModelSpec(**ref["model"])
ar = estimate_roofline(model, hardware, positions_per_step=1, denoise_steps=1, weight_bits=16)
for key, value in asdict(ar).items():
    assert ref["autoregressive"][key] == value
assert math.isclose(
    ar.compute_time_s / ar.step_time_s,
    ref["autoregressive"]["ideal_arithmetic_share_of_step"],
    rel_tol=0.0,
    abs_tol=1e-15,
)
assert ref["scenario"] == "generic_dense_8b_h100_lower_bound"

for steps in (16, 64, 128):
    estimate = estimate_roofline(
        model,
        hardware,
        positions_per_step=1024,
        denoise_steps=steps,
        weight_bits=16,
    )
    assert asdict(estimate) == ref["diffusion"][str(steps)]
    print(f"S={steps:3d} -> {estimate.tokens_per_s:8.3f} tok/s, {estimate.bottleneck}-bound")

diff64 = RooflineEstimate(**ref["diffusion"]["64"])
parity = maximum_steps_strictly_faster(
    positions_per_step=1024,
    diffusion_step_time_s=diff64.step_time_s,
    autoregressive_step_time_s=ar.step_time_s,
)
assert parity == ref["maximum_steps_strictly_faster"] == 295
assert 1024 / (295 * diff64.step_time_s) > ar.tokens_per_s
assert 1024 / (296 * diff64.step_time_s) <= ar.tokens_per_s

block_bf16 = estimate_roofline(model, hardware, positions_per_step=256, denoise_steps=16, weight_bits=16)
block_4bit = estimate_roofline(model, hardware, positions_per_step=256, denoise_steps=16, weight_bits=4)
assert asdict(block_bf16) == ref["block_bf16"]
assert asdict(block_4bit) == ref["block_4bit"]
assert block_bf16.bottleneck == "memory" and block_4bit.bottleneck == "compute"
print(
    f"bf16 payload {ar.weight_payload_decimal_gb:.3f} GB = {ar.weight_payload_gib:.3f} GiB; "
    f"4-bit payload {block_4bit.weight_payload_decimal_gb:.3f} GB = {block_4bit.weight_payload_gib:.3f} GiB"
)
print(f"AR ceiling {ar.tokens_per_s:.3f} tok/s | strict parity S={parity} | ratio@64={ref['batch1_throughput_ratio_at_64_steps']:.3f}x")
print('Generic dense-8B roofline and units verified')

ref = REFERENCE["diffusiongemma_passport"]
assert ref["total_parameters"] == 25_200_000_000
assert ref["active_parameters"] == 3_800_000_000
assert ref["routed_experts"] == 128 and ref["active_routed_experts"] == 8
assert ref["shared_experts"] == 1 and ref["canvas_tokens"] == 256
assert ref["committed_tokens_per_pass_min"] == 15
assert ref["committed_tokens_per_pass_max"] == 20
assert ref["h100_precision"] == "FP8"
assert ref["native_nvfp4_hardware_family"] == "NVIDIA Blackwell"
assert math.isclose(ref["idealized_full_bank_4bit_decimal_gb"], 12.6, rel_tol=0.0, abs_tol=1e-12)
assert math.isclose(ref["idealized_active_4bit_decimal_gb"], 1.9, rel_tol=0.0, abs_tol=1e-12)
print({
    "total_4bit_GB": ref["idealized_full_bank_4bit_decimal_gb"],
    "active_4bit_GB": ref["idealized_active_4bit_decimal_gb"],
    "committed_tokens_per_pass": [ref["committed_tokens_per_pass_min"], ref["committed_tokens_per_pass_max"]],
})
print('DiffusionGemma passport and payload bounds verified')

def expect_failure(exception_type, function, *args, **kwargs):
    """Assert that one invalid educational input fails with the declared exception."""

    try:
        function(*args, **kwargs)
    except exception_type:
        return
    raise AssertionError(f"expected {exception_type.__name__} from {function.__name__}")

expect_failure(ValueError, linear_mask_schedule, 0, 4)
expect_failure(ValueError, elbo_time_weight, 0.0)
expect_failure(ValueError, corrupt_absorbing, [1, 2, -1], time=0.5, mask_id=-1, rng=np.random.default_rng(1))
expect_failure(ValueError, corrupt_absorbing, [1, 2], time=0.5, mask_id=-1, rng=np.random.default_rng(1), eligible_positions=[1, 0])
expect_failure(ValueError, stable_log_softmax, [[0.0, np.nan]])
assert masked_diffusion_loss([[0.0, 1.0]], [1], [False], time=0.5) == 0.0
expect_failure(ValueError, masked_diffusion_loss, [[0.0, 1.0]], [2], [True], time=0.5)
expect_failure(ValueError, sample_masked_diffusion, tie_predict, length=4, steps=2, mask_id=-1, rng=np.random.default_rng(1), selection_policy="unknown")
expect_failure(ValueError, sample_masked_diffusion, tie_predict, length=4, steps=2, mask_id=1, rng=np.random.default_rng(1))
expect_failure(ValueError, sample_masked_diffusion, lambda x, s: np.zeros((3, 2)), length=4, steps=2, mask_id=-1, rng=np.random.default_rng(1))
expect_failure(ValueError, sample_masked_diffusion, lambda x, s: np.full((4, 2), np.nan), length=4, steps=2, mask_id=-1, rng=np.random.default_rng(1))
expect_failure(ValueError, block_diffusion_attention_mask, -1, 4)
expect_failure(ValueError, block_diffusion_attention_mask, 3, 0)
expect_failure(ValueError, HardwareSpec, 0.0, 990e12)
expect_failure(ValueError, DenseModelSpec, 0)
expect_failure(ValueError, estimate_roofline, model, hardware, positions_per_step=0, denoise_steps=4, weight_bits=16)
expect_failure(ValueError, estimate_roofline, model, hardware, positions_per_step=8, denoise_steps=4, weight_bits=0)
print('Negative tests passed')

# A coupled path stays masked; independent snapshots need not.
thresholds = np.array([.1, .35, .6, .95])
paths = [thresholds <= t for t in (.0, .25, .5, .75, 1.)]
assert all(np.all(~a | b) for a,b in zip(paths, paths[1:]))

# Enumerate all 16 mask subsets. With context-independent losses, the
# Bernoulli expectation of sum(masked NLL)/t equals the full NLL sum.
import itertools
row_losses = -np.log(np.array([.9,.7,.6,.8]))
for t in (.1,.5,.9):
    expectation = 0.0
    for flags in itertools.product((False,True), repeat=4):
        k = sum(flags)
        expectation += t**k * (1-t)**(4-k) * row_losses[list(flags)].sum()/t
    np.testing.assert_allclose(expectation, row_losses.sum(), rtol=0, atol=1e-12)

# Marginals do not recover a dependent joint law.
joint = np.array([[.5,0],[0,.5]])
parallel = np.outer(joint.sum(axis=1), joint.sum(axis=0))
assert np.isclose(parallel[0,1]+parallel[1,0], .5)
assert joint[0,1]+joint[1,0] == 0

# The model really sees remasked positions before predicting them.
seen = []
def observe_input(x, step):
    seen.append(x.copy())
    values = np.tile(np.array([1.,0.,-1.]), (len(x),1))
    return values
r = sample_masked_diffusion(observe_input, length=8, steps=4, mask_id=-1,
                           rng=np.random.default_rng(5), remask_fraction=.5)
assert len(seen) == r.model_calls == 4
assert any(step.remasked_positions for step in r.steps)
for x, step in zip(seen, r.steps):
    assert all(x[i] == -1 for i in step.remasked_positions)
# Reject a vocabulary change even if the previous IDs still fit.
expect_failure(ValueError, sample_masked_diffusion,
               lambda x,s: np.zeros((4, 3 if s == 1 else 4)),
               length=4, steps=2, mask_id=-1, rng=np.random.default_rng(1))
print("Coupled masking, zero-mask contribution, joint dependence, remask timing: PASS")

MINIMAL = {}
exec('"""M12: summed masked loss and deterministic reveal schedule."""\nimport numpy as np\n\n\ndef linear_mask_schedule(length: int, steps: int) -> tuple[int, ...]:\n    """Target counts from fully masked to fully revealed."""\n    if any(isinstance(v, (bool, np.bool_)) or\n           not isinstance(v, (int, np.integer)) or v <= 0\n           for v in (length, steps)):\n        raise ValueError("length and steps must be positive")\n    return tuple(\n        (int(length) * remaining + int(steps) - 1) // int(steps)\n        for remaining in range(steps, -1, -1)\n    )\n\n\ndef dllm_loss(logits, target_ids, loss_mask, time: float) -> float:\n    """(1/t) times the summed NLL on explicitly masked positions."""\n    logits = np.asarray(logits, dtype=np.float64)\n    target_ids = np.asarray(target_ids)\n    loss_mask = np.asarray(loss_mask)\n    if logits.ndim != 2 or target_ids.shape != (logits.shape[0],):\n        raise ValueError("expected logits [T,V] and target_ids [T]")\n    if logits.shape[1] == 0 or target_ids.dtype.kind not in "iu":\n        raise ValueError("target_ids must be integers and vocabulary non-empty")\n    if np.any(target_ids < 0) or np.any(target_ids >= logits.shape[1]):\n        raise ValueError("target IDs must be in [0, V)")\n    if loss_mask.shape != target_ids.shape or loss_mask.dtype.kind != "b":\n        raise ValueError("loss_mask must be boolean [T]")\n    if isinstance(time, (bool, np.bool_)) or not (0.0 < time <= 1.0):\n        raise ValueError("time must be in (0,1]")\n    if not np.isfinite(logits).all():\n        raise ValueError("logits must be finite")\n    shifted = logits - logits.max(axis=-1, keepdims=True)\n    log_probs = shifted - np.log(np.exp(shifted).sum(axis=-1, keepdims=True))\n    nll = -log_probs[np.arange(target_ids.size), target_ids]\n    return float(nll[loss_mask].sum() / time)\n\n\ndef sample_dllm(predict, length: int, steps: int, mask_id: int):\n    """Greedy confidence reveal with one model invocation per iteration."""\n    schedule = linear_mask_schedule(length, steps)\n    if isinstance(mask_id, (bool, np.bool_)) or not isinstance(mask_id, (int, np.integer)):\n        raise ValueError("mask_id must be an integer")\n    vocabulary = None\n    tokens = np.full(length, mask_id, dtype=np.int64)\n    trace = []\n    for step, target_masked in enumerate(schedule[1:], start=1):\n        logits = np.asarray(predict(tokens.copy(), step), dtype=np.float64)\n        if logits.ndim != 2 or logits.shape[0] != length or logits.shape[1] == 0 or not np.isfinite(logits).all():\n            raise ValueError("predict must return finite logits [T,V]")\n        if vocabulary is not None and logits.shape[1] != vocabulary:\n            raise ValueError("predict changed vocabulary size")\n        vocabulary = logits.shape[1]\n        if 0 <= mask_id < vocabulary:\n            raise ValueError("mask_id must lie outside the vocabulary")\n        shifted = logits - logits.max(axis=-1, keepdims=True)\n        probs = np.exp(shifted)\n        probs /= probs.sum(axis=-1, keepdims=True)\n        candidates = probs.argmax(axis=-1)\n        confidence = probs[np.arange(length), candidates]\n        masked = np.flatnonzero(tokens == mask_id)\n        open_count = masked.size - target_masked\n        order = np.lexsort((masked, -confidence[masked]))\n        opened = masked[order[:open_count]]\n        tokens[opened] = candidates[opened]\n        trace.append((tuple(opened), tokens.copy()))\n    return tokens, trace\n', MINIMAL)

assert MINIMAL["linear_mask_schedule"](7,11) == linear_mask_schedule(7,11)
assert MINIMAL["dllm_loss"]([[0.,1.]], [1], np.array([False]), .5) == 0.0
for bad in (True, 1.5, 0):
    expect_failure(ValueError, MINIMAL["linear_mask_schedule"], bad, 4)

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

import importlib.util,subprocess,sys
if importlib.util.find_spec("torch") is None:
    subprocess.check_call([sys.executable,"-m","pip","install","-q","torch"])
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from numbers import Integral,Real
from typing import Any,Callable,Literal
torch.set_num_threads(1)
torch.manual_seed(120)


def _positive_int(name: str, value: object) -> int:
    """Return a strictly positive Python integer and reject booleans explicitly."""

    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    result = int(value)
    if result <= 0:
        raise ValueError(f"{name} must be positive")
    return result


def _nonnegative_int(name: str, value: object) -> int:
    """Return a non-negative Python integer and reject booleans explicitly."""

    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    result = int(value)
    if result < 0:
        raise ValueError(f"{name} must be non-negative")
    return result


def _finite_real(name: str, value: object) -> float:
    """Return a finite real scalar and reject booleans and complex values."""

    if isinstance(value, bool) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real scalar")
    result = float(value)
    if not math.isfinite(result):
        raise ValueError(f"{name} must be finite")
    return result


def _bounded_real(
    name: str,
    value: object,
    *,
    lower: float,
    upper: float,
    lower_inclusive: bool = True,
    upper_inclusive: bool = True,
) -> float:
    """Validate one finite scalar against an explicit interval."""

    result = _finite_real(name, value)
    lower_ok = result >= lower if lower_inclusive else result > lower
    upper_ok = result <= upper if upper_inclusive else result < upper
    if not (lower_ok and upper_ok):
        left = "[" if lower_inclusive else "("
        right = "]" if upper_inclusive else ")"
        raise ValueError(f"{name} must lie in {left}{lower}, {upper}{right}")
    return result


def _finite_array(
    name: str,
    value: object,
    *,
    ndim: int | None = None,
    dtype: np.dtype[Any] | type = np.float64,
) -> np.ndarray:
    """Convert to a NumPy array and enforce shape rank and finite values."""

    array = np.asarray(value, dtype=dtype)
    if ndim is not None and array.ndim != ndim:
        raise ValueError(f"{name} must have rank {ndim}, got shape {array.shape}")
    if array.size == 0:
        raise ValueError(f"{name} must be non-empty")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    return array


def _integer_tokens(name: str, value: object, *, ndim: int = 1) -> np.ndarray:
    """Validate an integer token tensor without silently truncating floats."""

    raw = np.asarray(value)
    if raw.ndim != ndim:
        raise ValueError(f"{name} must have rank {ndim}, got shape {raw.shape}")
    if raw.size == 0:
        raise ValueError(f"{name} must be non-empty")
    if raw.dtype.kind not in "iu":
        raise TypeError(f"{name} must contain integer token IDs")
    return raw.astype(np.int64, copy=False)





@dataclass(frozen=True)
class CorruptionResult:
    """One absorbing-mask forward sample for a one-dimensional token sequence."""

    clean_tokens: np.ndarray
    noisy_tokens: np.ndarray
    masked_positions: np.ndarray
    eligible_positions: np.ndarray
    time: float
    mask_id: int


def linear_mask_schedule(length: int, steps: int) -> tuple[int, ...]:
    """Return monotone target mask counts from ``length`` to zero.

    ``ceil`` gives a declared policy for non-divisible lengths. Duplicate counts are allowed
    when ``steps > length``; such a step performs a model evaluation but opens no token.
    """

    n = _positive_int("length", length)
    s = _positive_int("steps", steps)
    counts = tuple((n * remaining + s - 1) // s for remaining in range(s, -1, -1))
    if counts[0] != n or counts[-1] != 0:
        raise AssertionError("schedule endpoints are inconsistent")
    if any(left < right for left, right in zip(counts, counts[1:])):
        raise AssertionError("mask schedule must be non-increasing")
    return counts


def elbo_time_weight(time: Real) -> float:
    """Return the absorbing-diffusion slice weight ``1 / t`` for ``0 < t <= 1``."""

    t = _bounded_real("time", time, lower=0.0, upper=1.0, lower_inclusive=False)
    return 1.0 / t


def corrupt_absorbing(
    tokens: object,
    *,
    time: Real,
    mask_id: int,
    rng: np.random.Generator,
    eligible_positions: object | None = None,
) -> CorruptionResult:
    """Replace eligible tokens independently by ``mask_id`` with probability ``time``.

    The mask token must be outside the clean sequence so that masked and observed positions
    remain unambiguous. ``eligible_positions`` can protect a fixed prefix or special tokens.
    """

    clean = _integer_tokens("tokens", tokens)
    if isinstance(mask_id, bool) or not isinstance(mask_id, Integral):
        raise TypeError("mask_id must be an integer")
    mask = int(mask_id)
    if np.any(clean == mask):
        raise ValueError("mask_id must not occur in the clean sequence")
    if not isinstance(rng, np.random.Generator):
        raise TypeError("rng must be numpy.random.Generator")
    t = _bounded_real("time", time, lower=0.0, upper=1.0)

    if eligible_positions is None:
        eligible = np.ones(clean.shape, dtype=bool)
    else:
        eligible = np.asarray(eligible_positions)
        if eligible.shape != clean.shape or eligible.dtype.kind != "b":
            raise ValueError("eligible_positions must be a boolean vector matching tokens")
        eligible = eligible.astype(bool, copy=False)

    selected = eligible & (rng.random(clean.shape[0]) < t)
    noisy = clean.copy()
    noisy[selected] = mask
    return CorruptionResult(
        clean_tokens=clean.copy(),
        noisy_tokens=noisy,
        masked_positions=selected,
        eligible_positions=eligible.copy(),
        time=t,
        mask_id=mask,
    )


def stable_log_softmax(logits: object) -> np.ndarray:
    """Compute row-wise log-softmax for a finite ``[positions, vocabulary]`` tensor."""

    values = _finite_array("logits", logits, ndim=2)
    maxima = np.max(values, axis=-1, keepdims=True)
    shifted = values - maxima
    log_normalizer = np.log(np.exp(shifted).sum(axis=-1, keepdims=True))
    result = shifted - log_normalizer
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("log-softmax produced a non-finite value")
    return result


def masked_diffusion_loss(
    logits: object,
    target_tokens: object,
    masked_positions: object,
    *,
    time: Real,
) -> float:
    """Return ``(1/t) * sum(-log p_theta(x_i | x_t))`` on masked positions only."""

    log_probs = stable_log_softmax(logits)
    targets = _integer_tokens("target_tokens", target_tokens)
    if targets.shape[0] != log_probs.shape[0]:
        raise ValueError("target_tokens and logits must have the same position axis")
    if np.any(targets < 0) or np.any(targets >= log_probs.shape[1]):
        raise ValueError("target token ID is outside the vocabulary")
    selected = np.asarray(masked_positions)
    if selected.shape != targets.shape or selected.dtype.kind != "b":
        raise ValueError("masked_positions must be a boolean vector matching target_tokens")
    selected = selected.astype(bool, copy=False)
    weight = elbo_time_weight(time)
    nll = -log_probs[np.arange(targets.size), targets]
    loss = weight * float(nll[selected].sum())
    if not math.isfinite(loss):
        raise FloatingPointError("diffusion loss is non-finite")
    return loss





@dataclass(frozen=True)
class DenoiseStep:
    """Inspectable state transition produced by one denoising model call."""

    step: int
    target_masked: int
    remasked_positions: tuple[int, ...]
    opened_positions: tuple[int, ...]
    masked_after: int
    sequence: tuple[int, ...]


@dataclass(frozen=True)
class DenoiseResult:
    """Final sequence and full trace of a masked-diffusion sampling run."""

    final_tokens: np.ndarray
    schedule: tuple[int, ...]
    steps: tuple[DenoiseStep, ...]
    model_calls: int


def stable_softmax(logits: object) -> np.ndarray:
    """Compute row-wise softmax for finite logits and verify probability normalization."""

    log_probs = stable_log_softmax(logits)
    probabilities = np.exp(log_probs)
    if not np.allclose(probabilities.sum(axis=-1), 1.0, atol=1e-12, rtol=0.0):
        raise FloatingPointError("softmax rows do not sum to one")
    return probabilities


def sample_masked_diffusion(
    predict: Callable[[np.ndarray, int], object],
    *,
    length: int,
    steps: int,
    mask_id: int,
    rng: np.random.Generator,
    selection_policy: Literal["confidence", "random", "entropy"] = "confidence",
    token_policy: Literal["greedy", "categorical"] = "greedy",
    remask_fraction: Real = 0.0,
) -> DenoiseResult:
    """Run a transparent masked-diffusion sampler with one model call per step.

    ``predict`` receives a copy of the current token vector and the one-based step number,
    and must return finite logits of shape ``[length, vocabulary]``. The mask ID must lie
    outside that vocabulary. Opened tokens remain fixed unless explicitly selected for
    remasking. Confidence ties are resolved by lower position index.
    """

    n = _positive_int("length", length)
    s = _positive_int("steps", steps)
    if isinstance(mask_id, bool) or not isinstance(mask_id, Integral):
        raise TypeError("mask_id must be an integer")
    mask = int(mask_id)
    if not isinstance(rng, np.random.Generator):
        raise TypeError("rng must be numpy.random.Generator")
    if selection_policy not in {"confidence", "random", "entropy"}:
        raise ValueError("selection_policy must be 'confidence', 'random', or 'entropy'")
    if token_policy not in {"greedy", "categorical"}:
        raise ValueError("token_policy must be 'greedy' or 'categorical'")
    remask = _bounded_real("remask_fraction", remask_fraction, lower=0.0, upper=1.0)

    schedule = linear_mask_schedule(n, s)
    sequence = np.full(n, mask, dtype=np.int64)
    trace: list[DenoiseStep] = []
    stored_confidence = np.zeros(n)
    expected_vocabulary = None

    for step_number, target_masked in enumerate(schedule[1:], start=1):
        # Select from scores saved when tokens were last predicted. Mask BEFORE
        # this call, so the network conditions on the actual corrupted input.
        opened_before = np.flatnonzero(sequence != mask)
        remasked = np.empty(0, dtype=np.int64)
        remask_count = math.floor(remask * opened_before.size)
        if remask_count:
            order = np.lexsort((opened_before, stored_confidence[opened_before]))
            remasked = opened_before[order[:remask_count]]
            sequence[remasked] = mask
        raw_logits = predict(sequence.copy(), step_number)
        logits = _finite_array("predict output", raw_logits, ndim=2)
        if logits.shape[0] != n:
            raise ValueError("predict output must have one row per sequence position")
        vocabulary = logits.shape[1]
        if vocabulary <= 0:
            raise ValueError("predict output vocabulary must be non-empty")
        if 0 <= mask < vocabulary:
            raise ValueError("mask_id must lie outside the model vocabulary")
        if expected_vocabulary is None:
            expected_vocabulary = vocabulary
        elif vocabulary != expected_vocabulary:
            raise ValueError("predict changed vocabulary size between calls")
        probabilities = stable_softmax(logits)

        if token_policy == "greedy":
            candidates = np.argmax(probabilities, axis=-1).astype(np.int64)
        else:
            candidates = np.array(
                [rng.choice(vocabulary, p=row) for row in probabilities],
                dtype=np.int64,
            )
        candidate_confidence = probabilities[np.arange(n), candidates]

        masked_now = np.flatnonzero(sequence == mask)
        open_count = masked_now.size - target_masked
        if open_count < 0:
            raise AssertionError("sampler has fewer masks than the schedule requires")
        if selection_policy == "confidence":
            order = np.lexsort((masked_now, -candidate_confidence[masked_now]))
            opened = masked_now[order[:open_count]]
        elif selection_policy == "entropy":
            log_probabilities = stable_log_softmax(logits)
            entropy = -(probabilities * log_probabilities).sum(axis=-1)
            order = np.lexsort((masked_now, entropy[masked_now]))
            opened = masked_now[order[:open_count]]
        else:
            opened = rng.permutation(masked_now)[:open_count]
        sequence[opened] = candidates[opened]
        stored_confidence[opened] = candidate_confidence[opened]

        masked_after = int(np.count_nonzero(sequence == mask))
        if masked_after != target_masked:
            raise AssertionError("sampler did not reach the scheduled mask count")
        trace.append(
            DenoiseStep(
                step=step_number,
                target_masked=target_masked,
                remasked_positions=tuple(sorted(map(int, remasked))),
                opened_positions=tuple(sorted(map(int, opened))),
                masked_after=masked_after,
                sequence=tuple(map(int, sequence)),
            )
        )

    if np.any(sequence == mask):
        raise AssertionError("final sequence still contains mask tokens")
    return DenoiseResult(
        final_tokens=sequence.copy(),
        schedule=schedule,
        steps=tuple(trace),
        model_calls=s,
    )


def make_contextual_toy_predictor(
    target_tokens: object,
    *,
    vocabulary_size: int,
) -> Callable[[np.ndarray, int], np.ndarray]:
    """Create a deterministic toy model whose odd positions need their left anchor.

    Even positions are high-confidence anchors. An odd position initially prefers a wrong
    token; once its left anchor is opened correctly, the target becomes high-confidence.
    This makes reveal order and remasking observable without training a neural network.
    """

    target = _integer_tokens("target_tokens", target_tokens)
    vocab = _positive_int("vocabulary_size", vocabulary_size)
    if np.any(target < 0) or np.any(target >= vocab):
        raise ValueError("target token ID is outside vocabulary_size")

    def predict(sequence: np.ndarray, step: int) -> np.ndarray:
        """Return deterministic contextual logits for one sampler step."""

        current = _integer_tokens("sequence", sequence)
        if current.shape != target.shape:
            raise ValueError("sequence shape differs from the toy target")
        _positive_int("step", step)
        logits = np.full((target.size, vocab), -2.5, dtype=np.float64)
        for position, target_id in enumerate(target):
            wrong_id = (int(target_id) + 1) % vocab
            if position % 2 == 0:
                logits[position, int(target_id)] = 3.5
                logits[position, wrong_id] = 0.0
            elif current[position - 1] == target[position - 1]:
                logits[position, int(target_id)] = 3.3
                logits[position, wrong_id] = 0.0
            else:
                logits[position, wrong_id] = 1.5
                logits[position, int(target_id)] = 1.1
        return logits

    return predict



# Synthetic grammar: a,b,a,b,a,b,a,b. Train and evaluation use disjoint pairs.
pairs=[(a,b) for a in range(4) for b in range(4)]
train_pairs=[p for p in pairs if p[0] != p[1]]
test_pairs=[p for p in pairs if p[0] == p[1]]
def examples(items):
    return torch.tensor([[a,b]*4 for a,b in items],dtype=torch.long)
train=examples(train_pairs);heldout=examples(test_pairs)
MASK=4
class TinyMaskedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed=nn.Embedding(5,32)
        self.position=nn.Parameter(torch.randn(1,8,32)*.02)
        self.layer=nn.TransformerEncoderLayer(32,4,64,dropout=0,batch_first=True)
        self.head=nn.Linear(32,4)
    def forward(self, ids):
        return self.head(self.layer(self.embed(ids)+self.position))
model=TinyMaskedModel()
optimizer=torch.optim.AdamW(model.parameters(),lr=.003)
# A fixed validation corruption is used only for evaluation, never optimization.
validation_input=heldout.clone();validation_input[:,2:]=MASK
@torch.no_grad()
def validation_loss():
    model.eval()
    values=model(validation_input)
    return float(F.cross_entropy(values[:,2:].reshape(-1,4),heldout[:,2:].reshape(-1)))
before=validation_loss()
train_losses=[]
for step in range(300):
    model.train()
    clean=train[torch.randint(len(train),(32,))]
    t=.05+.95*torch.rand(32,1)  # declared truncated training-time distribution
    hidden=torch.rand(clean.shape)<t
    hidden[:,:2]=False
    noisy=clean.masked_fill(hidden,MASK)
    logits=model(noisy)
    nll=F.cross_entropy(logits.reshape(-1,4),clean.reshape(-1),reduction="none").reshape_as(clean)
    # Fixed six-token answer normalization, not a random masked-token mean.
    loss=((nll*hidden)/t).sum(dim=1).mean()/6
    optimizer.zero_grad();loss.backward();optimizer.step()
    train_losses.append(float(loss.detach()))
after=validation_loss()
assert np.isfinite(train_losses).all()
assert np.mean(train_losses[-30:]) < np.mean(train_losses[:30])
print({"training_pairs":len(train_pairs),"heldout_pairs":len(test_pairs),
       "heldout_masked_CE_before":round(before,4),"after":round(after,4)})
# Do not tune against the held-out pairs or promise that every policy wins.
model.eval();records=[]
for pair in test_pairs:
    target=np.array(list(pair)*3)  # six answer positions after the two-token prompt
    for policy in ("random","confidence","entropy"):
        calls=[]
        def predict(sequence, step):
            ids=np.concatenate([np.array(pair),np.where(sequence==-1,MASK,sequence)])
            with torch.no_grad():
                logits=model(torch.tensor(ids)[None])[0,2:].numpy().copy()
            calls.append({"input":ids.copy(),"logits":logits.copy()})
            return logits
        result=sample_masked_diffusion(predict,length=6,steps=3,mask_id=-1,
              rng=np.random.default_rng(42),selection_policy=policy)
        assert len(calls)==result.model_calls==3
        accuracy=float(np.mean(result.final_tokens==target))
        records.append((pair,policy,accuracy))
for policy in ("random","confidence","entropy"):
    print(policy, np.mean([r[2] for r in records if r[1]==policy]), "token accuracy")

def work_counts(length, block_size, denoise_passes, finalize=True):
    if any(type(v) is not int or v <= 0 for v in (length,block_size,denoise_passes)):
        raise ValueError("positive integers required")
    sizes=[min(block_size,length-i) for i in range(0,length,block_size)]
    denoise_calls=len(sizes)*denoise_passes
    final_calls=len(sizes) if finalize else 0
    # No prompt prefill, prefix attention, or hardware timing is counted here.
    return dict(denoise_calls=denoise_calls,finalization_calls=final_calls,
                calls=denoise_calls+final_calls,
                positions=sum(sizes)*(denoise_passes+int(finalize)))
assert work_counts(10,4,3)==dict(denoise_calls=9,finalization_calls=3,calls=12,positions=40)
assert work_counts(8,8,4)["calls"]==5
for b in (2,4,8):
    counts=work_counts(8,b,3)
    # A declared dense-model ledger, not bytes measured from this Python run.
    hypothetical_parameters=8_000_000_000
    bytes_per_weight=2
    estimated_weight_bytes=counts["calls"]*hypothetical_parameters*bytes_per_weight
    print(b,counts,"estimated weight GB:",estimated_weight_bytes/1e9)
# The final forward pass must use the clean values, not the last noisy canvas.
last_input=np.array([-1,2,-1,0])
final_tokens=np.array([1,2,3,0])
def cache_features(x):  # visible counterexample, not a Transformer cache
    return np.stack([x,x*x],axis=1)
assert not np.array_equal(cache_features(last_input),cache_features(final_tokens))
print("Clean-block finalization is separate work; cached noisy states are stale.")