gemma/diffusion/_sampler_test.py

# Copyright 2026 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import dataclasses
from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import dialog
from gemma.diffusion import _chat_sampler
from gemma.diffusion import _early_stopping
from gemma.diffusion import _models
from gemma.diffusion import _sampler
from gemma.diffusion import _transformer
from gemma.gm.nn.gemma4 import _config
from gemma.gm.nn.gemma4 import _modules
from gemma.gm.nn.gemma4 import _transformer as _gemma4_transformer
from gemma.gm.text import _sampler_loop
from gemma.gm.text import _sampling
from gemma.gm.text import _tokenizer
import jax
import jax.numpy as jnp
import numpy as np


class _MockSpecialTokens(_tokenizer.SpecialTokens):
  PAD = 0
  EOS = 1
  BOS = 2
  UNK = 3
  MASK = 4
  CUSTOM = 5
  START_OF_TURN = 6
  END_OF_TURN = 7
  START_OF_IMAGE = 8
  END_OF_IMAGE = 9
  BEGIN_OF_TOOL_RESPONSE = 10
  END_OF_TOOL_RESPONSE = 11


@dataclasses.dataclass(frozen=True)
class _MockTokenizer:
  vocab_size: int = 32

  VERSION = None
  FORBIDDEN_TOKENS = ()
  FORMAT = dialog.Format.GEMMA3

  special_tokens = _MockSpecialTokens

  def encode(self, text, *, add_bos=False, add_eos=False):
    tokens = [ord(c) % self.vocab_size for c in text]
    if add_bos:
      tokens.insert(0, self.special_tokens.BOS)
    if add_eos:
      tokens.append(self.special_tokens.EOS)
    return tokens

  def decode(self, ids):
    if hasattr(ids, 'tolist'):
      ids = ids.tolist()
    if isinstance(ids, int):
      ids = [ids]
    return ''.join(chr(i + 65) if 0 <= i < 26 else '?' for i in ids)


_SMALL_CONFIG = _config.TransformerConfig(
    num_embed=32,
    embed_dim=8,
    num_heads=2,
    num_kv_heads=1,
    head_dim=4,
    hidden_dim=16,
    attention_types=[_modules.AttentionType.GLOBAL],
    kv_cache_sharing_config=None,
    use_post_attn_norm=True,
    use_post_ffw_norm=True,
    final_logit_softcap=None,
    global_rope_proportion=1.0,
)

_SMALL_SC_CONFIG = _transformer.SelfConditioningConfig(
    features=_SMALL_CONFIG.embed_dim,
    hidden_dim=_SMALL_CONFIG.hidden_dim,
)


def _sample_with_while_loop(
    sampler,
    params,
    cache,
    batch_size,
    canvas_length,
    max_denoising_steps,
    rng,
):
  """Runs sample_next_canvas which uses jax.lax.while_loop internally."""
  return sampler.sample_next_canvas(
      canvas_length=canvas_length,
      max_denoising_steps=max_denoising_steps,
      batch_size=batch_size,
      cache=cache,
      params=params,
      rng=rng,
  )


def _sample_with_for_loop(
    sampler,
    params,
    cache,
    batch_size,
    canvas_length,
    embed_dim,
    vocab_size,
    max_denoising_steps,
    rng,
    num_steps_override=None,
):
  """Reimplements the denoising loop with a Python for loop."""
  initial_canvas_rng, step_rng = jax.random.split(rng)

  if cache is not None:
    cache_layer = list(cache.values())[0]
    cache_len = cache_layer['k'].shape[1]
    samples_in_cache = cache_layer['end_index']
    positions = samples_in_cache[:, None] + jnp.arange(canvas_length)[None, :]
  else:
    cache_len = None
    samples_in_cache = None
    positions = jnp.broadcast_to(
        jnp.arange(canvas_length)[None, :], (batch_size, canvas_length)
    )

  attention_mask = _sampler._make_global_attention_mask(
      batch_size=batch_size,
      canvas_length=canvas_length,
      cache_length=cache_len,
      num_valid_tokens=samples_in_cache,
  )

  canvas = sampler.diffusion_process.get_initial_sample(
      rng=initial_canvas_rng,
      batch_size=batch_size,
      canvas_length=canvas_length,
      text_vocab_size=vocab_size,
  )
  sc_embeddings = jnp.zeros(
      (batch_size, canvas_length, embed_dim),
      dtype=jnp.bfloat16,
  )

  num_steps = (
      num_steps_override
      if num_steps_override is not None
      else max_denoising_steps
  )

  for step in range(num_steps):
    step_rng, sample_rng = jax.random.split(step_rng)
    current_noise_proportion = jnp.full(
        (batch_size,), 1.0 - step / max_denoising_steps
    )
    target_noise_proportion = jnp.full(
        (batch_size,), 1.0 - (step + 1) / max_denoising_steps
    )
    step_output = sampler.sample_step(
        canvas=canvas,
        sc_embeddings=sc_embeddings,
        cache=cache,
        positions=positions,
        attention_mask=attention_mask,
        current_noise_proportion=current_noise_proportion,
        target_noise_proportion=target_noise_proportion,
        params=params,
        rng=sample_rng,
    )
    canvas = step_output.sampled_tokens
    sc_embeddings = step_output.sc_embeddings

  return canvas


class SamplerTest(parameterized.TestCase):

  def test_sampler_step_runs(self):
    """Tests that DiffusionSampler.sample_step runs without errors.

    This is a basic smoke test for a single denoising step in isolation.
    Verifies that sample_step can be initialized and applied with minimal
    inputs (no cache, no positional encodings), and that the output shape
    matches [batch_size, canvas_length].
    """

    batch_size = 1
    canvas_length = 4
    vocab_size = _SMALL_CONFIG.num_embed
    embed_dim = _SMALL_CONFIG.embed_dim

    model = _models.DiffusionGemma_26B_A4B(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )

    sampler = _sampler.DiffusionSampler(
        model=model,
        end_tokens=(99,),
        forbidden_tokens=None,
        sampling=_sampling.Greedy(),
        cache_length=16,
        special_tokens=None,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=2,
        text_vocab_size=vocab_size,
    )

    rng = jax.random.PRNGKey(0)
    params = model.init(
        rngs=rng,
        tokens=jnp.ones((batch_size, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (batch_size, canvas_length, embed_dim), dtype=jnp.bfloat16
        ),
        attention_mask=jnp.ones(
            (batch_size, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        method=model.call_with_self_conditioning,
    )['params']

    output = sampler.sample_step(
        canvas=jnp.ones((batch_size, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (batch_size, canvas_length, embed_dim), dtype=jnp.bfloat16
        ),
        cache=None,
        positions=None,
        attention_mask=jnp.ones(
            (batch_size, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        current_noise_proportion=jnp.full(
            (batch_size,), 0.5, dtype=jnp.float32
        ),
        target_noise_proportion=jnp.full((batch_size,), 0.4, dtype=jnp.float32),
        params=params,
        rng=rng,
    )
    self.assertEqual(output.sampled_tokens.shape, (batch_size, canvas_length))

  def test_sample_from_predictions_final_step(self):
    """Tests that SampleFromPredictions commits all tokens on the final step."""
    logits = jnp.zeros((1, 4, 32), dtype=jnp.float32)
    logits = logits.at[:, :, 5].set(20.0)
    sampler_fn = _sampler.SampleFromPredictions(
        entropy_bound=0.0,  # Setting to 0.0 to accept no tokens.
        text_vocab_size=32,
    )
    tokens = sampler_fn(
        rng=jax.random.PRNGKey(0),
        denoiser_logits=logits,
        canvas=jnp.zeros((1, 4), dtype=jnp.int32),
        current_noise_proportion=jnp.array([0.25]),
        target_noise_proportion=jnp.array([0.0]),  # Final step.
    )
    np.testing.assert_array_equal(tokens, jnp.full((1, 4), 5, dtype=jnp.int32))

  @parameterized.named_parameters(
      dict(
          testcase_name='no_cache',
          use_cache=False,
      ),
      dict(
          testcase_name='with_cache',
          use_cache=True,
      ),
  )
  def test_sample_next_canvas(self, use_cache):
    """Tests that DiffusionSampler.sample_next_canvas runs end-to-end.

    Runs the full multi-step denoising loop that produces a complete canvas.
    Parameterized over two modes:
    - no_cache: The canvas is generated without a KV cache (standalone
      denoising, no prior context).
    - with_cache: The canvas is generated with a pre-allocated KV cache,
      simulating a conversational setting where prior tokens are cached.

    Args:
      use_cache: Whether to use a KV cache.
    """

    batch_size = 1
    canvas_length = 4
    cache_length = 10
    vocab_size = _SMALL_CONFIG.num_embed
    embed_dim = _SMALL_CONFIG.embed_dim

    model = _models.DiffusionGemma_26B_A4B(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )

    if use_cache:
      cache = _SMALL_CONFIG.init_cache(
          batch_size=batch_size,
          dtype=jnp.bfloat16,
          cache_length=cache_length,
      )
    else:
      cache = None

    sampler = _sampler.DiffusionSampler(
        model=model,
        end_tokens=(99,),
        forbidden_tokens=None,
        sampling=_sampling.Greedy(),
        cache_length=cache_length,
        special_tokens=None,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=2,
        text_vocab_size=vocab_size,
    )

    rng = jax.random.PRNGKey(0)
    params = model.init(
        rngs=rng,
        tokens=jnp.ones((batch_size, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (batch_size, canvas_length, embed_dim), dtype=jnp.bfloat16
        ),
        attention_mask=jnp.ones(
            (batch_size, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        method=model.call_with_self_conditioning,
    )['params']

    output = sampler.sample_next_canvas(
        canvas_length=canvas_length,
        max_denoising_steps=2,
        batch_size=batch_size,
        cache=cache,
        params=params,
        rng=rng,
    )
    self.assertEqual(output.shape, (batch_size, canvas_length))

  def test_sample_next_canvas_unpadded_positions(self):
    """Tests that sample_next_canvas uses unpadded offsets for positions."""
    batch_size = 2
    canvas_length = 4
    cache_length = 10
    vocab_size = _SMALL_CONFIG.num_embed
    embed_dim = _SMALL_CONFIG.embed_dim

    model = _models.DiffusionGemma_26B_A4B(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )
    cache = _SMALL_CONFIG.init_cache(
        batch_size=batch_size,
        dtype=jnp.bfloat16,
        cache_length=cache_length,
    )
    cache['layer_0']['end_index'] = jnp.array([5, 5], dtype=jnp.int32)

    sampler = _sampler.DiffusionSampler(
        model=model,
        end_tokens=(99,),
        forbidden_tokens=None,
        sampling=_sampling.Greedy(),
        cache_length=cache_length,
        special_tokens=None,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=1,
        text_vocab_size=vocab_size,
    )

    rng = jax.random.PRNGKey(0)
    params = model.init(
        rngs=rng,
        tokens=jnp.ones((batch_size, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (batch_size, canvas_length, embed_dim), dtype=jnp.bfloat16
        ),
        attention_mask=jnp.ones(
            (batch_size, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        method=model.call_with_self_conditioning,
    )['params']

    # Sequence 0 has 2 valid prompt tokens and 3 padding tokens.
    full_attention_mask = jnp.array(
        [
            [1, 1, 0, 0, 0, 1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        ],
        dtype=jnp.bool_,
    )

    recorded_positions = None
    original_sample_step = _sampler.DiffusionSampler.sample_step

    def tracking_sample_step(self_sampler, *args, **kwargs):
      nonlocal recorded_positions
      recorded_positions = kwargs['positions']
      return original_sample_step(self_sampler, *args, **kwargs)

    with mock.patch.object(
        _sampler.DiffusionSampler,
        'sample_step',
        new=tracking_sample_step,
    ):
      with jax.disable_jit():
        sampler.sample_next_canvas(
            canvas_length=canvas_length,
            max_denoising_steps=1,
            batch_size=batch_size,
            cache=cache,
            params=params,
            rng=rng,
            full_attention_mask=full_attention_mask,
        )

    expected_positions = jnp.array(
        [
            [2, 3, 4, 5],
            [5, 6, 7, 8],
        ],
        dtype=jnp.int32,
    )
    np.testing.assert_array_equal(recorded_positions, expected_positions)

  def test_append_tokens_to_cache(self):
    """Tests that append_tokens_to_cache correctly advances the cache index.

    Verifies two sequential appends: first inserting `canvas_length` tokens
    into an empty cache and checking the end_index advances to
    `canvas_length`, then inserting another batch and confirming it advances
    to `2 * canvas_length`. This ensures the cache write pointer accumulates
    correctly across multiple calls.
    """
    batch_size = 1
    canvas_length = 4
    cache_length = 16
    vocab_size = _SMALL_CONFIG.num_embed

    model = _models.DiffusionGemma_26B_A4B(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )

    sampler = _sampler.DiffusionSampler(
        model=model,
        end_tokens=(99,),
        forbidden_tokens=None,
        sampling=_sampling.Greedy(),
        cache_length=cache_length,
        special_tokens=None,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=2,
        text_vocab_size=vocab_size,
    )

    cache = _SMALL_CONFIG.init_cache(
        batch_size=batch_size,
        dtype=jnp.bfloat16,
        cache_length=cache_length,
    )

    tokens = jnp.ones((batch_size, canvas_length), dtype=jnp.int32)

    rng = jax.random.PRNGKey(0)
    params = model.init(
        rngs=rng,
        tokens=tokens,
        cache=cache,
        positions=jnp.arange(canvas_length)[None, :],
        attention_mask=jnp.ones(
            (batch_size, canvas_length, cache_length), dtype=jnp.bool_
        ),
    )['params']

    updated_cache = sampler.append_tokens_to_cache(
        tokens=tokens,
        cache=cache,
        params=params,
    )

    first_layer = list(updated_cache.values())[0]
    expected_end_index = jnp.full((batch_size,), canvas_length, dtype=jnp.int32)
    np.testing.assert_array_equal(first_layer['end_index'], expected_end_index)

    updated_cache_2 = sampler.append_tokens_to_cache(
        tokens=tokens,
        cache=updated_cache,
        params=params,
    )
    first_layer_2 = list(updated_cache_2.values())[0]
    expected_end_index_2 = jnp.full(
        (batch_size,), 2 * canvas_length, dtype=jnp.int32
    )
    np.testing.assert_array_equal(
        first_layer_2['end_index'], expected_end_index_2
    )

  def test_make_global_attention_no_cache(self):
    """Tests that the global attention mask is all-ones when no cache is used.

    Without a cache, the mask should be a simple [batch, canvas, canvas]
    tensor of ones, allowing full bidirectional self-attention.
    """
    mask = _sampler._make_global_attention_mask(
        batch_size=2,
        canvas_length=4,
        cache_length=None,
        num_valid_tokens=None,
    )
    np.testing.assert_array_equal(mask, jnp.ones((2, 4, 4), dtype=jnp.bool_))

  def test_make_global_attention_mask_batched_edge_cases(self):
    """Tests the global attention mask with a KV cache across edge cases.

    The mask shape is [batch, canvas_length, cache_length] with 1s for
    valid positions and 0s for padding. All canvas tokens attend to the
    same set of valid positions (no causal constraint). Edge cases tested
    as a single batched call:
    - no_wrap: Partially filled cache (3 of 10 slots); total valid
      positions = 3 + 4 = 7, so positions 7-9 are masked out.
    - with_wrap: Nearly full cache (8 valid); total valid = min(12, 10)
      = 10, so all positions are visible.
    - empty_cache: Zero valid cache tokens; only the 4 canvas tokens are
      visible, positions 4-9 are masked out.
    - overfull_cache: More tokens written than cache_length (12 > 10);
      all positions should be visible since the cache is saturated.
    """
    # cache_length=10, canvas_length=4
    # Test multiple global mask edge cases in a single batched call
    edge_cases = [
        (
            'no_wrap',
            3,
            jnp.array([
                # total_valid = min(3+4, 10) = 7
                # mask should be True for 0..6, False for 7..9
                [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
                [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
                [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
                [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
            ]).reshape(4, 10),
        ),
        (
            'with_wrap',
            8,
            jnp.array([
                # total_valid = min(8+4, 10) = 10
                # mask should be all True
                [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, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            ]).reshape(4, 10),
        ),
        (
            'empty_cache',
            0,
            jnp.array([
                # total_valid = min(0+4, 10) = 4
                # mask should be True for 0..3, False for 4..9
                [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
                [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
                [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
                [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
            ]).reshape(4, 10),
        ),
        (
            'overfull_cache',
            12,
            jnp.array([
                # total_valid = min(12+4, 10) = 10
                # mask should be all True
                [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, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            ]).reshape(4, 10),
        ),
    ]

    num_valid_cache_tokens = jnp.array([case[1] for case in edge_cases])

    mask = _sampler._make_global_attention_mask(
        batch_size=len(edge_cases),
        canvas_length=4,
        cache_length=10,
        num_valid_tokens=num_valid_cache_tokens,
    )

    for i, (name, _, expected) in enumerate(edge_cases):
      with self.subTest(name=name):
        np.testing.assert_array_equal(mask[i], expected)

  def test_make_global_attention_mask_with_full_attention_mask(self):
    """Tests global attention mask combined with a prompt full_attention_mask."""
    batch_size = 2
    canvas_length = 3
    cache_length = 8

    # 4 valid cache tokens (so total_valid = min(4 + 3, 8) = 7)
    num_valid_tokens = jnp.array([4, 4], dtype=jnp.int32)

    # Batch 0: has pad at index 2 (so mask is [1, 1, 0, 1, 1, 1, 1, 1])
    # Batch 1: has pad at index 1 and 3 (so mask is [1, 0, 1, 0, 1, 1, 1, 1])
    full_attention_mask = jnp.array(
        [
            [1, 1, 0, 1, 1, 1, 1, 1],
            [1, 0, 1, 0, 1, 1, 1, 1],
        ],
        dtype=jnp.bool_,
    )

    mask = _sampler._make_global_attention_mask(
        batch_size=batch_size,
        canvas_length=canvas_length,
        cache_length=cache_length,
        num_valid_tokens=num_valid_tokens,
        full_attention_mask=full_attention_mask,
    )

    # Base mask (without full_attention_mask) for total_valid = 7 would be:
    # [1, 1, 1, 1, 1, 1, 1, 0] for both batches.
    # With full_attention_mask, it should be:
    # Batch 0: [1, 1, 0, 1, 1, 1, 1, 0]
    # Batch 1: [1, 0, 1, 0, 1, 1, 1, 0]

    expected_batch_0 = jnp.array(
        [
            [1, 1, 0, 1, 1, 1, 1, 0],
            [1, 1, 0, 1, 1, 1, 1, 0],
            [1, 1, 0, 1, 1, 1, 1, 0],
        ],
        dtype=jnp.bool_,
    )

    expected_batch_1 = jnp.array(
        [
            [1, 0, 1, 0, 1, 1, 1, 0],
            [1, 0, 1, 0, 1, 1, 1, 0],
            [1, 0, 1, 0, 1, 1, 1, 0],
        ],
        dtype=jnp.bool_,
    )

    np.testing.assert_array_equal(mask[0], expected_batch_0)
    np.testing.assert_array_equal(mask[1], expected_batch_1)

  def test_make_causal_attention_mask_no_cache(self):
    """Tests that the causal mask is lower-triangular when no cache is used.

    Without a cache, the mask should be a standard lower-triangular matrix
    broadcast across the batch dimension, enforcing standard autoregressive
    left-to-right attention.
    """
    mask = _sampler._make_causal_attention_mask(
        batch_size=2,
        canvas_length=7,
        cache_length=None,
        num_valid_cache_tokens=None,
    )
    expected = jnp.tril(jnp.ones((7, 7), dtype=jnp.bool_))
    expected = jnp.broadcast_to(expected[jnp.newaxis, :, :], (2, 7, 7))
    np.testing.assert_array_equal(mask, expected)

  def test_make_causal_attention_mask_batched_edge_cases(self):
    """Tests the causal attention mask for cache insertion across edge cases.

    When appending canvas tokens to a ring-buffer KV cache, the mask must
    handle wrap-around correctly. Each canvas token should attend to all
    previously cached tokens plus all preceding tokens in the current
    canvas, but NOT future canvas tokens. Edge cases tested:
    - no_wrap: 2 valid cache tokens; new tokens write to indices 2-5,
      no wrap-around occurs.
    - exactly_full: 6 valid tokens; the last canvas token writes to
      index 9, exactly filling the cache with no wrap.
    - with_wrap: 8 valid tokens; writes start at index 8 and wrap to
      indices 0-1, testing ring-buffer wrap-around.
    - full_cache: 10 valid tokens (cache already full); writes wrap
      starting at index 0, overwriting the oldest entries.
    - overfull_cache: 12 valid tokens (more than cache capacity); writes
      start at index 2 (12 % 10), also wrapping, verifying the
      modular arithmetic is correct when the counter exceeds capacity.
    """
    # cache_length=10, seq_len=4
    # Test multiple edge cases in a single batched call
    edge_cases = [
        (
            'no_wrap',
            2,
            jnp.array([
                [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
                [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
                [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
                [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
            ]),
        ),
        (
            'exactly_full',
            6,
            jnp.array([
                [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
                [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            ]),
        ),
        (
            'with_wrap',
            8,
            jnp.array([
                [0, 0, 1, 1, 1, 1, 1, 1, 1, 0],
                [0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 0, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            ]),
        ),
        (
            'full_cache',
            10,
            jnp.array([
                [1, 0, 0, 0, 1, 1, 1, 1, 1, 1],
                [1, 1, 0, 0, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 0, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            ]),
        ),
        (
            'overfull_cache',
            12,
            jnp.array([
                [1, 1, 1, 0, 0, 0, 1, 1, 1, 1],
                [1, 1, 1, 1, 0, 0, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 0, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            ]),
        ),
    ]

    num_valid_cache_tokens = jnp.array([case[1] for case in edge_cases])

    mask = _sampler._make_causal_attention_mask(
        batch_size=len(edge_cases),
        canvas_length=4,
        cache_length=10,
        num_valid_cache_tokens=num_valid_cache_tokens,
    )

    for i, (name, _, expected) in enumerate(edge_cases):
      with self.subTest(name=name):
        np.testing.assert_array_equal(mask[i], expected)

  def test_diffusion_sample_step_runs(self):
    """Tests the DiffusionSamplerLoop._sample_step integration.

    Unlike test_sampler_step_runs which tests the raw DiffusionSampler,
    this tests the full sampling pipeline through DiffusionSamplerLoop,
    which orchestrates a canvas generation, stop-token truncation, and
    cache update in a single step. Verifies that:
    - predicted_tokens has the correct shape.
    - The step counter advances by canvas_length.
    - The batch is not prematurely marked as done.
    - The cache end_index advances by canvas_length.
    """
    batch_size = 1
    canvas_length = 4
    cache_length = 16
    max_out_length = 16
    max_denoising_steps = 2
    vocab_size = _SMALL_CONFIG.num_embed
    embed_dim = _SMALL_CONFIG.embed_dim

    model = _models.DiffusionGemma_26B_A4B(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )

    sampler_loop = _sampler.DiffusionSampler(
        model=model,
        end_tokens=(99,),
        forbidden_tokens=None,
        sampling=_sampling.Greedy(),
        cache_length=cache_length,
        special_tokens=None,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=max_denoising_steps,
        text_vocab_size=vocab_size,
    )

    cache = _SMALL_CONFIG.init_cache(
        batch_size=batch_size,
        dtype=jnp.bfloat16,
        cache_length=cache_length,
    )

    rng = jax.random.PRNGKey(0)
    init_rng, step_rng = jax.random.split(rng)

    params = model.init(
        rngs=init_rng,
        tokens=jnp.ones((batch_size, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (batch_size, canvas_length, embed_dim), dtype=jnp.bfloat16
        ),
        attention_mask=jnp.ones(
            (batch_size, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        method=model.call_with_self_conditioning,
    )['params']

    init_state = _sampler_loop.SamplingState(
        step=jnp.int32(0),
        done=jnp.zeros((batch_size,), dtype=jnp.bool_),
        last_token=jnp.zeros((batch_size,), dtype=jnp.int32),
        last_token_pos=jnp.zeros((batch_size,), dtype=jnp.int32),
        predicted_tokens=jnp.zeros(
            (batch_size, max_out_length), dtype=jnp.int32
        ),
        cache=cache,
        rng=step_rng,
        init_cache_length=jnp.int32(0),
        full_attention_mask=jnp.ones(
            (batch_size, cache_length), dtype=jnp.bool_
        ),
    )

    with jax.disable_jit():
      new_state = sampler_loop._sample_step(
          init_state,
          params=params,
      )

    self.assertEqual(
        new_state.predicted_tokens.shape, (batch_size, max_out_length)
    )
    self.assertEqual(new_state.step, canvas_length)
    self.assertFalse(jnp.all(new_state.done))

    cache_layer = list(new_state.cache.values())[0]
    expected_end_index = jnp.full((batch_size,), canvas_length, dtype=jnp.int32)
    np.testing.assert_array_equal(cache_layer['end_index'], expected_end_index)

  def test_chat_sampler_runs(self):
    """Tests the full ChatSampler end-to-end with a mock tokenizer.

    Exercises the complete inference stack: ChatSampler -> Sampler ->
    DiffusionSamplerLoop -> DiffusionSampler, including tokenization,
    prompt encoding, multi-step diffusion sampling, and decoding back to
    a string. Uses a mock tokenizer that maps characters to token IDs
    via their ordinal values (mod vocab_size). Verifies that the output
    is a valid string.
    """
    canvas_length = 4
    cache_length = 512
    max_out_length = 16
    max_denoising_steps = 2
    vocab_size = _SMALL_CONFIG.num_embed
    embed_dim = _SMALL_CONFIG.embed_dim

    model = _models.DiffusionGemma_26B_A4B(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )

    rng = jax.random.PRNGKey(0)

    params = model.init(
        rngs=rng,
        tokens=jnp.ones((1, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (1, canvas_length, embed_dim), dtype=jnp.bfloat16
        ),
        attention_mask=jnp.ones(
            (1, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        method=model.call_with_self_conditioning,
    )['params']

    tokenizer = _MockTokenizer(vocab_size=vocab_size)

    chat_sampler = _chat_sampler.ChatSampler(
        model=model,
        params=params,
        tokenizer=tokenizer,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=max_denoising_steps,
        cache_length=cache_length,
        max_out_length=max_out_length,
    )

    with jax.disable_jit():
      output = chat_sampler.chat('hello')

    self.assertIsInstance(output, str)

  def test_chat_sampler_gemma4_dispatch(self):
    """Tests that diffusion ChatSampler works when _is_gemma4 is True."""
    canvas_length = 4
    cache_length = 512
    max_out_length = 16
    max_denoising_steps = 2
    vocab_size = _SMALL_CONFIG.num_embed
    embed_dim = _SMALL_CONFIG.embed_dim

    class _DiffusionGemmaTransformer(_models.DiffusionGemma_26B_A4B):
      INFO = _gemma4_transformer.ModelInfo(tokenizer_version=4)

    model = _DiffusionGemmaTransformer(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )

    rng = jax.random.PRNGKey(0)

    params = model.init(
        rngs=rng,
        tokens=jnp.ones((1, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (1, canvas_length, embed_dim), dtype=jnp.float32
        ),
        attention_mask=jnp.ones(
            (1, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        method=model.call_with_self_conditioning,
    )['params']

    tokenizer = _MockTokenizer(vocab_size=vocab_size)

    chat_sampler = _chat_sampler.ChatSampler(
        model=model,
        params=params,
        tokenizer=tokenizer,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=max_denoising_steps,
        cache_length=cache_length,
        max_out_length=max_out_length,
    )

    # Verify _is_gemma4 is True — this is the condition that triggers the bug.
    self.assertTrue(chat_sampler._is_gemma4)  # pylint: disable=protected-access

    # Verify the diffusion sampler is used (not Gemma4Sampler) by tracking
    # calls to sample_next_canvas, which only DiffusionSampler invokes.
    original = _sampler.DiffusionSampler.sample_next_canvas
    sample_next_canvas_called = False

    def tracking_sample_next_canvas(*args, **kwargs):
      nonlocal sample_next_canvas_called
      sample_next_canvas_called = True
      return original(*args, **kwargs)

    with mock.patch.object(
        _sampler.DiffusionSampler,
        'sample_next_canvas',
        new=tracking_sample_next_canvas,
    ):
      with jax.disable_jit():
        output = chat_sampler.chat('hello')

    self.assertIsInstance(output, str)
    self.assertTrue(sample_next_canvas_called)

  @parameterized.named_parameters(
      dict(
          testcase_name='no_cache',
          use_cache=False,
      ),
      dict(
          testcase_name='with_cache',
          use_cache=True,
      ),
  )
  def test_sample_next_canvas_while_loop_matches_for_loop(self, use_cache):
    """Verifies while_loop-based denoising is numerically identical to a for loop.

    Runs sample_next_canvas (which uses jax.lax.while_loop internally) and an
    explicit for-loop reference that calls sample_step sequentially. Asserts
    the final token outputs are identical, catching any silent numerical
    divergence introduced by the while_loop refactor.

    Args:
      use_cache: Whether to use a KV cache.
    """
    batch_size = 2
    canvas_length = 4
    cache_length = 16
    max_denoising_steps = 8
    vocab_size = _SMALL_CONFIG.num_embed
    embed_dim = _SMALL_CONFIG.embed_dim

    model = _models.DiffusionGemma_26B_A4B(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )

    sampler = _sampler.DiffusionSampler(
        model=model,
        end_tokens=(99,),
        forbidden_tokens=None,
        sampling=_sampling.Greedy(),
        cache_length=cache_length,
        special_tokens=None,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=max_denoising_steps,
        text_vocab_size=vocab_size,
    )

    if use_cache:
      cache = _SMALL_CONFIG.init_cache(
          batch_size=batch_size,
          dtype=jnp.bfloat16,
          cache_length=cache_length,
      )
    else:
      cache = None

    rng = jax.random.PRNGKey(42)
    params = model.init(
        rngs=rng,
        tokens=jnp.ones((batch_size, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (batch_size, canvas_length, embed_dim), dtype=jnp.bfloat16
        ),
        attention_mask=jnp.ones(
            (batch_size, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        method=model.call_with_self_conditioning,
    )['params']

    while_loop_output = _sample_with_while_loop(
        sampler,
        params,
        cache,
        batch_size,
        canvas_length,
        max_denoising_steps,
        rng,
    )
    for_loop_output = _sample_with_for_loop(
        sampler,
        params,
        cache,
        batch_size,
        canvas_length,
        embed_dim,
        vocab_size,
        max_denoising_steps,
        rng,
    )

    np.testing.assert_array_equal(while_loop_output, for_loop_output)

  def test_early_stopping_terminates_after_one_step(self):
    """Tests that an EarlyStopFn that always returns True stops after step 0.

    Creates a sampler with an early_stop_fn that returns True unconditionally.
    Compares the output against a manual 1-iteration for-loop reference that
    uses the same noise schedule (max_denoising_steps=8), so the noise
    proportions at step 0 are identical.
    """

    class AlwaysStop:

      def should_stop(self, *, step, canvas, previous_canvas, logits):
        del step, canvas, previous_canvas, logits
        return jnp.bool_(True)

    batch_size = 2
    canvas_length = 4
    cache_length = 16
    max_denoising_steps = 8
    vocab_size = _SMALL_CONFIG.num_embed
    embed_dim = _SMALL_CONFIG.embed_dim

    model = _models.DiffusionGemma_26B_A4B(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )

    rng = jax.random.PRNGKey(42)
    params = model.init(
        rngs=rng,
        tokens=jnp.ones((batch_size, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (batch_size, canvas_length, embed_dim), dtype=jnp.float32
        ),
        attention_mask=jnp.ones(
            (batch_size, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        method=model.call_with_self_conditioning,
    )['params']

    cache = _SMALL_CONFIG.init_cache(
        batch_size=batch_size,
        dtype=jnp.float32,
        cache_length=cache_length,
    )

    sampler_early = _sampler.DiffusionSampler(
        model=model,
        end_tokens=(99,),
        forbidden_tokens=None,
        sampling=_sampling.Greedy(),
        cache_length=cache_length,
        special_tokens=None,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=max_denoising_steps,
        early_stop_fn=AlwaysStop(),
        text_vocab_size=vocab_size,
    )

    early_output = sampler_early.sample_next_canvas(
        canvas_length=canvas_length,
        max_denoising_steps=max_denoising_steps,
        batch_size=batch_size,
        cache=cache,
        params=params,
        rng=rng,
    )

    # Reference: run 1 iteration of the for-loop with the SAME noise schedule
    # (max_denoising_steps=8, but only 1 step executed).
    ref_output = _sample_with_for_loop(
        sampler_early,
        params,
        cache,
        batch_size,
        canvas_length,
        embed_dim,
        vocab_size,
        max_denoising_steps=max_denoising_steps,
        rng=rng,
        num_steps_override=1,
    )

    np.testing.assert_array_equal(early_output, ref_output)

  def test_token_stability_early_stop(self):
    """Tests that TokenStabilityEarlyStop runs end-to-end and produces valid output.

    Uses the built-in TokenStabilityEarlyStop to verify it integrates correctly
    with the while_loop. Also checks determinism: two calls with the same
    RNG produce identical results.
    """
    batch_size = 1
    canvas_length = 4
    cache_length = 16
    max_denoising_steps = 10
    vocab_size = _SMALL_CONFIG.num_embed
    embed_dim = _SMALL_CONFIG.embed_dim

    model = _models.DiffusionGemma_26B_A4B(
        config=_SMALL_CONFIG,
        self_conditioning_config=_SMALL_SC_CONFIG,
    )

    rng = jax.random.PRNGKey(123)
    params = model.init(
        rngs=rng,
        tokens=jnp.ones((batch_size, canvas_length), dtype=jnp.int32),
        sc_embeddings=jnp.ones(
            (batch_size, canvas_length, embed_dim), dtype=jnp.float32
        ),
        attention_mask=jnp.ones(
            (batch_size, canvas_length, canvas_length), dtype=jnp.bool_
        ),
        method=model.call_with_self_conditioning,
    )['params']

    cache = _SMALL_CONFIG.init_cache(
        batch_size=batch_size,
        dtype=jnp.float32,
        cache_length=cache_length,
    )

    sampler_stability = _sampler.DiffusionSampler(
        model=model,
        end_tokens=(99,),
        forbidden_tokens=None,
        sampling=_sampling.Greedy(),
        cache_length=cache_length,
        special_tokens=None,
        diffusion_process=_sampler.DiffusionProcess(),
        logit_shaper=_sampler.AnnealingTemperatureShaperConfig().make(),
        sample_from_predictions=_sampler.SampleFromPredictions(
            text_vocab_size=vocab_size,
        ),
        canvas_length=canvas_length,
        max_denoising_steps=max_denoising_steps,
        early_stop_fn=_early_stopping.TokenStabilityEarlyStop(),
        text_vocab_size=vocab_size,
    )

    output1 = sampler_stability.sample_next_canvas(
        canvas_length=canvas_length,
        max_denoising_steps=max_denoising_steps,
        batch_size=batch_size,
        cache=cache,
        params=params,
        rng=rng,
    )
    output2 = sampler_stability.sample_next_canvas(
        canvas_length=canvas_length,
        max_denoising_steps=max_denoising_steps,
        batch_size=batch_size,
        cache=cache,
        params=params,
        rng=rng,
    )

    self.assertEqual(output1.shape, (batch_size, canvas_length))
    np.testing.assert_array_equal(output1, output2)


if __name__ == '__main__':
  absltest.main()