Self-review

Run with the self-review skill against .ai/references/review-rules.md (plus code_style.md,
pipelines.md and testing.md for the touched areas), over the whole diff rather than the CI's
src/ + tests/ scope. Three rounds; every finding below is applied. Verdict: READY, with the
items in Left for the reviewer deliberately unresolved.

Blocking (fixed)

1. A bool attention_mask was mutated in place — examples/discrete_diffusion/train_llada2.py:231
valid = attention_mask.to(dtype=torch.bool) followed by valid[:, :prompt_length] = False.
   Tensor.to(dtype=...) returns self when the dtype already matches, so a dataset yielding a bool
   attention mask had its prompt columns zeroed in place — and that same tensor is handed to the model two
   lines later, so the prompt gets attended to as padding. Silent, and invisible with the long masks the
   script's own fallback produces. Now built without mutation
   (attention_mask.to(dtype=torch.bool) & (arange >= prompt_length)) and verified for both dtypes.

2. editing_threshold=0.0 meant opposite things in the two refinement schedulers —
   src/diffusers/schedulers/scheduling_uniform_refinement.py:256
UniformRefinementScheduler gated only on is not None, so 0.0 made confidence > 0.0 true almost
   everywhere and rewrote nearly every committed position, while BlockRefinementScheduler, both docs pages
   and LLaDA2Pipeline.__call__ all document 0.0 as disables editing. Now is not None and > 0.0 in both.

3. Validation stayed behind at the deprecated location — scheduling_block_refinement.py:112-116,
   scheduling_uniform_refinement.py:65
check_inputs still validated sampling_method / threshold, while the scheduler config — their new
   home — validated neither, so BlockRefinementScheduler(sampling_method="multinomail") was accepted and
   silently returned the greedy argmax. Both are now checked in __init__ (threshold in both schedulers,
   sampling_method in BlockRefinementScheduler, the only one that has it).

4. EntropyBoundScheduler produced NaN from a plausible config —
   src/diffusers/schedulers/scheduling_entropy_bound.py:62
   The annealed temperature divides the logits, so t_max=0.0, t_min=0.0 — a reasonable-looking "pure
   greedy" — gave -inf/nan in pred_logits, NaN entropy, garbage acceptance, and NaN self-conditioning
   fed back to the model. The first fix was itself insufficient and a third round caught it: validating
   t_max > 0 alone still lets t_min < 0 < t_max cross exactly zero mid-schedule (t_max=1, t_min=-1, n=4 lands on 0.0 at step 2, which the sampler's own temperature < 0 guard waves through). __init__
   now bounds both ends (t_max > 0, t_min >= 0), which keeps the line strictly positive for either
   ordering of the two. t_min=0 stays legal, and a test walks a whole schedule asserting the logits stay
   finite.

5. A sample whose shape changed mid-schedule was silently re-initialized rather than reported —
   scheduling_uniform_refinement.py:234-241
   The state reset had two redundant branches: the step_index == 0 one was the pre-refactor reset
   mechanism (the pipeline restarting its own counter), made redundant by set_timesteps clearing
   _committed. Now one _committed is None initializer plus an explicit ValueError naming both shapes.

Non-blocking (fixed)

6. examples/discrete_diffusion/sample_llada2.py:220 — still drove all six deprecated __call__ knobs. It
   worked, printing six FutureWarnings per run, and taught the old API. Now builds the scheduler from its
   CLI flags and passes only loop arguments (num_inference_steps, max_post_steps, eos_early_stop).
7. pipeline_llada2.py — dead defensive initializer timestep = self.scheduler.timesteps[-1];
   set_timesteps guarantees at least one iteration, so it was always overwritten (code_style.md, "no
   defensive code").
8. pipeline_llada2.py:279 — _update_finished(sampled_tokens=...) → pred_original_sample, finishing the
   output-field rename inside the pipeline.
9. Type annotations on the new private helpers (DiscreteSchedulerOutput, the callback union).
10. pipeline_diffusion_gemma.py — callback_kwargs dict comprehension → the loop form pipelines.md
    gotcha 9 prescribes. The comprehension was correct (it indexes an explicit dict, not locals()), but the
    rule reads categorically and LLaDA2Pipeline already used the loop.
11. pipeline_llada2.py:227 — added the comment explaining why the deprecated-argument checks in
    check_inputs are load-bearing rather than duplicated validation: this path reaches the config through
    register_to_config, which merges values in without re-running __init__, so those checks are what
    catches a bad argument. They go when the arguments do, at 1.0.0.
12. PR hygiene — a .gitignore entry for /.parity was reverted. The rubric names parity harnesses as
    ephemeral context, so the harness is not part of this PR and a committed ignore rule naming it would
    dangle; the rule lives in .git/info/exclude instead.

Dead code (advisory)

Nothing is likely-dead in the final state. (An earlier staged phase did land DiscreteSchedulerOutput and
the index bookkeeping with no callers at all — a self-review run at that point would rightly have flagged
them; they got their callers in the pipeline phases.)

Intended behaviour changes — the two worth a reviewer's attention

1. sampled_probs is now the probability under the unmodified denoiser distribution, in every
   scheduler. It used to be three different quantities: BlockRefinement's multinomial path gathered from
   the temperature-scaled, top-k/p-filtered distribution, its greedy path from the raw softmax,
   DiscreteDDIM always from the raw one, and EntropyBound from the scaled one. Since threshold /
   editing_threshold compare directly against this number, raising temperature used to silently rescale
   what those thresholds mean. Unifying it makes confidence thresholds temperature-independent and portable
   across schedulers, and changes BlockRefinementScheduler output at temperature != 0. Measured: 3/32
   tokens on the one parity case with temperature != 0, and no change at all to any LLaDA2 output,
   including temperature=0.7, top_k=16, top_p=0.9 — a token only moves when the threshold actually
   straddles the two confidences.
2. The pipelines' per-call defaults are gone; the scheduler's config defaults are what "unspecified" means
   now. The schedulers' own defaults are unchanged (threshold=0.95, editing_threshold=None), but
   LLaDA2Pipeline.__call__ used to default threshold=0.7, editing_threshold=0.5 — so an unqualified
   pipe(prompt=...) commits less eagerly and no longer runs the editing phase. There is no way to have
   this both ways: editing_threshold=None is documented as disable editing, so it can't double as
   "unspecified", and defaulting the config to 0.5 would silently override users who pass None to turn
   editing off. The migration is one line and it is in both the deprecation warning and the pipeline example:
   BlockRefinementScheduler(threshold=0.7, editing_threshold=0.5).

Four more findings came from the parity harness rather than from review, and are worth knowing:

• The released LLaDA2 loop ran one wasted model forward per block. It detected "no masks left" only at
  the top of the next iteration, so it always spent a full step (forward + sample) discovering it. The
  new loop checks after the write: every block is one forward cheaper, with identical output (29 → 28
  recorded steps on the greedy case).
• Editing never ran when the masks cleared exactly on the final scheduled step — step_idx >= n and
  masks_remaining both held, so the released while exited before any editing sweep. The new structure
  preserves this (the editing phase is guarded on the mask-filling loop having breaked early, not on the
  block being full), but it looks like an accident of the old control flow rather than an intended rule, so
  it is flagged rather than "fixed".
• The editing budget is max_post_steps + 1 sweeps, not max_post_steps — the released loop
  incremented before the step and compared after it. Preserved, and pinned by a test.
• A reconfigured scheduler copy is not an option. The first attempt honoured the deprecated per-call
  knobs by running the loop against from_config(self.scheduler.config, ...), which silently broke every
  consumer keyed on identity — pipe.scheduler was not the object being stepped. The knobs are now applied
  to self.scheduler's config for the duration of the call and restored in a finally
  (_scheduler_config_overrides, pipeline_llada2.py:257), which is also how the tokenizer's
  mask_token_id reaches the scheduler. A test asserts both the identity and the restore.

One more argument for a change this PR makes: the inspect.signature dispatch both pipelines used is what
made the harness lie. Its recorder wrapped step as (*args, **kwargs), so the dispatch saw
{args, kwargs} and quietly stopped passing generator — attaching the observer changed the run it was
observing, and produced 398 phantom divergences. Only a signature-sniffing call site can fail that way.

Left for the reviewer

1. The right scope for _scheduler_config_overrides. Restoring a key the config never had adds it as
   None, which save_pretrained would then persist. Harmless for BlockRefinementScheduler, which has
   every key, but the pipeline now accepts any discrete scheduler.
2. How much of the now-duplicated grid / index-recovery assertions to drop from the per-scheduler test
   files, given test_discrete_scheduler_contract.py asserts them for all four.
3. Whether the final-step editing quirk above should be preserved at all.

Verification

1072 passed, 15 skipped

• Parity harness (28 cases, CPU/float32) against a baseline captured before the refactor at the branch's
  base commit: Tier A and Tier B PARITY OK, with only the intended API changes allow-listed (t_in
  int index → float grid, committed_mask added to DiscreteDDIM's output, step_correct's
  (batch, corrector_k) → (batch, seq_len) shape fix, EntropyBound's sampled_probs at ~1e-7, and the
  two behaviour changes above). All eight DiffusionGemma cases and three of the four LLaDA2 cases are
  bit-identical in sequences; the fourth matches once its editing_threshold is stated explicitly. The
  harness itself is not part of this PR.
• check_copies, check_dummies, check_doc_toc, ruff check, ruff format and
  doc-builder style --check_only are clean over the touched files. utils/check_inits.py cannot run in
  this environment (it tries to import a sibling src/transformers) — pre-existing and unrelated.

Checkpoint-side migration

google/diffusiongemma-26B-A4B-it's scheduler_config.json names BlockRefinementScheduler in uniform
mode, which now raises on step; it needs repointing to UniformRefinementScheduler alongside this PR.