tests/models/diffusion_gemma/test_modeling_diffusion_gemma.py - main - githubmirror/transformers

tests/models/diffusion_gemma/test_modeling_diffusion_gemma.py
1 844 строки107 KB
Kashif Rasul
[DiffusionGemma] Cast the decoder padding mask to bool (#47295)
27 июл 2026, 12:53
27 июл 2026, 12:534ad6ba1
# Copyright 2026 the HuggingFace Team. All rights reserved.
#
# 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.
"""Testing suite for the PyTorch DiffusionGemma model."""

import unittest

from parameterized import parameterized

from transformers import is_torch_available
from transformers.testing_utils import (
    CaptureStdout,
    Expectations,
    cleanup,
    require_deterministic_for_xpu,
    require_torch,
    slow,
    tooslow,
    torch_device,
)

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor


if is_torch_available():
    import torch

    from transformers import (
        AutoConfig,
        AutoProcessor,
        AutoTokenizer,
        Cache,
        DiffusionGemmaConfig,
        DiffusionGemmaForBlockDiffusion,
        DiffusionGemmaGenerationConfig,
        DiffusionGemmaModel,
        DynamicCache,
        EntropyBoundSamplerConfig,
        StaticCache,
        TextDiffusionStreamer,
        set_seed,
    )


class DiffusionGemmaVisionText2TextModelTester:
    def __init__(
        self,
        parent,
        batch_size=3,
        text_config={
            # default text config test values
            "vocab_size": 99,
            "hidden_size": 32,
            "num_attention_heads": 2,
            "num_key_value_heads": 2,
            "intermediate_size": 32,
            "hidden_act": "gelu",
            "max_position_embeddings": 512,
            "initializer_range": 0.02,
            "pad_token_id": 0,
            "bos_token_id": 1,
            "eos_token_id": 2,
            "moe_intermediate_size": 8,
            "moe_num_shared_experts": 1,
            "num_experts": 4,
            "num_hidden_layers": 2,
            # model-specific text config values
            "layer_types": ["sliding_attention", "full_attention"],  # we want to test both types
            # full-attention layer overrides (head_dim is hidden_size // num_attention_heads)
            "per_layer_config": {1: {"head_dim": 32 // 2, "num_key_value_heads": 2}},
            "top_k_experts": 2,  # key introduced by the gemma4 family
            "use_bidirectional_attention": "vision",  # Test if bidirectional image mask path works
        },
        vision_config={
            "use_labels": True,
            "image_size": 20,
            "patch_size": 5,
            "num_channels": 3,
            "is_training": True,
            "hidden_size": 16,
            "num_key_value_heads": 1,
            "num_hidden_layers": 2,
            "num_attention_heads": 4,
            "intermediate_size": 37,
            "dropout": 0.1,
            "attention_dropout": 0.1,
            "initializer_range": 0.02,
        },
        mm_tokens_per_image=2,
        image_token_id=4,
        boi_token_id=5,
        eoi_token_id=6,
        seq_length=25,  # See: `test_generate_beyond_sliding_window` before changing this value!
        self_conditioning_size=16,
        canvas_length=16,
        is_training=False,
    ):
        self.parent = parent
        self.batch_size = batch_size
        # `image_token_id` is set to 0 to pass "resize_embeddings" test, do not modify
        self.mm_tokens_per_image = mm_tokens_per_image
        self.image_token_id = image_token_id
        self.boi_token_id = boi_token_id
        self.eoi_token_id = eoi_token_id
        self.text_config = text_config
        self.vision_config = vision_config
        self.seq_length = seq_length
        self.is_training = is_training
        self.self_conditioning_size = self_conditioning_size
        self.canvas_length = canvas_length

        self.pad_token_id = text_config["pad_token_id"]
        self.num_hidden_layers = text_config["num_hidden_layers"]
        self.vocab_size = text_config["vocab_size"]
        self.hidden_size = text_config["hidden_size"]
        self.num_attention_heads = text_config["num_attention_heads"]

        self.num_channels = vision_config["num_channels"]
        self.image_size = vision_config["image_size"]

    def get_config(self):
        return DiffusionGemmaConfig(
            text_config=self.text_config,
            vision_config=self.vision_config,
            image_token_id=self.image_token_id,
            boi_token_id=self.boi_token_id,
            eoi_token_id=self.eoi_token_id,
            mm_tokens_per_image=self.mm_tokens_per_image,
            self_conditioning_size=self.self_conditioning_size,
            canvas_length=self.canvas_length,
        )

    def prepare_config_and_inputs_for_common(self):
        # 1. config preparation
        config = self.get_config()
        config.vision_config.pooling_kernel_size = 2

        # 2. image inputs preparation
        # (num_images, max_num_patches, patch_size * patch_size * num_channels)
        patch_size = config.vision_config.patch_size
        pixel_values = floats_tensor(
            [
                self.batch_size,
                self.vision_config["image_size"],
                patch_size * patch_size * self.vision_config["num_channels"],
            ]
        )
        # (num_images, max_num_patches, 2) for height/width positions. Let it be all ones for testign
        pixel_position_ids = torch.ones(self.vision_config["image_size"], device=torch_device, dtype=torch.long)
        pixel_position_ids = pixel_position_ids[None, :, None].repeat(self.batch_size, 1, 2)

        # 3. text input preparation
        input_ids = ids_tensor([self.batch_size, self.seq_length], config.text_config.vocab_size - 1) + 1
        decoder_input_ids = ids_tensor([self.batch_size, self.canvas_length], config.text_config.vocab_size - 1) + 1
        attention_mask = input_ids.ne(self.pad_token_id).to(torch_device)

        # Ensure no tokens accidentally match special token IDs
        input_ids[input_ids == config.image_token_id] = self.pad_token_id
        input_ids[:, :1] = config.image_token_id

        mm_token_type_ids = torch.zeros_like(input_ids)
        mm_token_type_ids[input_ids == config.image_token_id] = 1

        # 4. build the input dict
        inputs_dict = {
            "pixel_values": pixel_values,
            "image_position_ids": pixel_position_ids,
            "input_ids": input_ids,
            "decoder_input_ids": decoder_input_ids,
            "attention_mask": attention_mask,
            "mm_token_type_ids": mm_token_type_ids,
        }
        return config, inputs_dict


@require_torch
class DiffusionGemmaVisionText2TextModelTest(ModelTesterMixin, unittest.TestCase):
    all_model_classes = (DiffusionGemmaModel, DiffusionGemmaForBlockDiffusion) if is_torch_available() else ()
    all_generative_model_classes = ()  # No class inherits `GenerationMixin`
    additional_model_inputs = ["mm_token_type_ids", "decoder_input_ids", "image_position_ids"]

    def setUp(self):
        self.model_tester = DiffusionGemmaVisionText2TextModelTester(self)
        self.config_tester = ConfigTester(self, config_class=DiffusionGemmaConfig, has_text_modality=False)

    # Tests adapted from `CausalLMModelTest` (when appropriate, as this is not an AR model). We can use the test
    # name to check the original equivalent that served as inspiration.
    def test_config(self):
        """Tests basic config properties"""
        self.config_tester.run_common_tests()

    def test_model(self):
        """Tests that we can run a forward pass with the base model"""
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaModel(config=config).to(torch_device).eval()
        expected_shape = (self.model_tester.batch_size, self.model_tester.canvas_length, self.model_tester.hidden_size)

        result = model(inputs_dict["input_ids"], attention_mask=inputs_dict["attention_mask"])
        self.assertEqual(result.last_hidden_state.shape, expected_shape)

        result = model(inputs_dict["input_ids"], decoder_input_ids=inputs_dict["decoder_input_ids"])
        self.assertEqual(result.last_hidden_state.shape, expected_shape)

        result = model(inputs_dict["input_ids"])
        self.assertEqual(result.last_hidden_state.shape, expected_shape)

    # Tests overwritten from `ModelTesterMixin`.
    @unittest.skip(reason="TODO")
    def test_attention_outputs(self):
        # TODO(joaogante): Write a custom test for this model. The base test can't be used because
        # the decoder has query len = canvas len and key/value shape = kv cache + canvas len.
        # On top of this, we should also check the encoder attentions.
        pass

    @unittest.skip(reason="TODO")
    def test_hidden_states_output(self):
        # TODO(joaogante): Same as `test_attention_outputs`
        pass

    @unittest.skip(reason="TODO")
    def test_can_init_all_missing_weights(self):
        # TODO(joaogante): failing on tied weights. Explore later.
        pass

    @unittest.skip(reason="TODO")
    def test_init_weights_can_init_buffers(self):
        # TODO(joaogante): probably related to `test_can_init_all_missing_weights`
        pass

    @unittest.skip(reason="TODO")
    def test_tp_plan_matches_params(self):
        # TODO(joaogante): explore TP config after other issues are sorted
        pass

    @unittest.skip(reason="Hard to specify `self.model_split_percents` due to tied weights. Skip for now.")
    def test_cpu_offload(self):
        pass

    @unittest.skip(reason="Hard to specify `self.model_split_percents` due to tied weights. Skip for now.")
    def test_disk_offload_bin(self):
        pass

    @unittest.skip(reason="Hard to specify `self.model_split_percents` due to tied weights. Skip for now.")
    def test_disk_offload_safetensors(self):
        pass

    def test_gradient_checkpointing_matches_no_checkpointing(self):
        """
        Gradient checkpointing relies on the encoder writing its KV cache outside the checkpointed layer calls and
        on the decoder keeping the read-only cache (`_can_checkpoint_with_cache`). Both must reproduce the exact
        same loss and gradients as the non-checkpointed forward, including encoder gradients through the cache.
        """
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        # Checkpointing only engages in train mode, so determinism requires disabling dropout instead of eval()
        config.text_config.attention_dropout = 0.0
        config.vision_config.dropout = 0.0
        config.vision_config.attention_dropout = 0.0
        model = DiffusionGemmaForBlockDiffusion(config).to(torch_device)
        inputs = self._prepare_for_class(inputs_dict, DiffusionGemmaForBlockDiffusion)

        def loss_and_grads(checkpointing):
            model.zero_grad(set_to_none=True)
            if checkpointing:
                model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
            else:
                model.gradient_checkpointing_disable()
            model.train()
            loss = model(**inputs).logits.float().mean()
            loss.backward()
            return loss.detach(), {name: p.grad.clone() for name, p in model.named_parameters() if p.grad is not None}

        loss_ref, grads_ref = loss_and_grads(checkpointing=False)
        loss_checkpointed, grads_checkpointed = loss_and_grads(checkpointing=True)

        # Snapshot/restore makes the recomputation bit-exact on CPU; accelerators can reorder accumulation in the
        # backward, so relax the tolerance there.
        rtol = atol = 0 if torch_device == "cpu" else 1e-5
        torch.testing.assert_close(loss_checkpointed, loss_ref, rtol=rtol, atol=atol)
        self.assertEqual(grads_checkpointed.keys(), grads_ref.keys())
        self.assertTrue(any("encoder" in name for name in grads_checkpointed))
        for name in grads_ref:
            torch.testing.assert_close(grads_checkpointed[name], grads_ref[name], rtol=rtol, atol=atol)

    @unittest.skip(reason="Hard to specify `self.model_split_percents` due to tied weights. Skip for now.")
    def test_model_parallelism(self):
        pass

    # Tests designed specifically for `DiffusionGemma`, excluding generation tests.
    def test_tied_weights(self):
        """
        Tests that randomly initialized models have their weights properly tied.

        There are two sets of ties in DiffusionGemma:
        1. The text encoder trainable params are the same as the decoder trainable params, if we exclude the
           decoder self-conditioning weights
        2. The LM Head weights are the same as the embedding weights -- in both the text encoder and in the decoder!
        """
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # 1. text encoder trainable params == decoder trainable params (excluding decoder self-conditioning)
        # In other words, all text encoder trainable params must be decoder trainable params.
        decoder_trainable_params = dict(model.model.decoder.named_parameters())
        for name, encoder_text_param in model.model.encoder.language_model.named_parameters():
            decoder_trainable_param = decoder_trainable_params[name]
            self.assertTrue(encoder_text_param is decoder_trainable_param)

        # 2. lm head weights == text encoder embedding weights == decoder embedding weights
        self.assertTrue(model.lm_head.weight is model.model.decoder.embed_tokens.weight)
        self.assertTrue(model.lm_head.weight is model.model.encoder.language_model.embed_tokens.weight)

    def test_use_cache_is_ignored(self):
        """
        DiffusionGemma's decoder never writes a cache, so the common kwarg `use_cache` is silently ignored instead
        of raising. This keeps the model usable with gradient checkpointing, which forces `use_cache=False`.
        """
        config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # Either value is accepted and has no effect: no exception, and the output matches omitting the kwarg.
        expected = model(**model_inputs).logits
        for use_cache in (False, True):
            logits = model(**model_inputs, use_cache=use_cache).logits
            torch.testing.assert_close(logits, expected)

    def test_diffusion_decoder_mask_no_cache_raises_exception(self):
        """
        DiffusionGemma has a custom function to create an attention mask for the decoder. Contrarily
        to other mask creation functions, it requires a KV cache
        """
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        dummy_canvas = torch.ones((2, 4), dtype=torch.int32, device=torch_device)
        with self.assertRaises(ValueError):
            _ = model.model.decoder.create_diffusion_decoder_attention_mask(
                config=config.text_config,
                inputs_embeds=dummy_canvas.unsqueeze(-1),
                past_key_values=None,  # this will trigger an exception
            )

    def test_diffusion_decoder_mask_dynamic_cache(self):
        """
        This is the simplest test: the resulting mask should be full of 1s. As a shortcut, in eager execution
        + dynamic cache, it simply returns `None` attention masks
        """

        prefill_length = 8
        canvas_length = 4
        concat_kv_length = prefill_length + canvas_length
        batch_size = 2

        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # Apply prefill (smaller than sliding window)
        past_key_values = DynamicCache(config=config)
        prefill_input_ids = torch.ones((batch_size, prefill_length), dtype=torch.int32, device=torch_device) * 50
        past_key_values = model.model.encoder(prefill_input_ids, past_key_values=past_key_values).past_key_values

        # Get the mask
        decoder_attention_mask = torch.ones((batch_size, concat_kv_length), dtype=torch.bool, device=torch_device)
        dummy_canvas = torch.ones((batch_size, canvas_length), dtype=torch.int32, device=torch_device)
        mask_mapping = model.model.decoder.create_diffusion_decoder_attention_mask(
            config=config.text_config,
            inputs_embeds=dummy_canvas.unsqueeze(-1),
            past_key_values=past_key_values,
            decoder_attention_mask=decoder_attention_mask,
        )
        self.assertIsNone(mask_mapping["full_attention"])
        self.assertIsNone(mask_mapping["sliding_attention"])

        # We get the same result if we pass `attention_mask = None`
        mask_mapping = model.model.decoder.create_diffusion_decoder_attention_mask(
            config=config.text_config,
            inputs_embeds=dummy_canvas.unsqueeze(-1),
            past_key_values=past_key_values,
            decoder_attention_mask=None,
        )
        self.assertIsNone(mask_mapping["full_attention"])
        self.assertIsNone(mask_mapping["sliding_attention"])

    def test_diffusion_decoder_mask_dynamic_cache_left_padding(self):
        """
        This test is similar to `test_diffusion_decoder_mask_dynamic_cache`, but with left-padding on one of the
        rows. There should be a few zeros on the final mask, corresponding to the left-padded items.
        """

        prefill_length = 8
        canvas_length = 4
        concat_kv_length = prefill_length + canvas_length
        batch_size = 2
        left_padding_length = 2  # only applied on batch item 0
        expected_non_zero = (
            ((concat_kv_length - left_padding_length) * canvas_length)  # batch item 0
            + (concat_kv_length * canvas_length)  # batch item 1
        )
        expected_attention_mask_shape = (batch_size, 1, canvas_length, concat_kv_length)

        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # Apply prefill (smaller than sliding window)
        past_key_values = DynamicCache(config=config)
        prefill_input_ids = torch.ones((batch_size, prefill_length), dtype=torch.int32, device=torch_device) * 50
        past_key_values = model.model.encoder(prefill_input_ids, past_key_values=past_key_values).past_key_values

        # Get the mask
        decoder_attention_mask = torch.ones((batch_size, concat_kv_length), dtype=torch.bool, device=torch_device)
        decoder_attention_mask[0, :left_padding_length] = 0
        dummy_canvas = torch.ones((batch_size, canvas_length), dtype=torch.int32, device=torch_device)
        mask_mapping = model.model.decoder.create_diffusion_decoder_attention_mask(
            config=config.text_config,
            inputs_embeds=dummy_canvas.unsqueeze(-1),
            past_key_values=past_key_values,
            decoder_attention_mask=decoder_attention_mask,
        )
        self.assertEqual(mask_mapping["full_attention"].shape, expected_attention_mask_shape)
        self.assertEqual(mask_mapping["full_attention"].sum(), expected_non_zero)
        self.assertEqual(mask_mapping["full_attention"][0, 0, :, :left_padding_length].sum(), 0)
        self.assertEqual(mask_mapping["sliding_attention"].shape, expected_attention_mask_shape)
        self.assertEqual(mask_mapping["sliding_attention"].sum(), expected_non_zero)
        self.assertEqual(mask_mapping["sliding_attention"][0, 0, :, :left_padding_length].sum(), 0)

    def test_diffusion_decoder_mask_dynamic_cache_int_padding_mask(self):
        """
        Same as `test_diffusion_decoder_mask_dynamic_cache_left_padding`, but the padding mask holds `1`/`0` integers,
        the way a tokenizer returns it, instead of booleans. It must produce the same boolean mask.
        """

        prefill_length = 8
        canvas_length = 4
        concat_kv_length = prefill_length + canvas_length
        batch_size = 2
        left_padding_length = 2  # only applied on batch item 0
        expected_non_zero = (
            ((concat_kv_length - left_padding_length) * canvas_length)  # batch item 0
            + (concat_kv_length * canvas_length)  # batch item 1
        )
        expected_attention_mask_shape = (batch_size, 1, canvas_length, concat_kv_length)

        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # Apply prefill (smaller than sliding window)
        past_key_values = DynamicCache(config=config)
        prefill_input_ids = torch.ones((batch_size, prefill_length), dtype=torch.int32, device=torch_device) * 50
        past_key_values = model.model.encoder(prefill_input_ids, past_key_values=past_key_values).past_key_values

        # Get the mask
        decoder_attention_mask = torch.ones((batch_size, concat_kv_length), dtype=torch.long, device=torch_device)
        decoder_attention_mask[0, :left_padding_length] = 0
        dummy_canvas = torch.ones((batch_size, canvas_length), dtype=torch.int32, device=torch_device)
        mask_mapping = model.model.decoder.create_diffusion_decoder_attention_mask(
            config=config.text_config,
            inputs_embeds=dummy_canvas.unsqueeze(-1),
            past_key_values=past_key_values,
            decoder_attention_mask=decoder_attention_mask,
        )
        for layer_pattern in ("full_attention", "sliding_attention"):
            self.assertEqual(mask_mapping[layer_pattern].dtype, torch.bool)
            self.assertEqual(mask_mapping[layer_pattern].shape, expected_attention_mask_shape)
            self.assertEqual(mask_mapping[layer_pattern].sum(), expected_non_zero)
            self.assertEqual(mask_mapping[layer_pattern][0, 0, :, :left_padding_length].sum(), 0)

    def test_diffusion_decoder_mask_dynamic_cache_beyond_sliding_window(self):
        """
        This tests builds upon `test_diffusion_decoder_mask_dynamic_cache_left_padding`: it tests the case with
        left-padding AND going beyond the sliding window. The prefill length is longer than the sliding window
        length. The short left-padding should become invisible in the sliding window layers, after prefill.
        """

        prefill_length = 16
        sliding_window_length = 8
        canvas_length = 4
        concat_kv_length_full = prefill_length + canvas_length
        # -1 -> the DYNAMIC sliding window kv cache has len=window-1
        concat_kv_length_sliding = sliding_window_length + canvas_length - 1
        batch_size = 2
        left_padding_length = 2  # only applied on batch item 0
        expected_non_zero_full = (
            ((concat_kv_length_full - left_padding_length) * canvas_length)  # batch item 0
            + (concat_kv_length_full * canvas_length)  # batch item 1
        )
        expected_attention_mask_shape_full = (batch_size, 1, canvas_length, concat_kv_length_full)
        expected_non_zero_sliding = concat_kv_length_sliding * batch_size * canvas_length
        expected_attention_mask_shape_sliding = (
            batch_size,
            1,
            canvas_length,
            sliding_window_length + canvas_length - 1,
        )
        # Double-check test assumption that left-padding should be past the sliding window
        self.assertTrue(prefill_length - left_padding_length > sliding_window_length)

        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
        config.text_config.sliding_window = sliding_window_length
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # Apply prefill (LARGER than sliding window)
        past_key_values = DynamicCache(config=config)
        prefill_input_ids = torch.ones((batch_size, prefill_length), dtype=torch.int32, device=torch_device) * 50
        past_key_values = model.model.encoder(prefill_input_ids, past_key_values=past_key_values).past_key_values

        # Get the mask
        decoder_attention_mask = torch.ones((batch_size, concat_kv_length_full), dtype=torch.bool, device=torch_device)
        decoder_attention_mask[0, :left_padding_length] = 0
        dummy_canvas = torch.ones((batch_size, canvas_length), dtype=torch.int32, device=torch_device)
        mask_mapping = model.model.decoder.create_diffusion_decoder_attention_mask(
            config=config.text_config,
            inputs_embeds=dummy_canvas.unsqueeze(-1),
            past_key_values=past_key_values,
            decoder_attention_mask=decoder_attention_mask,
        )
        self.assertEqual(mask_mapping["full_attention"].shape, expected_attention_mask_shape_full)
        self.assertEqual(mask_mapping["full_attention"].sum(), expected_non_zero_full)
        self.assertEqual(mask_mapping["full_attention"][0, 0, :, :left_padding_length].sum(), 0)
        self.assertEqual(mask_mapping["sliding_attention"].shape, expected_attention_mask_shape_sliding)
        self.assertEqual(mask_mapping["sliding_attention"].sum(), expected_non_zero_sliding)
        # sliding window got beyond padding, so the first tokens in the sliding window will not be masked
        self.assertEqual(
            mask_mapping["sliding_attention"][0, 0, :, :left_padding_length].sum(), left_padding_length * canvas_length
        )

    def test_diffusion_decoder_mask_static_cache(self):
        """
        Same as `test_diffusion_decoder_mask_dynamic_cache`, but with a Static Cache. Contrarily to the
        original test, the mask is materialized (a materialized mask is needed at compile time)
        """

        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        prefill_length = 8
        canvas_length = config.canvas_length
        total_input_length = prefill_length + canvas_length

        # More than input so we don't hit sliding window being cropped yet
        static_cache_length = total_input_length + 20
        concat_kv_length = static_cache_length + canvas_length
        batch_size = 2
        expected_non_zero = total_input_length * canvas_length * batch_size
        expected_attention_mask_shape = (batch_size, 1, canvas_length, concat_kv_length)

        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # Apply prefill (smaller than sliding window)
        past_key_values = StaticCache(config=config, max_cache_len=static_cache_length)
        prefill_input_ids = torch.ones((batch_size, prefill_length), dtype=torch.int32, device=torch_device) * 50
        past_key_values = model.model.encoder(prefill_input_ids, past_key_values=past_key_values).past_key_values

        # Get the mask (correctly designed for the static cache)
        decoder_attention_mask = torch.ones((batch_size, total_input_length), dtype=torch.bool, device=torch_device)
        dummy_canvas = torch.ones((batch_size, canvas_length), dtype=torch.int32, device=torch_device)
        mask_mapping = model.model.decoder.create_diffusion_decoder_attention_mask(
            config=config,
            inputs_embeds=dummy_canvas.unsqueeze(-1),
            past_key_values=past_key_values,
            decoder_attention_mask=decoder_attention_mask,
        )
        self.assertEqual(mask_mapping["full_attention"].shape, expected_attention_mask_shape)
        self.assertEqual(mask_mapping["full_attention"].sum(), expected_non_zero)
        self.assertEqual(mask_mapping["sliding_attention"].shape, expected_attention_mask_shape)
        self.assertEqual(mask_mapping["sliding_attention"].sum(), expected_non_zero)

    def test_diffusion_decoder_mask_static_cache_beyond_sliding_window(self):
        """
        Same as `test_diffusion_decoder_mask_dynamic_cache_beyond_sliding_window`, but with a Static Cache.
        This is the most complex mask preparation test, including static caches, left-padding, and mask preparation
        beyond the sliding window length.
        """
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        prefill_length = 16
        sliding_window_length = 8
        canvas_length = config.canvas_length
        total_input_length = prefill_length + canvas_length

        # Max-cache-len is less than `total_input_length` to trigger going beyond sliding window
        concat_kv_length_sliding = sliding_window_length + canvas_length - 1
        static_cache_length = total_input_length - 4
        concat_kv_length_full = static_cache_length + canvas_length
        batch_size = 2
        left_padding_length = 2  # only applied on batch item 0

        # Current query (`canvas`) attends to key (`prefill + itself`). Optionally, doesn't attend if pading present
        expected_non_zero_full = (
            ((total_input_length - left_padding_length) * canvas_length)  # batch item 0
            + (total_input_length * canvas_length)  # batch item 1
        )
        expected_attention_mask_shape_full = (batch_size, 1, canvas_length, concat_kv_length_full)
        expected_non_zero_sliding = concat_kv_length_sliding * batch_size * canvas_length
        expected_attention_mask_shape_sliding = (
            batch_size,
            1,
            canvas_length,
            concat_kv_length_sliding + 1,
        )
        # Double-check test assumption that left-padding should be past the sliding window
        self.assertTrue(prefill_length - left_padding_length > sliding_window_length)

        config.text_config.sliding_window = sliding_window_length
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # Apply prefill (LARGER than sliding window)
        past_key_values = StaticCache(config=config, max_cache_len=static_cache_length)
        prefill_input_ids = torch.ones((batch_size, prefill_length), dtype=torch.int32, device=torch_device) * 50
        past_key_values = model.model.encoder(prefill_input_ids, past_key_values=past_key_values).past_key_values

        # Get the mask (correctly designed for the static cache)
        decoder_attention_mask = torch.ones((batch_size, total_input_length), dtype=torch.bool, device=torch_device)
        decoder_attention_mask[0, :left_padding_length] = 0
        dummy_canvas = torch.ones((batch_size, canvas_length), dtype=torch.int32, device=torch_device)
        mask_mapping = model.model.decoder.create_diffusion_decoder_attention_mask(
            config=config,
            inputs_embeds=dummy_canvas.unsqueeze(-1),
            past_key_values=past_key_values,
            decoder_attention_mask=decoder_attention_mask,
        )
        self.assertEqual(mask_mapping["full_attention"].shape, expected_attention_mask_shape_full)
        self.assertEqual(mask_mapping["full_attention"].sum(), expected_non_zero_full)
        self.assertEqual(mask_mapping["full_attention"][0, 0, :, :left_padding_length].sum(), 0)
        self.assertEqual(mask_mapping["sliding_attention"].shape, expected_attention_mask_shape_sliding)
        self.assertEqual(mask_mapping["sliding_attention"].sum(), expected_non_zero_sliding)
        # sliding window got beyond padding, so the first tokens in the sliding window will not be masked
        self.assertEqual(
            mask_mapping["sliding_attention"][0, 0, :, :left_padding_length].sum(), left_padding_length * canvas_length
        )

    # Generation tests
    @parameterized.expand(
        [
            # (sub_test_name, {generate_kwargs}),
            # Cache interface
            ("dynamic_cache", {}),  # default cache
            ("static_cache", {"cache_implementation": "static"}),  # static cache -> triggers compilation
            # Diffusion-specific common kwargs
            ("custom_sampler", {"sampler_config": EntropyBoundSamplerConfig(0.1)}),
            (
                "custom_temperature",
                {"t_min": 0.1, "t_max": 0.2},
            ),
            (
                "custom_stopping",
                {"stability_threshold": 0, "confidence_threshold": 1e-1},
            ),
        ]
    )
    def test_generate(self, name, generate_kwargs):
        """Tests `generate` calls with common flags"""
        config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()
        model.generation_config.eos_token_id = None  # force generation up to `max_new_tokens`

        generation_outputs = model.generate(
            **model_inputs, max_new_tokens=16, max_denoising_steps=2, **generate_kwargs
        )

        expected_shape = (model_inputs["input_ids"].shape[0], model_inputs["input_ids"].shape[1] + 16)
        self.assertEqual(generation_outputs.sequences.shape, expected_shape)

    def test_generate_text_only(self):
        """Same as `test_generate`, but only with text inputs (no images)"""
        config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()
        model.generation_config.eos_token_id = None  # force generation up to `max_new_tokens`

        generation_outputs = model.generate(model_inputs["input_ids"], max_new_tokens=16, max_denoising_steps=2)

        expected_shape = (model_inputs["input_ids"].shape[0], model_inputs["input_ids"].shape[1] + 16)
        self.assertEqual(generation_outputs.sequences.shape, expected_shape)

    def test_generate_without_return_dict(self):
        """Same as `test_generate_text_only`, but return_dict_in_generate=False"""
        config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()
        model.generation_config.eos_token_id = None  # force generation up to `max_new_tokens`
        model.generation_config.return_dict_in_generate = False

        generation_outputs = model.generate(model_inputs["input_ids"], max_new_tokens=16, max_denoising_steps=2)

        expected_shape = (model_inputs["input_ids"].shape[0], model_inputs["input_ids"].shape[1] + 16)
        self.assertEqual(generation_outputs.shape, expected_shape)

    def test_generate_from_generation_config(self):
        """
        Same as the base case in `test_generate`, but we parameterize the generation call with a
        `DiffusionGemmaGenerationConfig`
        """
        config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # force `eos_token_id = None` to not stop before reaching `max_new_tokens`
        generation_config = DiffusionGemmaGenerationConfig(eos_token_id=None, max_new_tokens=16, max_denoising_steps=2)
        generation_outputs = model.generate(**model_inputs, generation_config=generation_config)

        expected_shape = (model_inputs["input_ids"].shape[0], model_inputs["input_ids"].shape[1] + 16)
        self.assertEqual(generation_outputs.sequences.shape, expected_shape)

    def test_generate_kwarg_overrides(self):
        """
        Same as `test_generate_from_generation_config`, but we override some parameters with kwargs in the
        `generate` call. Remember: kwargs > generation config
        """
        config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # force `eos_token_id = None` to not stop before reaching `max_new_tokens`
        generation_config = DiffusionGemmaGenerationConfig(eos_token_id=None, max_new_tokens=16, max_denoising_steps=2)
        generation_outputs = model.generate(**model_inputs, generation_config=generation_config, max_new_tokens=32)

        expected_shape = (model_inputs["input_ids"].shape[0], model_inputs["input_ids"].shape[1] + 32)
        self.assertEqual(generation_outputs.sequences.shape, expected_shape)

    def test_generate_with_past_key_values(self):
        """
        Tests that we can pass in past_key_values outputted from a previous `generate` into a subsequent call.
        This pattern is used in chat sessions.
        """
        config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()
        model.generation_config.eos_token_id = None  # force generation up to `max_new_tokens`

        # 1st generate call, without KV cache
        generation_1_outputs = model.generate(**model_inputs, max_new_tokens=16, max_denoising_steps=2)
        self.assertIsInstance(generation_1_outputs.past_key_values, Cache)

        expected_shape = (model_inputs["input_ids"].shape[0], model_inputs["input_ids"].shape[1] + 16)
        self.assertEqual(generation_1_outputs.sequences.shape, expected_shape)

        # 2nd generate call, with KV cache from previous call
        # NOTE: `input_ids` in `generate` can only contain **unprocessed** tokens
        new_input_ids = torch.cat((generation_1_outputs.sequences, model_inputs["input_ids"]), dim=-1)
        generation_2_outputs = model.generate(
            input_ids=new_input_ids,
            past_key_values=generation_1_outputs.past_key_values,
            max_new_tokens=16,
            max_diffusion_steps=2,
        )

        # in total, we feed in `model_inputs["input_ids"]` twice, and get 16 new tokens twice
        expected_shape = (model_inputs["input_ids"].shape[0], (model_inputs["input_ids"].shape[1] + 16) * 2)
        self.assertEqual(generation_2_outputs.sequences.shape, expected_shape)

    @parameterized.expand(
        [
            ("dynamic_cache", None),  # default cache
            ("static_cache", "static"),  # static cache -> triggers compilation under the hood
        ]
    )
    @slow
    def test_generate_beyond_sliding_window(self, name, cache_implementation):
        """Tests that generate can run beyond the sliding window length"""
        config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        # Canvas length has to be much smaller than sliding window to correctly check the mask
        # and input-length has to be smaller than sliding length! We will check three cases:
        # 1) input/mask is smaller than window
        # 2) input length approaching full cache length
        # 3) input length goes beyond window length
        # Input length is 25 in model tester, so we will set sliding window to 48 tokens
        # After generating `max_new_tokens=32`, the final length will go beyond sliding window
        config.text_config.sliding_window = 48
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()
        model.generation_config.eos_token_id = None  # force generation up to `max_new_tokens`

        generation_outputs = model.generate(
            **model_inputs,
            max_new_tokens=32,
            max_denoising_steps=2,
            disable_compile=True,
            cache_implementation=cache_implementation,
        )
        self.assertTrue(generation_outputs.sequences.shape[1] > config.text_config.sliding_window)

        expected_shape = (model_inputs["input_ids"].shape[0], model_inputs["input_ids"].shape[1] + 32)
        self.assertEqual(generation_outputs.sequences.shape, expected_shape)

    @unittest.skip(reason="TODO(joaogante): red in CI, fix me")
    def test_diffusion_streaming(self):
        """Tests `TextDiffusionStreamer`"""
        # TODO(huggingface team): after a tiny diffusion model is created, this test can be moved into
        # its right place, i.e. tests/generation/test_streamers.py

        config, model_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        input_ids = model_inputs["input_ids"][:1, :]  # streaming requires bsz=1
        model = DiffusionGemmaForBlockDiffusion(config=config).to(torch_device).eval()

        # Any tokenizer will do, as long as it has the same vocab size.
        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")

        with CaptureStdout() as cs:
            streamer = TextDiffusionStreamer(tokenizer)
            model_out = model.generate(input_ids, max_new_tokens=16, max_denoising_steps=2, streamer=streamer)
        streamer_text = cs.out
        generate_text = tokenizer.decode(model_out.sequences[0])

        # the streamer ends with a `\n`
        self.assertTrue(streamer_text.endswith("\n"))
        # Important for this test: max_new_tokens==canvas_length. The last piece of text thrown by the streamer
        # is the finalized canvas.
        self.assertEqual(streamer_text[-(len(generate_text) + 1) : -1], generate_text)
        # `\x1b7` = save cursor position; `\x1b8` = restore cursor position;
        # Used to overwrite draft tokens while streaming.
        self.assertIn("\x1b7", streamer_text)
        self.assertIn("\x1b8", streamer_text)


@require_torch
class DiffusionGemmaIntegrationTest(unittest.TestCase):
    _model_path = "google/diffusiongemma-26B-A4B-it"

    def setup(self):
        cleanup(torch_device, gc_collect=True)

    def tearDown(self):
        cleanup(torch_device, gc_collect=True)

    @slow
    def test_diffusion_gemma_chat_template(self):
        # fmt: off
        expected_input_ids = [
            [
                2,    105,   2364,    107,   6974,    496,   1440,  20494,   1003,
                23613, 236761,    106,    107,    105,   4368,    107,    100,  45518,
                107,    101
            ]
        ]
        # fmt: on
        expected_decoded_tokens = [
            "<bos><|turn>user\nWrite a long essay about Portugal.<turn|>\n<|turn>model\n<|channel>thought\n<channel|>"
        ]

        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            {"role": "user", "content": "Write a long essay about Portugal."},
        ]
        input_ids = processor.apply_chat_template(chat, tokenize=True, return_tensors="pt", add_generation_prompt=True)
        self.assertEqual(input_ids.tolist(), expected_input_ids)

        decoded_tokens = processor.decode(input_ids)
        self.assertIn("<|turn>", decoded_tokens[0])
        self.assertIn("<turn|>", decoded_tokens[0])
        self.assertIn("<|channel>", decoded_tokens[0])
        self.assertIn("<channel|>", decoded_tokens[0])
        self.assertNotIn("<|think|>", decoded_tokens[0])  # shouldn't have the think token
        self.assertEqual(decoded_tokens, expected_decoded_tokens)

    @slow
    def test_diffusion_gemma_chat_template_image(self):
        image_tokens = [255999, 258880, 258882]  # These tokens must be present in the `input_ids`
        # TODO(joao): this should be 280! Something is wrong with processing?
        image_token_count = 256

        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "url": "https://raw.githubusercontent.com/google-gemma/cookbook/refs/heads/main/apps/sample-data/GoldenGate.png",
                    },
                    {"type": "text", "text": "What is shown in this image?"},
                ],
            }
        ]
        model_inputs = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, return_dict=True
        )
        for token in image_tokens:
            self.assertIn(token, model_inputs["input_ids"])
        self.assertTrue((model_inputs["input_ids"] == 258880).sum() == image_token_count)

        for expected_model_input in ("attention_mask", "pixel_values", "image_position_ids", "mm_token_type_ids"):
            self.assertIn(expected_model_input, model_inputs)

    @slow
    def test_diffusion_gemma_chat_template_with_thinking(self):
        # fmt: off
        expected_input_ids = [
            [
                2,    105,   9731,    107,     98,    107,    106,    107,    105,
                2364,    107,   6974,    496,   1440,  20494,   1003,  23613, 236761,
                106,    107,    105,   4368,    107
            ]
        ]
        # fmt: on
        expected_decoded_tokens = [
            "<bos><|turn>system\n<|think|>\n<turn|>\n<|turn>user\nWrite a long essay about Portugal.<turn|>\n<|turn>model\n"
        ]

        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            {"role": "user", "content": "Write a long essay about Portugal."},
        ]
        input_ids = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, enable_thinking=True
        )
        self.assertEqual(input_ids.tolist(), expected_input_ids)

        decoded_tokens = processor.decode(input_ids)
        self.assertIn("<|turn>", decoded_tokens[0])
        self.assertIn("<turn|>", decoded_tokens[0])
        self.assertIn("<|think|>", decoded_tokens[0])  # must have the think token
        self.assertEqual(decoded_tokens, expected_decoded_tokens)

    # Only 1 checkpoint got released: the 26B MoE. At release time, MoE quantization was not working properly,
    # and the HF CI didn't have 80GB GPUs. As such, tests against the original model are tagged with `@tooslow`.
    # Meaning: these tests won't run in the `py.test` command, and you have to uncomment the line. These tests
    # exist to keep track of the model's original numerics.
    def _load_model(self, minified=False):
        if minified:
            config = AutoConfig.from_pretrained(self._model_path)
            config.text_config.num_hidden_layers = 6  # loaded model requires ~12.7 GB
            config.text_config.layer_types = config.text_config.layer_types[:6]
            model = DiffusionGemmaForBlockDiffusion.from_pretrained(
                self._model_path,
                config=config,
                dtype=torch.bfloat16,
                device_map="auto",
            )
        else:
            model = DiffusionGemmaForBlockDiffusion.from_pretrained(
                self._model_path,
                dtype=torch.bfloat16,
                device_map="auto",
            )
        return model

    @tooslow
    def test_diffusion_gemma_forward_text_only(self):
        """bsz=1, no image"""
        # fmt: off
        # print(model_out.logits[:, :10, :12])
        # Printed after running `torch.set_printoptions(precision=8, threshold=200, linewidth=100)`
        expected_logits = Expectations(
            {
                ("cuda", None): [
                    [
                        [ 4.12986183, 20.66659164,  6.75712347,  4.16051483, 13.11566544,  3.59206009,  3.46879172,
                        3.86895919,  3.14465618,  2.47871161,  4.43598938,  4.77164030],
                        [ 3.20645595, 22.52679825,  7.78647614,  2.71133161, 14.73279667,  4.55818033,  4.16051483,
                        3.03644228,  3.91504526,  1.58446193,  4.52764702,  4.46655130],
                        [ 4.77164030, 23.55075455,  5.80103159,  3.65364885, 15.66178513,  4.19115925,  5.01498747,
                        2.92814803,  4.55818033,  2.92814803,  5.28794241,  7.87383986],
                        [ 4.64972258, 24.27903175,  6.01130199,  3.66904116, 15.57065392,  3.17555952,  5.25765753,
                        3.40711212,  4.95421267,  4.28303957,  5.71073294,  8.10610390],
                        [ 5.16673803, 24.27903175,  6.42998266,  3.36083388, 16.89584923,  4.49710369,  5.77094412,
                        4.40541792,  5.22736216,  4.12986183,  7.05306482,  8.68214703],
                        [ 5.28794241, 24.36471176,  7.58207989,  3.53044105, 17.31759644,  4.49710369,  5.31821585,
                        4.40541792,  4.98460531,  3.88432312,  7.46494198,  8.33732700],
                        [ 5.10607004, 24.27903175,  7.17103910,  3.34540391, 16.55123901,  3.96111226,  4.09919930,
                        4.25242186,  4.34424734,  4.46655130,  5.71073294,  8.27962112],
                        [ 6.81642532, 24.19219208,  8.56747913,  5.37872887, 17.15010071,  5.80103159,  5.86117029,
                        4.25242186,  6.34047985,  4.49710369,  7.61132526,  9.36425304],
                        [ 6.99399042, 24.44923782,  8.96760273,  5.34847832, 16.81029892,  5.77094412,  4.95421267,
                        4.49710369,  5.68060923,  4.68021679,  6.93485880,  8.85363007],
                        [ 7.93200207, 23.83295822,  7.81561375,  6.51936626, 17.72932053,  5.98130083,  5.83110666,
                        5.43919706,  6.60862780,  5.28794241,  7.66977119,  9.86893463]
                    ]
                ]
            }
        )
        # fmt: on

        model = self._load_model()
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            {"role": "user", "content": "Write a long essay about Portugal."},
        ]
        model_inputs = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, return_dict=True
        ).to(model.device)

        # hardcodes input canvas to avoid depending on random canvas sampling
        decoder_input_ids = torch.arange(model.config.canvas_length, dtype=torch.int32, device=model.device).unsqueeze(
            0
        )
        model_inputs["decoder_input_ids"] = decoder_input_ids

        with torch.no_grad():
            model_out = model(**model_inputs)
        self.assertEqual(
            list(model_out.logits.shape), [1, model.config.canvas_length, model.config.text_config.vocab_size]
        )
        torch.testing.assert_close(
            model_out.logits[:, :10, :12].cpu(), torch.tensor(expected_logits.get_expectation())
        )

    @tooslow
    def test_diffusion_gemma_forward_with_image(self):
        """bsz=1, with image"""
        # fmt: off
        # print(model_out.logits[:, :10, :12])
        # Printed after running `torch.set_printoptions(precision=8, threshold=200, linewidth=100)`
        expected_logits = Expectations(
            {
                ("cuda", None): [
                    [
                        [ 6.93485880, 20.60072327,  3.89968514,  7.49424887,  6.99399042, -0.03344725,  0.44332713,
                        -0.53900456,  2.18363142,  0.35545215,  1.24927723,  3.53044105],
                        [ 6.81642532, 19.92110252,  4.92381001,  5.74084425, 10.31235313,  4.83253860,  0.94890219,
                        1.05425322,  0.87475199,  0.16503741,  0.06933582,  2.74232340],
                        [ 6.40016174, 22.63516045,  2.88171268,  6.75712347, 12.08642769,  5.25765753,  2.60281467,
                        2.60281467,  5.55999660,  2.66483307,  2.05925679,  5.98130083],
                        [ 9.86893463, 21.78874207,  2.35452294,  7.78647614, 15.01560497,  7.11208105,  5.07571983,
                        5.80103159,  4.68021679,  4.86297274,  3.45337439,  7.78647614],
                        [ 5.01498747, 21.18004990,  1.31166339,  4.95421267, 11.13002872,  2.54077387,  0.40427244,
                        2.40110350,  0.42770541,  1.23367894, -0.29100651,  1.06205606],
                        [ 5.74084425, 21.42788887,  2.12145352,  4.77164030, 11.82355690,  3.12920213,  3.66904116,
                        3.43795562,  2.64933062,  2.88171268,  2.71133161,  4.89339685],
                        [ 6.28074551, 22.30597687,  4.16051483,  4.52764702, 12.08642769,  4.64972258,  4.31364775,
                        5.89122152,  1.52212954,  3.74597287,  2.13699985,  6.31061935],
                        [ 6.37032795, 23.92457008,  9.02448273,  3.28366685, 10.64161205,  3.22190189,  3.00550938,
                        4.83253860,  0.88255781,  1.63899159,  1.81808197,  5.49961996],
                        [ 7.46494198, 25.78687286,  9.36425304,  3.63825440, 13.06506538,  2.02815175,  3.23734570,
                        5.37872887,  3.02097654,  2.27686334,  2.16808867,  6.63835430],
                        [ 6.16112089, 26.16215897,  5.71073294,  3.34540391, 11.92896748,  1.55329728,  4.77164030,
                        4.12986183,  1.02304065,  3.49961996,  1.73244548,  6.45979071]
                ]
                ]
            }
        )
        # fmt: on

        model = self._load_model()
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "url": "https://raw.githubusercontent.com/google-gemma/cookbook/refs/heads/main/apps/sample-data/GoldenGate.png",
                    },
                    {"type": "text", "text": "What is shown in this image?"},
                ],
            }
        ]
        model_inputs = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, return_dict=True
        ).to(model.device)

        # hardcodes input canvas to avoid depending on random canvas sampling
        decoder_input_ids = torch.arange(model.config.canvas_length, dtype=torch.int32, device=model.device).unsqueeze(
            0
        )
        model_inputs["decoder_input_ids"] = decoder_input_ids

        # sanity check: has image inputs
        self.assertIn("pixel_values", model_inputs)

        with torch.no_grad():
            model_out = model(**model_inputs)
        self.assertEqual(
            list(model_out.logits.shape), [1, model.config.canvas_length, model.config.text_config.vocab_size]
        )
        torch.testing.assert_close(
            model_out.logits[:, :10, :12].cpu(), torch.tensor(expected_logits.get_expectation())
        )

    @tooslow
    def test_diffusion_gemma_forward_batched(self):
        """bsz>1, no image"""
        # fmt: off
        # print(model_out.logits[:, :5, :12])
        # Printed after running `torch.set_printoptions(precision=8, threshold=200, linewidth=100)`
        expected_logits = Expectations(
            {
                ("cuda", None): [
                    [
                        [-1.56108880e+00,  1.99211025e+01,  1.60004282e+00,  1.01073846e-01,  7.58207989e+00,
                        -5.98130083e+00, -4.40541792e+00, -1.73244548e+00, -1.42860627e+00, -5.97577214e-01,
                        -1.81808197e+00,  8.74021053e-02],
                        [-1.02304065e+00,  2.26888294e+01,  9.53315258e+00, -3.37876350e-01,  8.91065121e+00,
                        1.29394531e-02,  5.35099506e-01, -2.52526045e+00, -7.91013837e-02, -2.02815175e+00,
                        2.27686334e+00,  8.86460662e-01],
                        [-2.33393744e-01,  2.34054356e+01,  5.04535866e+00, -4.94095922e-01,  1.10760899e+01,
                        1.23367894e+00,  9.56706762e-01, -1.94259071e+00,  1.51433694e+00,  1.03864717e+00,
                        2.61832142e+00,  5.65047359e+00],
                        [-1.21808004e+00,  2.25267982e+01,  4.43598938e+00, -1.23367894e+00,  8.62484741e+00,
                        6.79571211e-01,  7.10804522e-01, -2.02815175e+00, -4.82380331e-01,  7.57651389e-01,
                        2.01259756e+00,  4.92381001e+00],
                        [ 4.56995904e-01,  2.34541931e+01,  4.09919930e+00, -5.39550222e-02,  1.30143728e+01,
                        7.14708507e-01,  1.21028030e+00,  5.11669159e-01,  1.81029797e+00,  1.63120234e+00,
                        4.68021679e+00,  7.72815514e+00]
                    ],
                    [
                        [ 1.15722090e-01,  2.06007233e+01,  3.12920213e+00, -2.12145352e+00,  5.37872887e+00,
                        -3.46879172e+00, -5.80103159e+00, -4.12986183e+00, -2.47871161e+00, -5.52981424e+00,
                        -5.01498747e+00, -8.16204786e-01],
                        [ 2.88171268e+00,  2.27421646e+01,  8.62484741e+00,  5.11669159e-01,  8.51004219e+00,
                        2.41662765e+00,  5.89767814e-01,  2.88171268e+00, -5.19479334e-01, -2.66483307e+00,
                        -2.52526045e+00,  2.33899355e+00],
                        [ 5.37872887e+00,  2.53779068e+01,  8.96760273e+00,  1.32725811e+00,  1.13449488e+01,
                        1.28826988e+00, -6.79571211e-01,  3.45337439e+00,  3.66904116e+00, -2.52526045e+00,
                        -1.18688023e+00,  2.61832142e+00],
                        [ 5.04535866e+00,  2.42790318e+01,  8.68214703e+00,  4.43598938e+00,  9.47692776e+00,
                        3.78886133e-01, -4.61921787e+00,  6.28814161e-01,  4.99953747e-01, -3.69981956e+00,
                        -6.49413019e-02,  3.63825440e+00],
                        [ 6.13118267e+00,  2.50096397e+01,  8.68214703e+00,  4.52764702e+00,  1.17707224e+01,
                        -9.52148531e-03, -3.34540391e+00,  2.16808867e+00,  3.08282948e+00, -3.32997227e+00,
                        -6.32718682e-01,  3.26822829e+00]
                    ]
                ]
            }
        )
        # fmt: on

        model = self._load_model()
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            [
                {"role": "user", "content": "Write a long essay about Portugal."},
            ],
            [
                {"role": "user", "content": "Why is 6 affraid of 7?"},
            ],
        ]
        model_inputs = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, return_dict=True
        ).to(model.device)

        # hardcodes input canvas to avoid depending on random canvas sampling
        decoder_input_ids = torch.arange(model.config.canvas_length, dtype=torch.int32, device=model.device)
        decoder_input_ids = (
            torch.ones((2, model.config.canvas_length), dtype=torch.int32, device=model.device) * decoder_input_ids
        )
        model_inputs["decoder_input_ids"] = decoder_input_ids

        # sanity check: the input has bsz=2
        self.assertEqual(model_inputs["input_ids"].shape[0], 2)
        # sanity check: has padding, and it is left-padding
        self.assertIn(model.config.text_config.pad_token_id, model_inputs["input_ids"])
        self.assertFalse((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, -1]).any())
        self.assertTrue((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, 0]).any())

        with torch.no_grad():
            model_out = model(**model_inputs)
        self.assertEqual(
            list(model_out.logits.shape), [2, model.config.canvas_length, model.config.text_config.vocab_size]
        )
        torch.testing.assert_close(model_out.logits[:, :5, :12].cpu(), torch.tensor(expected_logits.get_expectation()))

    @tooslow
    def test_diffusion_gemma_generate_with_image_batched(self):
        """
        bsz>1, with image, calling `generate`.

        This test has a LOT of non-determinism (sampling canvases, sampling tokens), so it is normal
        if different platforms need to set different expected outcomes.
        """
        # fmt: off
        # print(generated_tokens)
        # Printed after running `torch.set_printoptions(threshold=10000, linewidth=100)`
        expected_sequences = Expectations(
            {
                ("cuda", None): [
                    [
                        1018,    818,  21065,  41025, 236787,    562,  61774,  77985,    529,  23613,   1018,
                        108,   5141,   3059,    580,    506,  16425,   3374,   7377,    529,  49566,   3879,
                        236764,   1298,    506, 211050,   2601,   6861,  15741,    531,    506,  12529,  21065,
                        18414, 236764,  23613,    563,    496,   7097,   5221,    684,   1061,   4191,    607,
                        506,   5442, 236761,   1701,  24744, 236764,    625,    815,   7779,    618,    496,
                        44701,   1534,    506,  10687,   4109,    532,    506,   1799, 236764,    496,  11825,
                        1534,  20117,    532,  70548, 236761,   2282,   3050,  23613,    563,    531,   3050,
                        496,   3996, 105362,    529,  49094,  27877, 236764,  44040,  11858,  24637, 236764,
                        532,    496,   4709,   9226,  78729,   3224,    618,    808,   4834,    689,   1007,
                        236829, 237028, 236746,   5268, 236764,  85539,   7833,  86508,    573,   2613,   5745,
                        653,   8229,   2752,   1765, 236761,    108,  10354,    562,  34641,  84181, 236787,
                        4934, 117838,    576,   4848,    531,  41254,    108,    818,   4083,    529,  23613,
                        563,    886,    529,  35691,    532,  10534, 236761,  42251,    618,    496,  21880,
                        528,    506, 236743, 236770, 236778,    594,   7691,   1913,    506, 117838,    576,
                        4848,    699,  78763,   1044,   6157, 236764,  23613,   5452,    886,    529,    506,
                        24625,   7097, 236772,  26582,    528,   3879, 236761,   5978,   1061,  23789,    964,
                        3187, 204045,    528,   6145,  77486,  42469, 236764,  23613,   6976,  44630, 236761,
                        236743,    108,    818, 236743, 236770, 236810,    594,    532, 236743, 236770, 236825,
                        594,  24744,  11373,   1061,  16522,  17884, 236764,    506,  17884,    529,  41254,
                        236761,   8382,    506,  91295,    529,   9732,   1133,  18934,  12297,    506,  60137,
                        236764,  41082,   2556,  41516,  19482,   6998,    506,    623,  46811,    529, 106819,
                        2098,  15468,    506,   9992,   9383,    532,  17154,   4673, 236761, 125620,   1776,
                        150225,  11788,    506,   5442,   9116,    531,   4673,    528, 236743, 236770, 236812,
                        236819, 236828, 236764
                    ],
                    [
                        2094,   2471,   3831,    506,  16522,  26010,  17936,    528,   5054,  14322, 236764,
                        7151, 236761,    669,   4429,    563,   3523,    699,    506,  15891,   2678, 236764,
                        18482,   9975,  10654,    528,    506,  44529,    532,    506,  69480,  10901,   6615,
                        528,    506,   1695, 236761,    106,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0
                    ]
                ]
            }
        )
        # fmt: on

        model = self._load_model()
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            [
                {"role": "user", "content": "Write a long essay about Portugal."},
            ],
            [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "url": "https://raw.githubusercontent.com/google-gemma/cookbook/refs/heads/main/apps/sample-data/GoldenGate.png",
                        },
                        {"type": "text", "text": "What is shown in this image?"},
                    ],
                }
            ],
        ]
        model_inputs = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, return_dict=True
        ).to(model.device)

        # sanity check: the input has bsz=2
        self.assertEqual(model_inputs["input_ids"].shape[0], 2)
        # sanity check: the input has an image component
        self.assertIn("pixel_values", model_inputs)
        # sanity check: has padding, and it is left-padding
        self.assertIn(model.config.text_config.pad_token_id, model_inputs["input_ids"])
        self.assertFalse((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, -1]).any())
        self.assertTrue((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, 0]).any())

        set_seed(42)
        gen_out = model.generate(**model_inputs, max_new_tokens=256)
        generated_tokens = gen_out.sequences[:, -256:]
        self.assertEqual(generated_tokens.tolist(), expected_sequences.get_expectation())

    @tooslow
    def test_diffusion_gemma_generate_with_image_batched_long(self):
        """
        Ultimate test: bsz>1, with image, with thinking on, calling `generate` with
        max_new_tokens > sliding attention length.

        This test has a LOT of non-determinism (sampling canvases, sampling tokens), so it is normal
        if different platforms need to set different expected outcomes.
        """
        # fmt: off
        # print(generated_tokens)
        # Printed after running `torch.set_printoptions(threshold=10000, linewidth=100)`
        expected_sequences = Expectations(
            {
                ("cuda", None): [
                    [
                        21042,   2846, 236764,  23613,    691,    506,   3988,    529,    506,   1902, 236764,
                        614,  38613,    600,  39349,    699,  14600,    531, 126382, 236764,    532,  73420,
                        531,  72449, 236761,   1174,   6933,   6111,  36628,  12821,    532,   9226,   8770,
                        236764, 175676,    528,    506,  94079,  37704,    688,  31035,   3117, 236764,  17202,
                        684,   1061,  82647,  59491,   1133,  70186, 236764,  35809, 236764,    532,  83798,
                        47397, 236761,  16961, 236764,    506,   3825,    529,  38613,  10734,   5378,    531,
                        16670, 236764,   6641,    684,   1440,  13443,    529,   6658,  23049,    532,    506,
                        56845,    808,  66856, 107130, 236829, 102708,   1208, 187920,    569,  81237, 118882,
                        236761,   1030,    691,    711,   3097,    506,  33275,    567,  24817,    529, 236743,
                        236770, 236819, 236832, 236812, 237028, 236746,  73887, 236764,    569, 236772,  18377,
                        236748,   4806,   1933,   7820,  28622,   1298,  11838,   7006,  37067,    847,    528,
                        506,    520, 204532,    529,  18187, 236789,  79689, 237028,   7705,  23613, 130661,
                        1131,    506,  28239, 236764,  29810,   6693,   8927, 236772,  14531,   7097,    625,
                        563,   3124, 236761,    108,   8551,  69565, 236764,  23613,    563,  33887,  12801,
                        9785,   1061,  30997,   2425, 236761,    669,   4695,    563,  52803,    532,   3826,
                        236764,   2033,    531,    506, 176165,  15706,    530,  11849, 236764,    506,   1902,
                        236858, 236751,  24625, 115368,    774,  10135,   4128, 236764,   1298,   6944,   3776,
                        81026, 103534,    531,  25465,  26607,   1949, 236761,    669,   3988,    563,  27222,
                        684,    506,   1429,   3194,   5400, 236764,    496,  12529, 236764,  13935, 152990,
                        529,  26012, 154863,    532,  64192,  32049,  24599,   1298,   1972,  12489,    657,
                        496,  31878,  17723, 236761,   2282,    506,   8710, 236764,    506, 201172,   5469,
                        20997,  62260, 236764,  11497,   5442,    505,  21294, 236764,    532,    496,   6962,
                        32920,  10022, 236761,  43725,    506,   2891, 236764,    506,  21065,  18414,   7474,
                        506,   4512,  61102
                    ],
                    [
                        563,  88076,    580,   1903,    529,    672,   5441, 236761,    108,    902,    506,
                        1695, 236764,   3418,    506,   1813, 236764,    506,  19519, 236764, 127150, 236772,
                        30497,  26607,    529,    506,  69480,  10901,   6615,    659,  11325,   1208,    506,
                        6425,   1933,   7217, 236761,    669,  15408,    563,   7804,    532,   1982, 236764,
                        22169,    496,  17635,   1719, 236764,  30439,   3538,  37676,    580,    506,  31035,
                        4065,    529,    506,  11825,    532,    506,   7845, 236761,    106,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
                        0,      0,      0
                    ]
                ]
            }
        )
        # fmt: on

        model = self._load_model()
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            [
                {"role": "user", "content": "Write a long essay about Portugal."},
            ],
            [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "url": "https://raw.githubusercontent.com/google-gemma/cookbook/refs/heads/main/apps/sample-data/GoldenGate.png",
                        },
                        {"type": "text", "text": "Describe the image in detail."},
                    ],
                }
            ],
        ]
        model_inputs = processor.apply_chat_template(
            chat,
            tokenize=True,
            return_tensors="pt",
            add_generation_prompt=True,
            return_dict=True,
            enable_thinking=True,
        ).to(model.device)

        # sanity check: the input has bsz=2
        self.assertEqual(model_inputs["input_ids"].shape[0], 2)
        # sanity check: the input has an image component
        self.assertIn("pixel_values", model_inputs)
        # sanity check: has padding, and it is left-padding
        self.assertIn(model.config.text_config.pad_token_id, model_inputs["input_ids"])
        self.assertFalse((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, -1]).any())
        self.assertTrue((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, 0]).any())

        set_seed(42)
        gen_out = model.generate(**model_inputs, max_new_tokens=1280)

        # sanity check: we went beyond the sliding window length
        self.assertTrue(gen_out.sequences.shape[1] > model.config.text_config.sliding_window)
        generated_tokens = gen_out.sequences[:, -256:]
        self.assertEqual(generated_tokens.tolist(), expected_sequences.get_expectation())

    # The tests below use a botched model, using the first six layers, to respect memory constraints of HF's CI.
    # Six layers = minimum number of layers to have both sliding window attention and full attention.
    @slow
    def test_minified_diffusion_gemma_forward_text_only(self):
        """[Minified (6 first layers)] bsz=1, no image"""
        # fmt: off
        # print(model_out.logits[:, :10, :12])
        # Printed after running `torch.set_printoptions(precision=8, threshold=200, linewidth=100)`
        expected_logits = Expectations(
            {
                ("cuda", None): [
                    [[ 20.99029922, -29.92623520,  -3.96111226,  29.14220428,  29.95949554,  28.15386963,
                    27.72417641,   1.10106778, -27.28695297, -29.00568390,  -1.60783303, -18.21017647],
                    [ 29.49378586,  17.06575012,  25.72111511,  29.87012291,  29.99975395, -17.31759644,
                    -29.92115784, -20.06009483, -14.73279667, -28.84785843, -29.78616714, -29.99817467],
                    [ -1.54550576, -29.99688721, -29.98895645,  26.56055069,  13.76502132, -24.44923782,
                    -4.55818033, -17.40074539, -29.42227936, -29.94157219, -25.58675003, -29.89002419],
                    [ -1.14787686, -29.44104195,  16.55123901,  26.77043152,  29.86572838,  13.56690025,
                    21.36647987, -10.31235313, -29.22277641, -28.15386963,   2.49422908, -24.27903175],
                    [ 28.09328270, -28.88492966,  27.24345016,  29.90688133,  29.90688133,  18.13103485,
                    27.99878311,   4.00716066, -24.53262520, -28.88492966,  24.44923782, -22.52679825],
                    [  0.10205039, -29.93098831,   7.55281830,  28.48003769,  29.93098831, -14.25369835,
                    21.11716843, -19.27641296, -29.42227936, -29.02185059,   8.16400814, -27.41356468],
                    [ 22.19349098, -27.10892868, -24.10417366,  27.37199974,  26.82098389, -10.03579903,
                        4.71070147, -16.46408844, -29.68155670, -25.30630493,  24.77605438,   1.80251384],
                    [ 20.73208427, -29.44104195,  17.97154999,  29.67082214,  29.82481384,  -9.47692776,
                    26.77043152, -12.86174679, -29.40289116, -29.44104195,   4.77164030, -22.13672829],
                    [  7.81561375, -29.63641167,  15.66178513,  29.14220428,  29.40289116, -15.93279934,
                    28.37808609, -13.11566544, -28.35157776, -29.76376724,  -7.49424887, -19.20282745],
                    [ 17.72932053, -29.98895645,  29.31869316,  29.31869316,  29.69194221, -25.00963974,
                        7.58207989, -28.12381363, -29.51024628, -29.93098831, -25.58675003, -26.87078667]]
                    ],
                ("xpu", 5): [
                    [[ 21.18004799, -29.92373848,  -3.05190611,  29.14220428,  29.96082115,  28.18346024,
                    27.83115768,   0.69128382, -27.28695488, -28.98925400,  -0.99962980, -17.81046677],
                    [ 29.47677803,  17.40074348,  25.85169601,  29.87012291,  29.99976730, -16.98100090,
                    -29.91848755, -18.97970009, -14.15672588, -28.78997993, -29.76376724, -29.99829102],
                    [ -2.10590601, -29.99708939, -29.98819351,  26.50611115,  14.20525932, -24.77605629,
                    -4.55818033, -17.81046677, -29.42227936, -29.94157028, -25.37790871, -29.89002419],
                    [ -1.60004294, -29.42227936,  16.46408844,  26.77043152,  29.86119080,  14.59008121,
                    22.02214432, -10.96796513, -29.24793625, -28.24126434,   2.33899355, -24.10417557],
                    [ 28.06227303, -28.97256088,  27.15444565,  29.90688133,  29.90992928,  18.52275848,
                    27.86572647,   4.64972305, -24.69602013, -28.92082787,  24.36471176, -22.36170006],
                    [  2.86623120, -29.89710236,   5.16673756,  28.57573891,  29.90688133, -13.76502132,
                    21.11716843, -18.82897186, -29.38286209, -29.12799644,   4.98460531, -27.57362747],
                    [ 22.07961273, -24.93285179, -22.47210693,  27.19928551,  27.49481773,  -9.58930302,
                        5.13640928, -16.81029701, -29.71171188, -25.30630493,  24.77605438,   1.87256217],
                    [ 20.86193848, -29.44104195,  17.72932243,  29.67082405,  29.83608246,  -9.13802814,
                    26.82098389, -13.41732407, -29.42227936, -29.44104195,   5.34847832, -22.13672447],
                    [  8.73937702, -29.52617455,  14.54231644,  29.24793625,  29.36216545, -16.72434616,
                    28.72925377, -11.98154259, -28.32464409, -29.70199013,  -5.98130035, -17.72932053],
                    [ 17.89120865, -29.98819351,  29.29587173,  29.31869316,  29.69194031, -25.00963974,
                        7.55281830, -28.15386772, -29.51024628, -29.92865372, -25.58675194, -26.71912193]]
                    ]
            }
        )
        # fmt: on

        model = self._load_model(minified=True)
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            {"role": "user", "content": "Write a long essay about Portugal."},
        ]
        model_inputs = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, return_dict=True
        ).to(model.device)

        # hardcodes input canvas to avoid depending on random canvas sampling
        decoder_input_ids = torch.arange(model.config.canvas_length, dtype=torch.int32, device=model.device).unsqueeze(
            0
        )
        model_inputs["decoder_input_ids"] = decoder_input_ids

        with torch.no_grad():
            model_out = model(**model_inputs)
        self.assertEqual(
            list(model_out.logits.shape), [1, model.config.canvas_length, model.config.text_config.vocab_size]
        )
        torch.testing.assert_close(
            model_out.logits[:, :10, :12].cpu(), torch.tensor(expected_logits.get_expectation())
        )

    @slow
    @require_deterministic_for_xpu
    def test_minified_diffusion_gemma_forward_with_image(self):
        """[Minified (6 first layers)] bsz=1, with image"""
        # fmt: off
        # print(model_out.logits[:, :10, :12])
        # Printed after running `torch.set_printoptions(precision=8, threshold=200, linewidth=100)`
        expected_logits = Expectations(
            {
                ("cuda", None): [[
                    [ 18.59990692, -29.96792030, -26.71912193,  27.65005684,  29.36216545,  24.36471176,
                        6.40016174,   6.81642532, -19.49482346, -22.19349098,  -0.41794172, -24.77605438],
                    [ 28.35157776,  26.91985321,  28.80958176,  29.63641167,  29.99991417,  -9.47692776,
                    -29.90373039, -25.30630493,  -0.17089660, -29.27229309, -29.90688133, -29.99860191],
                    [ -2.52526045, -29.38286209, -29.96454811,  25.08536720,   7.08258057, -27.01581001,
                    -21.11716843, -26.45086479, -28.62135887, -29.92865181, -28.90302277, -29.99154091],
                    [ -3.26822829, -27.10892868, -26.91985321,  25.91559219,  29.90373039,   1.14007568,
                        1.52212954,  -5.98130083, -27.61213493,  -8.79653835,  -8.22184658, -25.91559219],
                    [ 28.74981880, -29.79991150,  28.92082787,  29.91287804,  29.98117638,   9.47692776,
                    27.68740082,  -8.51004219, -12.13874054, -26.56055069,  28.45515060, -22.13672829],
                    [  0.53900456, -29.77148056,  -0.21874613,  26.22159958,  29.67082214, -20.92630577,
                        3.65364885, -21.18004990, -29.55650711, -28.18346024,  -3.20645595, -28.64363098],
                    [-18.52275848, -26.16215897, -28.62135887,  14.05937004, -12.29514503, -25.58675003,
                    -15.38721085, -26.33790970, -29.47677803,  -7.93200207,  27.37199974,  -2.44767213],
                    [  2.43215036, -29.29586983,  12.39897156,  28.50453377,  26.66704369, -24.44923782,
                        3.40711212, -24.10417366, -29.74754906, -29.64826202,  -6.42998266, -26.66704369],
                    [-14.49445248, -29.22277641,   9.98025513,  27.01581001,  25.30630493, -23.40543556,
                    14.01054859, -17.48348808, -28.62135887, -29.36216545,  -9.08129215, -27.61213493],
                    [  4.64972258, -29.98651123,  28.42986679,  28.88492966,  26.39480209, -28.40418053,
                    -15.38721085, -28.42986679, -29.09888268, -29.77148056, -26.61419106, -28.35157776]
                    ]],
                ("xpu", 5): [[
                    [ 18.97970009, -29.96791840, -26.61419106,  27.76038742,  29.54159164,  24.27903175,
                        8.22184658,   8.04813480, -18.97970009, -22.74216461,  -1.06205606, -24.36471176],
                    [ 28.35157585,  27.45450211,  28.68712234,  29.65972710,  29.99992943,  -9.02448273,
                    -29.89710236, -25.23368645,   0.80449456, -29.29587173, -29.90688324, -29.99850464],
                    [ -3.60746002, -29.76376724, -29.95053482,  24.01497269,   3.29910398, -26.22159958,
                    -19.49482346, -27.19928741, -28.90302277, -29.94886017, -27.72417450, -29.98458862],
                    [ -7.81561375, -26.96819115, -25.85169601,  26.16215897,  29.90688133,  -3.36083364,
                        3.63825440,  -3.60746002, -27.83115768,  -9.02448273, -11.02206802, -25.65441132],
                    [ 28.70835495, -29.78616714,  29.05341339,  29.90992928,  29.98179245,   9.58930302,
                    27.57362556,  -8.96760178, -13.01437187, -26.82098389,  28.55237770, -22.36170006],
                    [  0.66395408, -29.73022270,  -1.74023187,  26.50611115,  29.76376534, -19.05447006,
                        0.96841347, -21.78874397, -29.47677803, -27.68740082,   1.40522099, -28.55237961],
                    [-17.15010071, -26.28017807, -28.86654472,  17.48348999,  -6.93485880, -26.50610924,
                    -17.81046677, -24.85499382, -29.47677803,  -6.81642532,  27.53452492,  -3.60746002],
                    [  3.59206033, -29.29587173,  14.05937004,  28.59873009,  26.71912193, -24.10417557,
                        4.19115925, -23.59856415, -29.74754906, -29.62416267,  -5.25765705, -26.71912193],
                    [-13.76502132, -29.19678307,  10.58693695,  27.10892868,  26.16215897, -23.45419121,
                    14.01054859, -16.63799286, -28.57573891, -29.36216545,  -8.68214703, -27.57362747],
                    [  6.87567043, -29.98605347,  28.42986679,  29.00568199,  26.82098389, -28.06227493,
                    -15.57065392, -28.40418053, -29.15618324, -29.77894592, -26.71912193, -28.40418053]
                    ]]
            }
        )
        # fmt: on

        model = self._load_model(minified=True)
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "url": "https://raw.githubusercontent.com/google-gemma/cookbook/refs/heads/main/apps/sample-data/GoldenGate.png",
                    },
                    {"type": "text", "text": "What is shown in this image?"},
                ],
            }
        ]
        model_inputs = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, return_dict=True
        ).to(model.device)

        # hardcodes input canvas to avoid depending on random canvas sampling
        decoder_input_ids = torch.arange(model.config.canvas_length, dtype=torch.int32, device=model.device).unsqueeze(
            0
        )
        model_inputs["decoder_input_ids"] = decoder_input_ids

        # sanity check: has image inputs
        self.assertIn("pixel_values", model_inputs)

        with torch.no_grad():
            model_out = model(**model_inputs)
        self.assertEqual(
            list(model_out.logits.shape), [1, model.config.canvas_length, model.config.text_config.vocab_size]
        )
        torch.testing.assert_close(
            model_out.logits[:, :10, :12].cpu(), torch.tensor(expected_logits.get_expectation())
        )

    @slow
    @require_deterministic_for_xpu
    def test_minified_diffusion_gemma_forward_batched(self):
        """[Minified (6 first layers)] bsz>1, no image"""
        # print(model_out.logits[:, :5, :12])
        # Printed after running `torch.set_printoptions(precision=8, threshold=200, linewidth=100)`
        expected_logits = Expectations(
            {
                ("cuda", None): [
                    [[ 21.60994148, -29.94348526,  -9.81315899,  29.09888268,  29.96454811,  28.35157776,
                    27.37199974,   5.13640928, -26.56055069, -28.70835686,  -0.13183509, -18.97969818],
                    [ 29.54159164,  12.19096375,  24.93285370,  29.89002419,  29.99975395, -18.44520950,
                    -29.92623520, -21.05392075, -17.23404694, -29.08397102, -29.74754906, -29.99850464],
                    [ -0.04980464, -29.99727631, -29.99208641,  26.61419106,  17.81046486, -23.92457008,
                    -6.10123110, -18.13103485, -29.27229309, -29.93543434, -25.30630493, -29.86572838],
                    [  0.15429553, -29.38286209,  14.54231548,  26.77043152,  29.90373039,  15.93279934,
                    21.30470467, -10.20196629, -29.19678497, -27.99878311,   4.03784895, -24.36471176],
                    [ 27.99878311, -28.93834877,  27.01581001,  29.90992928,  29.92865181,  18.36726570,
                    27.53452492,   3.69981956, -24.27903175, -28.82887459,  25.23368835, -22.58115005]],

                    [[ 20.73208427, -29.81889534, -14.01054859,  29.42227936,  29.85649490,  27.61213493,
                    23.05516243,   9.81315899, -28.26949692, -28.06227112,  -1.21028030, -25.08536720],
                    [ 29.09888268,  20.12901306,  19.99079323,  29.89002419,  29.99984360, -17.64777756,
                    -29.86572838, -14.54231548, -17.64777756, -28.52864838, -29.72112274, -29.99791336],
                    [ 12.08642769, -29.97701073, -29.97776413,  28.92082787,  18.90453339, -20.33346367,
                        4.09919930, -11.18388462, -29.52617455, -29.87437057, -22.74216461, -29.72112274],
                    [-10.03579903, -28.59873009,  15.47912979,  27.06271935,  29.69194221,  15.01560497,
                    16.81029892, -15.66178513, -29.49378586, -27.86572838,  -3.56125450, -24.19219208],
                    [ 25.37790680, -29.47677803,  27.32979774,  29.89710426,  29.67082214,  13.41732597,
                    24.01497269,   4.95421267, -26.96818924, -28.42986679,  19.92110252, -24.44923782]]
                ],
                ("xpu", 5): [
                    [[ 21.96432686, -29.94348526,  -9.53315163,  29.14220428,  29.96683502,  28.37808609,
                    27.45450211,   6.10123110, -26.50610924, -28.70835495,   0.21093403, -18.36726570],
                    [ 29.54159164,  12.39897156,  25.08536530,  29.88630295,  29.99973488, -18.05149078,
                    -29.91848755, -21.18004990, -16.63799286, -29.15618324, -29.74754906, -29.99850464],
                    [ -0.82791477, -29.99761581, -29.99153900,  26.50611115,  17.72932243, -23.92457008,
                    -6.54913378, -17.81046677, -29.24793625, -29.93543434, -25.08536720, -29.87012291],
                    [  1.48316491, -29.42227936,  14.10809517,  27.01580811,  29.90688133,  16.02234268,
                    21.30470467,  -9.64538002, -29.16993713, -27.93329048,   4.34424686, -24.19219208],
                    [ 28.06227303, -28.84785843,  27.06271935,  29.91287804,  29.92623711,  18.21017647,
                    27.65005684,   4.16051531, -24.19219208, -28.77005959,  25.30630302, -22.95215416]],

                    [[ 20.86193848, -29.81277466, -14.30204105,  29.42228127,  29.85649490,  27.65005684,
                    23.05516243,   9.81316090, -28.26949692, -28.06227493,  -0.92158508, -25.08536720],
                    [ 29.11355782,  20.79719925,  20.73208427,  29.89710236,  29.99983406, -17.81046677,
                    -29.87012291, -15.01560688, -17.72932053, -28.42986488, -29.73022270, -29.99791336],
                    [ 11.61169815, -29.97701073, -29.97623253,  28.88492966,  18.28891945, -20.19754601,
                        4.12986183, -11.18388557, -29.52617455, -29.87437248, -22.79516220, -29.73022270],
                    [-10.42242527, -28.57573891,  15.57065582,  26.91985321,  29.70199203,  15.38721275,
                    16.46408844, -14.82745552, -29.47677803, -27.76038551,  -4.00716066, -24.19219208],
                    [ 25.44850922, -29.49378586,  27.41356277,  29.89710236,  29.64826012,  13.46727657,
                    24.19219208,   4.40541744, -27.10892868, -28.50453377,  19.42241287, -24.36471176]]
                ],
            }
        )  # fmt: skip

        model = self._load_model(minified=True)
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            [
                {"role": "user", "content": "Write a long essay about Portugal."},
            ],
            [
                {"role": "user", "content": "Why is 6 affraid of 7?"},
            ],
        ]
        model_inputs = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, return_dict=True
        ).to(model.device)

        # hardcodes input canvas to avoid depending on random canvas sampling
        decoder_input_ids = torch.arange(model.config.canvas_length, dtype=torch.int32, device=model.device)
        decoder_input_ids = (
            torch.ones((2, model.config.canvas_length), dtype=torch.int32, device=model.device) * decoder_input_ids
        )
        model_inputs["decoder_input_ids"] = decoder_input_ids

        # sanity check: the input has bsz=2
        self.assertEqual(model_inputs["input_ids"].shape[0], 2)
        # sanity check: has padding, and it is left-padding
        self.assertIn(model.config.text_config.pad_token_id, model_inputs["input_ids"])
        self.assertFalse((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, -1]).any())
        self.assertTrue((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, 0]).any())

        with torch.no_grad():
            model_out = model(**model_inputs)
        self.assertEqual(
            list(model_out.logits.shape), [2, model.config.canvas_length, model.config.text_config.vocab_size]
        )
        torch.testing.assert_close(model_out.logits[:, :5, :12].cpu(), torch.tensor(expected_logits.get_expectation()))

    @slow
    @require_deterministic_for_xpu
    def test_minified_diffusion_gemma_generate_with_image_batched(self):
        """
        [Minified (6 first layers)] bsz>1, with image, calling `generate`.

        This test has a LOT of non-determinism (sampling canvases, sampling tokens), so it is normal
        if different platforms need to set different expected outcomes.
        """
        # print(generated_tokens)
        # Printed after running `torch.set_printoptions(threshold=10000, linewidth=100)`
        expected_sequences = Expectations(
            {
                ("cuda", None): [
                    [48, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 728, 48, 48, 48, 48, 598, 371, 48, 48, 595, 48, 48,
                     983, 371, 48, 598, 371, 48, 48, 595, 48, 48, 371, 595, 371, 516, 371, 595, 516, 516, 516, 379, 48, 371, 48, 48, 371, 371, 173, 516, 4,
                     595, 516, 371, 595, 48, 516, 48, 371, 516, 495, 371, 516, 516, 983, 371, 516, 516, 48, 371, 516, 371, 48, 516, 48, 371, 516, 48, 983,
                     48, 841, 371, 48, 516, 48, 379, 516, 832, 48, 371, 48, 371, 371, 371, 48, 48, 48, 371, 371, 48, 48, 48, 48, 48, 371, 371, 48, 48, 173,
                     48, 48, 48, 516, 48, 48, 48, 371, 379, 48, 371, 48, 516, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48,
                     48, 841, 48, 48, 48, 48, 516, 173, 371, 48, 48, 48, 516, 48, 48, 516, 48, 48, 48, 48, 371, 48, 371, 48, 48, 371, 48, 379, 48, 371, 48,
                     832, 595, 48, 688, 48, 48, 371, 48, 371, 48, 48, 48, 48, 48, 516, 371, 371, 48, 48, 48, 371, 48, 405, 48, 48, 48, 48, 48, 48, 48, 48,
                     48, 48, 48, 371, 48, 371, 48, 371, 48, 48, 371, 371, 516, 48, 371, 48, 516, 516, 48, 48, 48, 173, 48, 48, 48, 516, 371, 516, 48, 160,
                     516, 48, 371, 173, 48, 48, 48, 371, 48, 48, 48, 48, 516],
                    [48, 371, 627, 371, 48, 651, 371, 379, 379, 371, 371, 371, 595, 48, 48, 595, 371, 48, 371, 379, 595, 595, 371, 977, 48, 48, 48, 517,
                     371, 841, 48, 48, 517, 832, 516, 371, 418, 48, 48, 48, 48, 371, 48, 48, 48, 48, 371, 516, 379, 371, 371, 48, 48, 371, 371, 656, 371,
                     48, 371, 371, 371, 48, 48, 48, 48, 371, 371, 841, 371, 371, 841, 516, 371, 371, 841, 371, 656, 48, 48, 371, 983, 48, 48, 379, 48, 48,
                     595, 48, 371, 983, 371, 983, 48, 48, 48, 371, 48, 48, 48, 48, 516, 48, 48, 48, 48, 48, 48, 48, 48, 48, 172, 48, 48, 371, 598, 48, 48,
                     48, 371, 656, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 371, 841, 48, 48, 48, 48, 832, 48, 48,
                     498, 832, 371, 48, 371, 371, 48, 48, 371, 48, 371, 371, 48, 48, 371, 48, 654, 516, 48, 48, 48, 172, 371, 634, 48, 48, 48, 48, 516,
                     48, 48, 371, 48, 48, 595, 371, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 371, 48, 371, 379, 595, 48, 48, 48, 371, 48, 48, 634, 48,
                     48, 48, 172, 48, 48, 48, 48, 379, 48, 48, 48, 371, 48, 371, 48, 48, 48, 371, 48, 48, 371, 48, 48, 48, 371, 48, 48, 371, 48, 371, 48,
                     48, 549, 48, 48, 48, 48, 48, 634, 499, 172, 48, 48, 48, 48, 48, 48, 48
                    ]
                ],
                ("xpu", 5): [
                    [48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 173, 4, 48, 48, 48, 48, 48, 48, 371, 48, 48, 371, 371, 48, 48, 48,
                     371, 371, 371, 48, 371, 516, 499, 516, 371, 371, 371, 173, 516, 48, 371, 371, 48, 371, 48, 595, 173, 371, 379, 371, 516, 379, 48, 48,
                     48, 371, 48, 371, 371, 48, 516, 516, 48, 595, 48, 48, 48, 841, 371, 516, 48, 516, 371, 48, 48, 371, 48, 500, 859, 48, 48, 48, 48, 48,
                     48, 48, 371, 48, 371, 48, 371, 371, 48, 48, 48, 48, 48, 371, 48, 48, 48, 48, 172, 48, 48, 371, 371, 371, 48, 48, 48, 48, 48, 48, 48,
                     48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 405, 48, 48, 48, 48, 48, 48, 48, 48, 371, 48, 48, 48, 48, 48, 48,
                     48, 48, 371, 48, 48, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 516, 48, 48, 549, 48, 48, 371, 48, 48, 48, 48, 371, 48, 48, 48,
                     48, 371, 371, 48, 48, 48, 48, 48, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 371, 48, 48,
                     516, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48, 172, 48, 48, 173, 48, 48, 516, 371, 48, 516, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48,
                     48, 48, 48, 48, 48, 48],
                    [5, 371, 173, 48, 371, 48, 48, 371, 48, 48, 371, 48, 371, 48, 371, 379, 841, 48, 48, 48, 371, 48, 48, 48, 371, 371, 371, 48, 48, 379,
                     656, 48, 379, 48, 371, 172, 48, 371, 48, 48, 371, 48, 48, 371, 48, 371, 371, 48, 48, 48, 48, 48, 48, 371, 371, 48, 48, 48, 48, 371, 48,
                     371, 48, 371, 48, 371, 48, 48, 48, 48, 48, 371, 173, 371, 371, 379, 48, 48, 48, 48, 48, 48, 48, 48, 371, 48, 48, 841, 48, 371, 983,
                     371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 371, 48, 48, 48, 48,
                     48, 48, 48, 48, 48, 48, 48, 371, 48, 48, 48, 371, 48, 371, 48, 48, 371, 48, 48, 48, 48, 48, 48, 371, 48, 48, 48, 48, 48, 48, 371, 48,
                     48, 48, 48, 48, 48, 516, 48, 48, 48, 371, 371, 841, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 371, 371, 48, 371, 48, 48, 48,
                     48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48,
                     371, 48, 48, 48, 48, 48, 371, 371, 48, 48, 48, 48, 371, 48, 48, 371, 172, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48,
                     4, 48, 371
                    ]
                ]
            }
        )  # fmt: skip

        model = self._load_model(minified=True)
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            [
                {"role": "user", "content": "Write a long essay about Portugal."},
            ],
            [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "url": "https://raw.githubusercontent.com/google-gemma/cookbook/refs/heads/main/apps/sample-data/GoldenGate.png",
                        },
                        {"type": "text", "text": "What is shown in this image?"},
                    ],
                }
            ],
        ]
        model_inputs = processor.apply_chat_template(
            chat, tokenize=True, return_tensors="pt", add_generation_prompt=True, return_dict=True
        ).to(model.device)

        # sanity check: the input has bsz=2
        self.assertEqual(model_inputs["input_ids"].shape[0], 2)
        # sanity check: the input has an image component
        self.assertIn("pixel_values", model_inputs)
        # sanity check: has padding, and it is left-padding
        self.assertIn(model.config.text_config.pad_token_id, model_inputs["input_ids"])
        self.assertFalse((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, -1]).any())
        self.assertTrue((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, 0]).any())

        set_seed(42)
        gen_out = model.generate(**model_inputs, max_new_tokens=256)
        generated_tokens = gen_out.sequences[:, -256:]
        self.assertEqual(generated_tokens.tolist(), expected_sequences.get_expectation())

    @slow
    @require_deterministic_for_xpu
    def test_minified_diffusion_gemma_generate_with_image_batched_long(self):
        """
        [Minified (6 first layers)] Ultimate test: bsz>1, with image, with thinking on, calling `generate` with
        max_new_tokens > sliding attention length.

        This test has a LOT of non-determinism (sampling canvases, sampling tokens), so it is normal
        if different platforms need to set different expected outcomes.
        """
        # print(generated_tokens)
        # Printed after running `torch.set_printoptions(threshold=10000, linewidth=100)`
        expected_sequences = Expectations(
            {
                ("cuda", None): [
                    [48, 48, 48, 567, 48, 405, 48, 48, 379, 48, 379, 379, 379, 379, 379, 379, 48, 48, 48, 48, 48, 379, 379, 379, 379, 48, 48, 379, 48, 48,
                     379, 379, 379, 379, 379, 379, 379, 379, 379, 379, 379, 595, 48, 48, 379, 48, 379, 48, 379, 379, 48, 516, 379, 48, 48, 48, 379, 379,
                     48, 516, 379, 595, 379, 379, 379, 48, 516, 418, 48, 379, 48, 379, 48, 48, 379, 48, 48, 379, 48, 379, 379, 379, 48, 48, 48, 48, 379,
                     48, 48, 379, 379, 379, 48, 48, 48, 48, 48, 48, 379, 379, 48, 48, 379, 48, 48, 516, 48, 48, 379, 48, 379, 48, 48, 48, 48, 48, 48, 48,
                     379, 379, 48, 379, 48, 48, 379, 379, 48, 48, 48, 48, 48, 48, 48, 379, 379, 48, 48, 379, 48, 48, 379, 48, 379, 48, 48, 379, 48, 48, 48,
                     48, 48, 48, 48, 977, 48, 48, 379, 48, 379, 379, 379, 379, 379, 48, 48, 48, 48, 48, 379, 379, 379, 379, 48, 48, 48, 379, 48, 48, 48, 48,
                     379, 379, 379, 48, 379, 48, 48, 48, 48, 48, 48, 379, 48, 379, 48, 48, 48, 48, 48, 516, 379, 48, 48, 379, 48, 48, 48, 516, 48, 379, 48,
                     48, 48, 48, 379, 48, 48, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 48, 48, 516, 48, 48, 48, 379, 379, 48, 379, 48, 371,
                     48, 379, 48, 48, 379, 48, 48, 379, 48, 48, 48, 48, 48
                    ],
                    [48, 48, 48, 379, 48, 379, 379, 48, 48, 379, 379, 516, 516, 379, 48, 48, 379, 379, 379, 379, 379, 379, 379, 379, 48, 48, 379, 379, 379,
                     48, 48, 48, 379, 379, 379, 595, 379, 379, 379, 379, 379, 48, 48, 516, 48, 379, 379, 379, 379, 48, 48, 379, 595, 48, 379, 379, 379, 48,
                     379, 379, 379, 48, 379, 48, 516, 379, 379, 48, 379, 379, 379, 379, 379, 379, 379, 379, 48, 48, 48, 379, 48, 379, 48, 48, 379, 516, 48,
                     48, 379, 379, 379, 379, 379, 379, 48, 48, 48, 516, 379, 379, 48, 595, 48, 379, 48, 48, 516, 48, 499, 48, 48, 48, 379, 379, 48, 48, 379,
                     379, 48, 379, 516, 379, 379, 48, 379, 379, 379, 379, 48, 379, 379, 48, 379, 379, 379, 379, 371, 48, 379, 48, 48, 371, 48, 48, 48, 48,
                     379, 379, 379, 379, 48, 379, 379, 48, 48, 48, 48, 48, 379, 379, 48, 379, 379, 379, 48, 379, 48, 379, 379, 379, 48, 379, 48, 841, 379,
                     379, 379, 48, 48, 379, 379, 48, 48, 48, 379, 48, 48, 379, 499, 379, 48, 516, 48, 48, 379, 48, 48, 379, 48, 379, 379, 48, 379, 379, 379,
                     48, 379, 48, 379, 379, 48, 379, 516, 48, 48, 516, 379, 379, 48, 379, 48, 379, 48, 379, 379, 379, 48, 48, 48, 48, 379, 48, 379, 379, 48,
                     48, 48, 48, 379, 48, 48, 379, 371, 48, 48, 371, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48
                    ]
                ],
                ("xpu", 5): [
                    [48, 48, 48, 48, 48, 48, 48, 379, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 379, 48, 48, 48, 48, 48,
                     48, 48, 48, 48, 48, 379, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 379, 48, 48, 48, 48, 48,
                     48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48, 379, 48, 48, 371, 48, 48, 48, 48, 48, 48, 48, 48,
                     48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 405, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48,
                     48, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48, 48, 379, 48, 379, 48, 48, 379,
                     48, 48, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 379, 48, 379, 48, 48, 48, 379, 48,
                     48, 379, 48, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48, 379, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48,
                     48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 48, 371, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48
                    ],
                    [21, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48, 48, 379, 48, 48, 379,
                     48, 48, 48, 379, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48, 48, 379, 48, 48, 48, 379, 48, 48, 48, 48, 48, 48, 379, 48,
                     48, 48, 48, 48, 48, 48, 48, 48, 418, 48, 48, 379, 48, 379, 48, 48, 48, 379, 48, 48, 379, 48, 48, 48, 48, 48, 379, 48, 48, 379, 48, 48,
                     379, 379, 48, 48, 48, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 405, 379, 48, 48, 48, 48, 48, 379, 48, 48, 48, 379, 48, 48,
                     379, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 379, 48, 48, 379, 48, 48, 48, 48, 48, 48, 379, 48, 48, 379, 48, 48, 48,
                     48, 379, 499, 48, 48, 48, 379, 48, 48, 48, 48, 48, 499, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 379, 48, 48, 48, 48, 48, 379, 48, 48,
                     48, 48, 48, 379, 48, 48, 48, 48, 48, 379, 379, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 418, 48, 48, 48, 379, 48, 48, 48, 48, 48,
                     48, 48, 48, 48, 371, 48, 48, 48, 48, 379, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48
                    ]
                ]
            }
        )  # fmt: skip

        model = self._load_model(minified=True)
        processor = AutoProcessor.from_pretrained(self._model_path)
        chat = [
            [
                {"role": "user", "content": "Write a long essay about Portugal."},
            ],
            [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "url": "https://raw.githubusercontent.com/google-gemma/cookbook/refs/heads/main/apps/sample-data/GoldenGate.png",
                        },
                        {"type": "text", "text": "Describe the image in detail."},
                    ],
                }
            ],
        ]
        model_inputs = processor.apply_chat_template(
            chat,
            tokenize=True,
            return_tensors="pt",
            add_generation_prompt=True,
            return_dict=True,
            enable_thinking=True,
        ).to(model.device)

        # sanity check: the input has bsz=2
        self.assertEqual(model_inputs["input_ids"].shape[0], 2)
        # sanity check: the input has an image component
        self.assertIn("pixel_values", model_inputs)
        # sanity check: has padding, and it is left-padding
        self.assertIn(model.config.text_config.pad_token_id, model_inputs["input_ids"])
        self.assertFalse((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, -1]).any())
        self.assertTrue((model.config.text_config.pad_token_id == model_inputs["input_ids"][:, 0]).any())

        set_seed(42)
        gen_out = model.generate(**model_inputs, max_new_tokens=1280)

        # sanity check: we went beyond the sliding window length
        self.assertTrue(gen_out.sequences.shape[1] > model.config.text_config.sliding_window)
        generated_tokens = gen_out.sequences[:, -256:]
        self.assertEqual(generated_tokens.tolist(), expected_sequences.get_expectation())