diff --git "a/index.html" "b/index.html"
--- "a/index.html"
+++ "b/index.html"
@@ -61,6 +61,7875 @@
         }
     </style>
 </head>
+<script id="main-code" type="text/plain">
+    #                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+    #           This file was automatically generated from src/transformers/models/gemma3n/modular_gemma3n.py.
+    #               Do NOT edit this file manually as any edits will be overwritten by the generation of
+    #             the file from the modular. If any change should be done, please apply the change to the
+    #                          modular_gemma3n.py file directly. One of our CI enforces this.
+    #                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+    # coding=utf-8
+    # Copyright 2025 Google Inc. HuggingFace Inc. 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.
+    import copy
+    import math
+    from collections.abc import Callable, Sequence
+    from dataclasses import dataclass
+    from typing import Optional, Union
+    
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    
+    from ...activations import ACT2FN
+    from ...cache_utils import Cache, DynamicCache, HybridCache
+    from ...generation import GenerationMixin
+    from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
+    from ...modeling_flash_attention_utils import FlashAttentionKwargs
+    from ...modeling_layers import GradientCheckpointingLayer
+    from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+    from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+    from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+    from ...processing_utils import Unpack
+    from ...utils import (
+        ModelOutput,
+        auto_docstring,
+        can_return_tuple,
+        is_torchdynamo_compiling,
+        logging,
+    )
+    from ...utils.deprecation import deprecate_kwarg
+    from ..auto import AutoModel
+    from .configuration_gemma3n import Gemma3nAudioConfig, Gemma3nConfig, Gemma3nTextConfig, Gemma3nVisionConfig
+    
+    
+    logger = logging.get_logger(__name__)
+    
+    
+    @dataclass
+    @auto_docstring(
+        custom_intro="""
+        Base class for Gemma3n outputs, with hidden states and attentions.
+        """
+    )
+    class Gemma3nModelOutputWithPast(BaseModelOutputWithPast):
+        r"""
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+    
+            Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+            `past_key_values` input) to speed up sequential decoding.
+        image_hidden_states (`torch.FloatTensor`, *optional*):
+            A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+            image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+        audio_hidden_states (`torch.FloatTensor`, *optional*):
+            A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+            audio_hidden_states of the model produced by the audio encoder and after projecting the last hidden state.
+        """
+    
+        image_hidden_states: Optional[torch.FloatTensor] = None
+    
+        audio_hidden_states: Optional[torch.FloatTensor] = None
+    
+    
+    @dataclass
+    @auto_docstring(
+        custom_intro="""
+        Base class for Gemma3n causal language model (or autoregressive) outputs.
+        """
+    )
+    class Gemma3nCausalLMOutputWithPast(ModelOutput):
+        r"""
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (for next-token prediction).
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.text_config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+    
+            Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+            `past_key_values` input) to speed up sequential decoding.
+        image_hidden_states (`torch.FloatTensor`, *optional*):
+            A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+            image_hidden_states of the model produced by the vision encoder after projecting last hidden state.
+        audio_hidden_states (`torch.FloatTensor`, *optional*):
+            A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+            audio_hidden_states of the model produced by the audio encoder and after projecting the last hidden state.
+        """
+    
+        loss: Optional[torch.FloatTensor] = None
+        logits: Optional[torch.FloatTensor] = None
+        past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None
+        hidden_states: Optional[tuple[torch.FloatTensor]] = None
+        attentions: Optional[tuple[torch.FloatTensor]] = None
+        image_hidden_states: Optional[torch.FloatTensor] = None
+    
+        audio_hidden_states: Optional[torch.FloatTensor] = None
+    
+    
+    class Gemma3nRMSNorm(nn.Module):
+        def __init__(self, dim: int, eps: float = 1e-6, with_scale: bool = True):
+            super().__init__()
+            self.eps = eps
+            self.with_scale = with_scale
+    
+            if self.with_scale:
+                self.weight = nn.Parameter(torch.ones(dim))
+            else:
+                self.register_buffer("weight", torch.tensor(1.0), persistent=False)
+    
+        def _norm(self, x):
+            return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+    
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            # Llama does x.to(float16) * w whilst Gemma2 is (x * w).to(float16)
+            # See https://github.com/huggingface/transformers/pull/29402
+            output = self._norm(x.float()) * self.weight.float()
+            return output.type_as(x)
+    
+        def extra_repr(self):
+            return f"{tuple(self.weight.shape)}, eps={self.eps}"
+    
+    
+    # ==== Audio Encoder ====
+    
+    
+    class Gemma3nAudioRelativePositionEmbedding(nn.Module):
+        def __init__(self, config: Gemma3nAudioConfig):
+            super().__init__()
+            self.config = config
+    
+            self.num_heads = self.config.conf_num_attention_heads
+            self.channels = self.config.hidden_size
+            self.head_dim = self.channels // self.num_heads
+            self.max_backward = max(0, self.config.conf_attention_context_left - 1)
+            self.max_forward = self.config.conf_attention_context_right
+    
+            self.pos_proj = nn.Linear(self.channels, self.num_heads * self.head_dim, bias=False)
+    
+            min_timescale = 1.0
+            max_timescale = 1.0e4
+            num_timescales = self.channels // 2
+            log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / max(num_timescales - 1, 1)
+            inv_timescales = min_timescale * torch.exp(torch.arange(num_timescales) * -log_timescale_increment)
+            self.register_buffer(
+                "inv_timescales",
+                inv_timescales.float().unsqueeze(0).unsqueeze(0),
+                persistent=False,
+            )
+    
+        def _get_timing_signal_1d_pos(self, position: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+            position = position.float().unsqueeze(-1)
+            scaled_time = position * self.inv_timescales.to(device=position.device, dtype=torch.float32)
+            timing_signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=-1)
+            return timing_signal.type(dtype)
+    
+        def _relative_shift(
+            self,
+            term_bd_before_shift: torch.Tensor,
+            batch_size: int,
+            num_heads: int,
+            num_query_blocks: int,
+            query_block_size: int,
+            key_context_size: int,
+            max_span_plus_1: int,
+        ) -> torch.Tensor:
+            """Performs the relative shift.
+    
+            Args:
+              term_bd_before_shift: Tensor of shape [B, N, U, W, F_span]. batch_size
+                (B), num_heads (N), num_query_blocks (U), query_block_size (W),
+                key_context_size (C = W+L+R), max_span_plus_1 (F_span = L+R+1).
+    
+            Returns:
+              Tensor of shape [B, N, U, W, C].
+            """
+            # term_bd_before_shift shape: [B, N, U, W, F_span]
+            # Target shape after shift:  [B, N, U, W, C]
+    
+            # Padding amount for the last dimension (F_span) to become (C + 1)
+            # C = key_context_size
+            # F_span = max_span_plus_1
+            pad_amount_last_dim = (key_context_size + 1) - max_span_plus_1
+    
+            # PyTorch F.pad expects (pad_left, pad_right, pad_top, pad_bottom ...)
+            # We only pad the last dimension on the right.
+            padding_tuple = (0, pad_amount_last_dim)
+    
+            term_bd_padded = nn.functional.pad(term_bd_before_shift, padding_tuple)
+            # Shape after pad: [B, N, U, W, C+1]
+    
+            # Reshape for slicing (emulating JAX's behavior)
+            # [B, N, U, W * (C+1)]
+            term_bd_reshaped = term_bd_padded.reshape(
+                (
+                    batch_size,
+                    num_heads,
+                    num_query_blocks,
+                    query_block_size * (key_context_size + 1),
+                )
+            )
+    
+            # Slice to effective [B, N, U, W * C]
+            term_bd_sliced = term_bd_reshaped[:, :, :, : query_block_size * key_context_size]
+    
+            # Reshape back to [B, N, U, W, C]
+            term_bd_shifted = term_bd_sliced.reshape(
+                (
+                    batch_size,
+                    num_heads,
+                    num_query_blocks,
+                    query_block_size,
+                    key_context_size,
+                )
+            )
+            return term_bd_shifted
+    
+        def forward(self, queries: torch.Tensor, keys: torch.Tensor) -> torch.Tensor:
+            # queries: [B, U, W, N, H] (batch, num_query_blocks, query_block_size, num_heads, head_dim)
+            # keys:    [B, U, C, N, H] (batch, num_query_blocks, key_context_size, num_heads, head_dim)
+            # C = W + L + R (key_context_size)
+            # F_span = L + R + 1 (max_span + 1)
+    
+            batch_size, num_query_blocks, query_block_size, num_heads, head_dim = queries.shape
+            _, _, key_context_size, _, _ = keys.shape
+    
+            # Relative positions for sinusoidal embeddings: [L, L-1, ..., -R]
+            # Length is L+R+1 = self.max_span + 1
+            pos_indices = torch.arange(self.max_backward, -self.max_forward - 1, -1, device=queries.device).unsqueeze(
+                0
+            )  # Shape [1, F_span]
+    
+            max_span_plus_1 = pos_indices.shape[1]  # F_span
+    
+            sin_emb_timing_signal = self._get_timing_signal_1d_pos(
+                pos_indices, dtype=queries.dtype
+            )  # Shape [1, F_span, self.channels]
+    
+            # Project sinusoidal embeddings: [1, F_span, self.channels] -> [1, F_span, N*H]
+            projected_sin_emb = self.pos_proj(sin_emb_timing_signal)
+            # Reshape to [1, F_span, N, H] then squeeze to [F_span, N, H]
+            sin_emb = projected_sin_emb.reshape(1, max_span_plus_1, self.num_heads, self.head_dim).squeeze(
+                0
+            )  # Shape [F, N, H]
+    
+            # term_ac: Query-Key content interaction
+            # queries: [B, U, W, N, H] -> permute to [B, N, U, W, H] for matmul
+            # keys:    [B, U, C, N, H] -> permute to [B, N, U, H, C] for matmul
+            queries_p = queries.permute(0, 3, 1, 2, 4)  # [B, N, U, W, H]
+            keys_p_t = keys.permute(0, 3, 1, 4, 2)  # [B, N, U, H, C]
+            term_ac = torch.matmul(queries_p, keys_p_t)  # [B, N, U, W, C]
+    
+            # term_bd: Query-Position interaction
+            # Original einsum: term_bd_unshifed = torch.einsum('buwnh,fnh->bnuwf', queries, sin_emb)
+            # queries shape: [B, U, W, N, H]
+            # sin_emb shape: [F, N, H]
+            # Target output shape: [B, N, U, W, F]
+    
+            # Permute queries to [B, N, U, W, H] for easier broadcasting with sin_emb
+            q_permuted = queries.permute(0, 3, 1, 2, 4)
+    
+            # Permute sin_emb to [N, H, F] to prepare for matmul
+            # sin_emb original is [F, N, H]
+            s_permuted = sin_emb.permute(1, 2, 0)  # Shape: [N, H, F]
+    
+            # Reshape queries for matmul: [B, N, U*W, H]
+            q_reshaped = q_permuted.reshape(batch_size, num_heads, num_query_blocks * query_block_size, head_dim)
+    
+            # Perform matmul: [B, N, U*W, H] @ [N, H, F]
+            # s_permuted ([N, H, F]) will be broadcast to [B, N, H, F]
+            # Result: [B, N, U*W, F]
+            term_bd_unshifed_matmul = torch.matmul(q_reshaped, s_permuted)
+    
+            # Reshape to target [B, N, U, W, F]
+            term_bd_unshifed = term_bd_unshifed_matmul.reshape(
+                batch_size,
+                num_heads,
+                num_query_blocks,
+                query_block_size,
+                max_span_plus_1,
+            )
+    
+            # Apply relative shift to term_bd_unshifed
+            term_bd_shifted = self._relative_shift(
+                term_bd_unshifed,
+                batch_size,
+                num_heads,
+                num_query_blocks,
+                query_block_size,
+                key_context_size,
+                max_span_plus_1,
+            )  # Shape [B, N, U, W, C]
+    
+            return term_ac + term_bd_shifted
+    
+    
+    class Gemma3nAudioAttention(nn.Module):
+        def __init__(self, config: Gemma3nAudioConfig):
+            super().__init__()
+            self.config = config
+    
+            self.num_heads = self.config.conf_num_attention_heads
+            self.hidden_size = self.config.hidden_size
+            self.head_dim = self.hidden_size // self.num_heads
+    
+            self.chunk_size = self.config.conf_attention_chunk_size
+            self.max_future_horizon = self.config.conf_attention_context_right
+            self.max_past_horizon = max(0, self.config.conf_attention_context_left - 1)
+            self.attention_logits_soft_cap = self.config.conf_attention_logit_cap
+            self.context_size = self.chunk_size + self.max_past_horizon + self.max_future_horizon
+    
+            self.relative_position_embedding = Gemma3nAudioRelativePositionEmbedding(config)
+            self.per_dim_scale = nn.Parameter(torch.zeros((self.head_dim,)))
+    
+            self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+            self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+            self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+    
+            q_scale = self.head_dim**-0.5
+            r_softplus_0 = 1.0 / torch.nn.functional.softplus(torch.tensor(0.0))
+            self.register_buffer("q_scale", (q_scale * r_softplus_0).clone().detach(), persistent=False)
+    
+            lower_causal_mask = torch.tril(
+                torch.ones((self.context_size, self.chunk_size), dtype=torch.bool),
+                diagonal=0,
+            ).T
+            upper_causal_mask = torch.tril(
+                torch.ones((self.chunk_size, self.context_size), dtype=torch.bool),
+                diagonal=self.max_past_horizon + self.max_future_horizon,
+            )
+            local_causal_valid_mask = torch.ones((self.chunk_size, self.context_size), dtype=torch.bool)
+            local_causal_valid_mask = local_causal_valid_mask * lower_causal_mask * upper_causal_mask
+            self.register_buffer("local_causal_valid_mask", local_causal_valid_mask, persistent=False)
+    
+            self.register_buffer(
+                "softcap",
+                torch.tensor(self.attention_logits_soft_cap).float(),
+                persistent=False,
+            )
+    
+        def _pad_dim1(self, x: torch.Tensor, pad_left: int, pad_right: int) -> torch.Tensor:
+            batch, _, *tail_shape = x.shape
+            left = x.new_zeros((batch, pad_left, *tail_shape))
+            right = x.new_zeros((batch, pad_right, *tail_shape))
+            x = torch.cat([left, x, right], dim=1)
+            return x
+    
+        def _convert_to_block(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            """Turns a sequence to non overlapping blocks.
+    
+            Args:
+                hidden_states: a tensor of [batch, time, ...].
+    
+            Returns:
+                A tensor of [batch, num_blocks, block_size, ...], with necessary
+                paddings,
+                where output[:, i, ...] are x[:, i*block_size:(i+1)*block_size, ...].
+            """
+            shape = hidden_states.shape
+            b, t = shape[:2]
+            num_blocks = (t + self.chunk_size - 1) // self.chunk_size
+    
+            if (padding_len := num_blocks * self.chunk_size - t) > 0:
+                hidden_states = self._pad_dim1(hidden_states, 0, padding_len)
+    
+            permute_dims = (b, num_blocks, self.chunk_size) + shape[2:]
+            hidden_states = hidden_states.reshape(permute_dims).contiguous()
+            return hidden_states
+    
+        def _extract_block_context(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            """Extracts temporal context for every block.
+    
+            Args:
+                hidden_states: a tensor of [batch, time, ...].
+    
+            Returns:
+                A tensor of [batch, num_blocks, context_size, ...], with necessary
+                paddings,
+                where context_size = block_size + left_context + right_context,
+                and output[:, i, ...] are x[:, start-left_context:end+right_context,
+                ...],
+                start = i * block_size, end = (i + 1) * block_size.
+            """
+            pad_left = self.max_past_horizon
+            # The JAX equivalent padding for signal.frame with pad_mode='valid' is
+            # (left_context, right_context + block_size - 1) on the time dimension.
+            # PyTorch's _pad_dim1 applies padding symmetrically if only one value is given,
+            # or (pad_dim_start, pad_dim_end) if two are given.
+            # Our _pad_dim1(x, pad_left, pad_right) pads dim -2 (time for [B,T,N,H])
+            # or dim 1 (time for [B,T]).
+            # The current pad_right calculation matches the JAX effective padding.
+            pad_right = self.max_future_horizon + self.chunk_size - 1
+            hidden_states = self._pad_dim1(hidden_states, pad_left, pad_right)
+    
+            frame_len = self.context_size
+            frame_step = self.chunk_size
+    
+            # Directly use unfold without the subframe_factor logic
+            # x.unfold(dimension, size, step)
+            # dimension=1 (time dimension, assuming x is [B, T_padded, ...])
+            # size=frame_len (context_size)
+            # step=frame_step (chunk_size)
+            x_unfolded = hidden_states.unfold(dimension=1, size=frame_len, step=frame_step)
+    
+            # If x was [B, T_padded], x_unfolded is [B, num_blocks, frame_len]
+            # If x was [B, T_padded, N, H], x_unfolded is [B, num_blocks, N, H, frame_len]
+            # We want to match JAX's typical output for such operations which might be
+            # [B, num_blocks, frame_len, N, H] if N, H are present.
+            # The relative_position_embedding expects keys as [B, U, C, N, H].
+            # If x_unfolded is [B, U, N, H, C(frame_len)], we need to move C.
+            if hidden_states.ndim > 2 and x_unfolded.ndim > 3:  # Check if inner dimensions (like N, H) exist
+                # Current shape after unfold for [B, T_pad, N, H] is [B, U, N, H, C]
+                # Target shape for keys in RPE: [B, U, C, N, H]
+                x_unfolded = torch.movedim(x_unfolded, source=-1, destination=2)
+    
+            return x_unfolded.contiguous()
+    
+        def forward(self, hidden_states: torch.Tensor, mask: torch.BoolTensor) -> torch.Tensor:
+            # sl.Dense uses jax.numpy.einsum("...a,abcd->...bcd") and jax.numpy.select()
+            qkv_shape = (*hidden_states.shape[:-1], self.num_heads, self.head_dim)
+            query_states = self.q_proj(hidden_states).reshape(qkv_shape).contiguous()
+            key_states = self.k_proj(hidden_states).reshape(qkv_shape).contiguous()
+            value_states = self.v_proj(hidden_states).reshape(qkv_shape).contiguous()
+    
+            per_dim_scale_sp = torch.nn.functional.softplus(self.per_dim_scale)
+    
+            broadcast_shape = (1, 1, 1, self.head_dim)
+            per_dim_scale_sp_broadcast = per_dim_scale_sp.view(broadcast_shape)
+            query_states = query_states * self.q_scale * per_dim_scale_sp_broadcast
+    
+            batch_size, q_time = query_states.shape[:2]
+    
+            query_blocks = self._convert_to_block(query_states)
+            key_blocks = self._extract_block_context(key_states)
+            value_blocks = self._extract_block_context(value_states)
+            num_query_blocks = query_blocks.shape[1]
+    
+            # 1. Create a mask indicating originally valid positions.
+            original_valid_mask = ~mask  # True for valid, False for padded
+    
+            # 2. Extract blocks from this validity mask.
+            extracted_valid_mask_blocks = self._extract_block_context(original_valid_mask)
+    
+            # If subframe_factor was used in _extract_block_context for a [B, T] input mask,
+            # the shape might be [B, U, C/SF, SF]. Reshape to [B, U, C].
+            # batch_size and num_query_blocks are known from query_blocks.
+            # self.context_size is C.
+            if (
+                extracted_valid_mask_blocks.ndim == 4
+                and extracted_valid_mask_blocks.shape[2] * extracted_valid_mask_blocks.shape[3] == self.context_size
+            ):
+                extracted_valid_mask_blocks = extracted_valid_mask_blocks.reshape(
+                    batch_size, num_query_blocks, self.context_size
+                )
+            # After potential reshape, ensure it's [B, U, C] if it was from a [B,T] mask.
+            # This assertion might be too strict if _extract_block_context handles higher-rank inputs differently,
+            # but for the mask case, this should hold.
+            if extracted_valid_mask_blocks.shape != (
+                batch_size,
+                num_query_blocks,
+                self.context_size,
+            ):
+                raise ValueError(
+                    "Shape of extracted_valid_mask_blocks"
+                    f" {extracted_valid_mask_blocks.shape} is not ({batch_size},"
+                    f" {num_query_blocks}, {self.context_size}) after potential reshape."
+                )
+    
+            # 3. Expand dimensions for broadcasting with logits and causal mask.
+            # Target shape for broadcasting with logits [B,N,U,W,C]
+            # extracted_valid_mask_blocks to [B, 1, U, 1, C]
+            condition_from_input_validity = extracted_valid_mask_blocks.unsqueeze(1).unsqueeze(-2)
+    
+            # self.local_causal_valid_mask is [W, C], True where allowed by local window.
+            # Expand to [1, 1, 1, W, C]
+            condition_from_causality = self.local_causal_valid_mask.unsqueeze(0).unsqueeze(0).unsqueeze(0)
+    
+            # 4. Combine the two conditions.
+            # final_condition will be True where a key is *both* originally valid *and* causally accessible.
+            # Broadcasts to [B, 1, U, W, C]
+            final_condition_for_where = torch.logical_and(
+                condition_from_input_validity,
+                condition_from_causality.to(condition_from_input_validity.device),  # Ensure same device
+            )
+    
+            # Embed queries and keys
+            logits = self.relative_position_embedding(query_blocks, key_blocks)
+    
+            # Apply attention logit softcap
+            # Ensure softcap is on the same device as logits
+            softcap_val = self.softcap.to(logits.device)
+            logits = logits / softcap_val
+            logits = torch.tanh(logits)
+            logits = logits * softcap_val
+    
+            # Apply the combined mask.
+            # final_condition_for_where will broadcast with logits [B,N,U,W,C]
+            logits = torch.where(final_condition_for_where, logits, torch.finfo(logits.dtype).min)
+            probabilities = torch.nn.functional.softmax(logits, dim=-1, dtype=torch.float32).to(dtype=value_blocks.dtype)
+    
+            # context_vectors is adapted from jax.numpy.einsum("BNuwc,BucNH->BuwNH", ...)
+            b_dim, n_dim, u_dim, w_dim, c_dim = probabilities.shape
+            h_dim = value_blocks.shape[-1]
+            prob_bun = probabilities.permute(0, 2, 1, 3, 4).reshape(-1, w_dim, c_dim)
+            v_bun = value_blocks.permute(0, 1, 3, 2, 4).reshape(-1, c_dim, h_dim)
+            result_bmm = torch.bmm(prob_bun, v_bun)
+            context_vectors = result_bmm.reshape(b_dim, u_dim, n_dim, w_dim, h_dim).permute(0, 1, 3, 2, 4)
+            context_vectors = context_vectors.reshape(
+                (
+                    batch_size,
+                    num_query_blocks * self.chunk_size,
+                    self.num_heads,
+                    self.head_dim,
+                )
+            )
+            context_vectors = context_vectors[:, :q_time]
+    
+            return context_vectors
+    
+    
+    class Gemma3nAudioCumulativeGroupNorm(nn.Module):
+        """Applies Group Normalization cumulatively over the time dimension.
+    
+        This layer normalizes the input by calculating the mean and variance
+        cumulatively over the time dimension (dim 1). The statistics are computed
+        over all feature dimensions (specified by `feature_dims` and `num_channels`)
+        for elements marked as valid by the optional `mask`.
+    
+        If a `mask` is provided (True for valid, False for invalid/padded),
+        invalid time steps do not contribute to the statistics calculation, and
+        their corresponding output values are zeroed out.
+    
+        Scale and bias, if enabled, are applied per-channel (last dimension).
+        This behavior is similar to JAX's `GroupNormalization` with `num_groups=1`
+        and `cumulative=True`.
+        """
+    
+        def __init__(
+            self,
+            num_channels: int,  # Number of channels (size of the last dimension)
+            feature_dims: Sequence[int],  # Sizes of non-channel feature dimensions, e.g., (H, W) for input [B,T,H,W,C]
+            eps: float = 1e-3,
+        ):
+            super().__init__()
+            self.num_channels = num_channels
+            self.feature_dims = tuple(feature_dims)
+            self.eps = eps
+    
+            # Scale parameter depends only on the channel dimension
+            self.weight = nn.Parameter(torch.ones(num_channels))
+    
+            # Axes for normalization: all dimensions except Batch (0) and Time (1).
+            # For input [B, T, *feature_dims, C], these are dims from 2 onwards.
+            self.reduction_axes = tuple(range(2, 2 + len(self.feature_dims) + 1))
+    
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            """Applies cumulative group norm, optionally using a mask.
+    
+            Args:
+              hidden_states: Input tensor, shape [B, T, *feature_dims, C].
+    
+            Returns:
+              Normalized tensor with the same shape as x.
+            """
+            expected_input_suffix = self.feature_dims + (self.num_channels,)
+            if hidden_states.shape[2:] != expected_input_suffix:
+                raise ValueError(
+                    f"Input tensor shape suffix {hidden_states.shape[2:]} does not match expected"
+                    f" suffix (feature_dims + num_channels) {expected_input_suffix}"
+                )
+    
+            input_dtype = hidden_states.dtype
+            # Calculations are performed in float32 for numerical stability.
+            calc_dtype = torch.float32
+            x_calc = hidden_states.to(calc_dtype)
+    
+            # Prepare a broadcastable mask (`mask_calc`).
+            # If no mask is provided, treat all elements as valid
+            # (mask_calc is all ones).
+            # Otherwise, expand the [B, T] mask to [B, T, 1, ..., 1] for broadcasting.
+            mask_calc = torch.ones_like(x_calc, dtype=calc_dtype)
+    
+            # Cumulative Statistics Calculation
+            # 1. Sum of values over reduction axes at each time step.
+            sum_values_at_t = torch.sum(x_calc, dim=self.reduction_axes, keepdim=True)
+            # 2. Cumulative sum of values over time.
+            cum_sum_values = torch.cumsum(sum_values_at_t, dim=1)
+    
+            # 3. Count of valid elements in the normalization group at each time step.
+            #    (A "group" here consists of all features at a given Batch, Time).
+            elements_in_group_at_t = torch.sum(mask_calc, dim=self.reduction_axes, keepdim=True)
+            # 4. Cumulative count of valid elements over time.
+            cum_count_elements = torch.cumsum(elements_in_group_at_t, dim=1)
+            # Avoid division by zero if all preceding elements were masked.
+            safe_cum_count_elements = torch.clamp(cum_count_elements, min=1.0)
+    
+            # 5. Cumulative mean.
+            cum_mean = cum_sum_values / safe_cum_count_elements
+    
+            # 6. Sum of squared differences from the cumulative mean.
+            #    Only sum for valid elements: (x_calc - cum_mean)^2 * mask_calc.
+            #    Using x_calc here for the difference, as cum_mean already accounts for masking.
+            squared_diff_from_mean = (x_calc - cum_mean).pow(2)
+            sum_sq_diff_at_t = torch.sum(squared_diff_from_mean, dim=self.reduction_axes, keepdim=True)
+    
+            # 7. Cumulative sum of squared differences over time.
+            cum_sum_sq_diff = torch.cumsum(sum_sq_diff_at_t, dim=1)
+    
+            # 8. Cumulative variance.
+            cum_variance = cum_sum_sq_diff / safe_cum_count_elements
+    
+            # Normalize the input using the calculated cumulative statistics:
+            # (x - E[x]) / sqrt(Var[x] + eps)
+            normalized_x = (x_calc - cum_mean) * torch.rsqrt(cum_variance + self.eps)
+    
+            # Apply affine transformation (scale and bias) if enabled.
+            # Scale and bias are applied per-channel (last dimension).
+            scale = self.weight.to(calc_dtype)
+            # Reshape for broadcasting: [C] -> [1, ..., 1, C]
+            scale_view_shape = [1] * (hidden_states.dim() - 1) + [self.num_channels]
+            normalized_x = normalized_x * scale.view(scale_view_shape)
+    
+            # Zero out outputs for time steps that were originally masked (where mask_calc is 0).
+            # This ensures padded/invalid positions in the input result in zero output.
+            final_output = normalized_x * mask_calc
+    
+            return final_output.to(input_dtype)
+    
+    
+    class Gemma3nAudioSSCPConvBlock(nn.Module):
+        """A single convolution block for the SubSampleConvProjection.
+    
+        This block consists of a 2D convolution, followed by CumulativeGroupNorm,
+        and a ReLU activation. It handles manual padding for the convolution.
+        """
+    
+        def __init__(
+            self,
+            config: Gemma3nAudioConfig,
+            idx: int,
+            input_freq_dim: int,  # Changed from input_spatial_dim
+            manual_padding: tuple[int, int, int, int] = (0, 0, 0, 0),
+        ):
+            super().__init__()
+            self.config = config
+            self.manual_padding = manual_padding
+    
+            # in_channels is 1 for the first block, or C_out from previous block's conv
+            in_channels = 1 if idx == 0 else self.config.sscp_conv_channel_size[idx - 1]
+            out_channels = self.config.sscp_conv_channel_size[idx]
+            kernel_h, kernel_w = self.config.sscp_conv_kernel_size[idx]
+            stride_h, stride_w = self.config.sscp_conv_stride_size[idx]
+    
+            self.conv = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(
+                    kernel_h,
+                    kernel_w,
+                ),  # Kernel (kH, kW) operates on (Time, Freq_dim)
+                stride=(stride_h, stride_w),
+                padding=(0, 0),  # Manual padding is used
+                bias=False,
+            )
+    
+            # Calculate output frequency dimension (f_out_conv) after this convolution.
+            # input_freq_dim is the unpadded width (feature dimension).
+            # self.manual_padding is (pad_F_left, pad_F_right, pad_T_top, pad_T_bottom)
+            f_in_padded = input_freq_dim + self.manual_padding[0] + self.manual_padding[1]
+            f_out_conv = (f_in_padded - kernel_w) // stride_w + 1
+    
+            self.norm = Gemma3nAudioCumulativeGroupNorm(
+                num_channels=out_channels,  # Channels of the conv output
+                feature_dims=(f_out_conv,),  # The frequency dimension size after conv
+                eps=self.config.sscp_conv_group_norm_eps,
+            )
+    
+            self.activation = nn.ReLU()
+    
+        def forward(self, audio_encodings: torch.Tensor) -> torch.Tensor:
+            # Input audio_encodings is [B, C_in, T_in, F_in] (e.g., C_in=1)
+            # manual_padding is (pad_F_left, pad_F_right, pad_T_top, pad_T_bottom)
+            # F.pad applies to last two dims: F_in then T_in
+            audio_encodings_padded = F.pad(audio_encodings, self.manual_padding, mode="constant", value=0.0)
+            # Expected padded shape for F_in, k_w=3, pad_F=(1,1) -> F_padded = F_in+2
+            # Expected padded shape for T_in, k_h=3, pad_T=(0,2) -> T_padded = T_in+2
+            audio_encodings_conv = self.conv(audio_encodings_padded)
+            # Expected conv output shape: [B, C_out, T_out, F_out]
+            # Input to norm is [B, T_out, F_out, C_out]
+            x_for_norm = audio_encodings_conv.permute(0, 2, 3, 1).contiguous()
+            x_normed = self.norm(x_for_norm)
+            # Output of norm is [B, T_out, F_out, C_out], permute back to [B, C_out, T_out, F_out]
+            audio_encodings_normed = x_normed.permute(0, 3, 1, 2).contiguous()
+            return self.activation(audio_encodings_normed)
+    
+    
+    class Gemma3nAudioSubSampleConvProjection(nn.Module):
+        def __init__(self, config: Gemma3nAudioConfig):
+            super().__init__()
+            self.config = config
+    
+            current_f_for_block_input = config.input_feat_size  # Start with original feature dim
+            calculated_block_padding = []
+            calculated_f_out_dims = []  # Tracking frequency dimension output sizes
+    
+            for i in range(2):  # Assuming 2 conv layers as per sscp_conv_... arrays
+                kernel_h, kernel_w = config.sscp_conv_kernel_size[i]
+                stride_h, stride_w = config.sscp_conv_stride_size[i]
+    
+                # Padding for Time (Height for Conv2d) - REVERSE_CAUSAL like
+                # JAX 'reverse_causal' padding is (0, kernel_size - 1)
+                pad_t_top = 0
+                pad_t_bottom = kernel_h - 1
+    
+                # Frequency Padding (Width for Conv2d)
+                # Based on JAX effective padding (1,1) for F_in=10, K_w=3, S_w=2
+                # and the successful test configuration.
+                # If kernel/stride/input_freq for frequency changes, this might need re-evaluation
+                # to match generic JAX 'SAME' behavior if it differs.
+                pad_f_left = 1
+                pad_f_right = 1
+    
+                manual_padding_tuple = (
+                    pad_f_left,
+                    pad_f_right,
+                    pad_t_top,
+                    pad_t_bottom,
+                )
+                calculated_block_padding.append(manual_padding_tuple)
+    
+                # Calculate output frequency dimension after this convolution
+                # This uses the actual padding applied and kernel/stride.
+                f_in_padded = current_f_for_block_input + pad_f_left + pad_f_right
+                f_out_after_conv = (f_in_padded - kernel_w) // stride_w + 1  # Assuming dilation_w = 1
+                calculated_f_out_dims.append(f_out_after_conv)
+                current_f_for_block_input = f_out_after_conv
+    
+            self.conv_0 = Gemma3nAudioSSCPConvBlock(
+                idx=0,
+                input_freq_dim=config.input_feat_size,  # Pass original feature dim
+                config=config,
+                manual_padding=calculated_block_padding[0],
+            )
+            self.conv_1 = Gemma3nAudioSSCPConvBlock(
+                idx=1,
+                input_freq_dim=calculated_f_out_dims[0],  # Output freq dim from conv_0
+                config=config,
+                manual_padding=calculated_block_padding[1],
+            )
+            final_c_out = config.sscp_conv_channel_size[-1]
+            final_f_out = calculated_f_out_dims[-1]  # Final frequency dimension
+            self.input_proj_in_features = final_c_out * final_f_out
+            self.input_proj_linear = nn.Linear(self.input_proj_in_features, self.config.hidden_size, bias=False)
+    
+        def forward(self, audio_encodings: torch.Tensor) -> torch.Tensor:
+            # audio_encodings is [B, T, F_in]
+            # Reshape to [B, 1, T, F_in] (Batch, Channels=1, Height=Time, Width=F_in)
+            audio_encodings_reshaped = audio_encodings.unsqueeze(1)
+            x = self.conv_0(audio_encodings_reshaped)
+            x = self.conv_1(x)
+            # x from conv_1 is [B, C_out_1, T_out_1, F_out_1]
+            b, c_out, t_out, f_out = x.shape
+            # Permute to [B, T_out_1, F_out_1, C_out_1] then flatten F_out_1 and C_out_1
+            x_permuted = x.permute(0, 2, 3, 1).contiguous()
+            output_flattened = x_permuted.view(b, t_out, f_out * c_out)
+            output = self.input_proj_linear(output_flattened)
+            return output
+    
+    
+    class Gemma3nAudioConformerAttention(nn.Module):
+        def __init__(self, config: Gemma3nAudioConfig):
+            super().__init__()
+            self.config = config
+            self.post_in_features = self.config.hidden_size
+            self.register_buffer("gradient_clipping", torch.tensor(self.config.gradient_clipping), persistent=False)
+            self.pre_attn_norm = Gemma3nRMSNorm(self.config.hidden_size)
+            self.attn = Gemma3nAudioAttention(config)
+            self.post = nn.Linear(self.post_in_features, self.config.hidden_size, bias=False)
+            self.post_norm = Gemma3nRMSNorm(self.config.hidden_size)
+    
+        def forward(self, audio_encodings: torch.Tensor, audio_mel_mask: torch.BoolTensor) -> torch.Tensor:
+            audio_encodings_input_to_attn = audio_encodings
+            audio_encodings = torch.clamp(audio_encodings, -self.gradient_clipping, self.gradient_clipping)
+            audio_encodings_norm = self.pre_attn_norm(audio_encodings)
+            # Output of self.attn is [B, T, NumHeads, HeadDim]
+            audio_encodings_attn_out = self.attn(audio_encodings_norm, audio_mel_mask)
+    
+            # Reshape from [B, T, NumHeads, HeadDim] to [B, T, NumHeads * HeadDim]
+            # NumHeads * HeadDim = hidden_size
+            b, t, num_heads, head_dim = audio_encodings_attn_out.shape
+            audio_encodings_reshaped = audio_encodings_attn_out.reshape(b, t, num_heads * head_dim)
+    
+            audio_encodings = self.post(audio_encodings_reshaped)
+            audio_encodings = torch.clamp(audio_encodings, -self.gradient_clipping, self.gradient_clipping)
+            return audio_encodings_input_to_attn + self.post_norm(audio_encodings)
+    
+    
+    class Gemma3nAudioConformerFeedForward(nn.Module):
+        def __init__(self, config: Gemma3nAudioConfig):
+            super().__init__()
+            self.config = config
+    
+            self.register_buffer("gradient_clipping", torch.tensor(self.config.gradient_clipping), persistent=False)
+    
+            self.pre_layer_norm = Gemma3nRMSNorm(self.config.hidden_size)
+            self.ffw_layer_1 = nn.Linear(self.config.hidden_size, self.config.hidden_size * 4, bias=False)
+            self.ffw_layer_2 = nn.Linear(self.config.hidden_size * 4, self.config.hidden_size, bias=False)
+            self.post_layer_norm = Gemma3nRMSNorm(self.config.hidden_size)
+            self.post_layer_scale = torch.tensor(self.config.conf_residual_weight)
+    
+        def forward(self, audio_encodings: torch.Tensor) -> torch.Tensor:
+            residual = audio_encodings
+            audio_encodings = torch.clamp(audio_encodings, -self.gradient_clipping, self.gradient_clipping)
+            audio_encodings = self.pre_layer_norm(audio_encodings)
+            audio_encodings: torch.Tensor = self.ffw_layer_1(audio_encodings)
+            audio_encodings = nn.functional.silu(audio_encodings)
+            audio_encodings: torch.Tensor = self.ffw_layer_2(audio_encodings)
+            audio_encodings = torch.clamp(audio_encodings, -self.gradient_clipping, self.gradient_clipping)
+            audio_encodings = self.post_layer_norm(audio_encodings)
+            return residual + (audio_encodings * self.post_layer_scale)
+    
+    
+    class Gemma3nAudioConformerLightConv1d(nn.Module):
+        def __init__(self, config: Gemma3nAudioConfig):
+            super().__init__()
+            self.config = config
+    
+            self.pre_layer_norm = Gemma3nRMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+            self.linear_start = nn.Linear(self.config.hidden_size, self.config.hidden_size * 2, bias=False)
+            self.depthwise_conv1d = nn.Conv1d(
+                in_channels=self.config.hidden_size,
+                out_channels=self.config.hidden_size,
+                kernel_size=self.config.conf_conv_kernel_size,
+                stride=1,
+                padding=0,  # Manual causal padding
+                groups=self.config.hidden_size,  # Depthwise
+                bias=False,
+            )
+            self.register_buffer("gradient_clipping", torch.tensor(self.config.gradient_clipping), persistent=False)
+            self.conv_norm = Gemma3nRMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+            self.linear_end = nn.Linear(self.config.hidden_size, self.config.hidden_size, bias=False)
+    
+            self.causal_padding = self.config.conf_conv_kernel_size - 1
+    
+        def forward(self, audio_encodings: torch.Tensor) -> torch.Tensor:
+            audio_encodings_residual = audio_encodings  # Save for residual connection
+    
+            audio_encodings = self.pre_layer_norm(audio_encodings)
+            audio_encodings = self.linear_start(audio_encodings)
+            audio_encodings = torch.nn.functional.glu(audio_encodings, dim=-1)
+            # Permute for Conv1d: [B, T, D] -> [B, D, T]
+            audio_encodings_permuted = audio_encodings.permute(0, 2, 1)
+            # Apply manual causal padding
+            audio_encodings_permuted_padded = F.pad(audio_encodings_permuted, (self.causal_padding, 0))
+            audio_encodings = self.depthwise_conv1d(audio_encodings_permuted_padded)
+            # Permute back: [B, D, T_out] -> [B, T_out, D]
+            audio_encodings = audio_encodings.permute(0, 2, 1)
+            audio_encodings = torch.clamp(audio_encodings, -self.gradient_clipping, self.gradient_clipping)
+            audio_encodings = self.conv_norm(audio_encodings)
+            audio_encodings = nn.functional.silu(audio_encodings)
+            audio_encodings = self.linear_end(audio_encodings)
+            output = audio_encodings + audio_encodings_residual
+            return output
+    
+    
+    class Gemma3nAudioConformerBlock(nn.Module):
+        def __init__(self, config: Gemma3nAudioConfig):
+            super().__init__()
+            self.config = config
+    
+            self.ffw_layer_start = Gemma3nAudioConformerFeedForward(self.config)
+            self.attention = Gemma3nAudioConformerAttention(self.config)
+            self.lconv1d = Gemma3nAudioConformerLightConv1d(self.config)
+            self.ffw_layer_end = Gemma3nAudioConformerFeedForward(self.config)
+            self.register_buffer("gradient_clipping", torch.tensor(self.config.gradient_clipping), persistent=False)
+            self.norm = Gemma3nRMSNorm(self.config.hidden_size)
+    
+        def forward(self, audio_encodings: torch.Tensor, audio_mel_mask: torch.BoolTensor) -> torch.Tensor:
+            audio_encodings = self.ffw_layer_start(audio_encodings)
+            audio_encodings = self.attention(audio_encodings, audio_mel_mask)
+            validity_mask_for_lconv = ~audio_mel_mask  # True for valid
+            audio_encodings_for_lconv_input = audio_encodings * validity_mask_for_lconv.unsqueeze(-1).to(
+                audio_encodings.dtype
+            )
+            audio_encodings = self.lconv1d(audio_encodings_for_lconv_input)
+    
+            audio_encodings = self.ffw_layer_end(audio_encodings)
+            audio_encodings = torch.clamp(audio_encodings, -self.gradient_clipping, self.gradient_clipping)
+            output = self.norm(audio_encodings)
+            return output
+    
+    
+    class Gemma3nAudioEncoder(PreTrainedModel):
+        """An audio encoder based on the [Universal Speech Model](https://arxiv.org/abs/2303.01037) architecture."""
+    
+        config_class = Gemma3nAudioConfig
+    
+        main_input_name = "audio_mel"
+    
+        def __init__(self, config: Gemma3nAudioConfig):
+            super().__init__(config)
+            self.config = config
+    
+            self.subsample_conv_projection = Gemma3nAudioSubSampleConvProjection(config)
+            self.conformer = nn.ModuleList(
+                [Gemma3nAudioConformerBlock(config) for _ in range(config.conf_num_hidden_layers)]
+            )
+    
+        def forward(
+            self, audio_mel: torch.Tensor, audio_mel_mask: torch.BoolTensor
+        ) -> tuple[torch.Tensor, torch.BoolTensor]:
+            """Encodes a batch of MELs.
+    
+            Args:
+                audio_mel: a torch.Tensor of shape [batch, num_frames, num_channels,
+                  mel_bins].
+    
+            Returns:
+                audio_encodings: a torch.Tensor of shape
+                    `[batch_size, self.config.audio_soft_tokens_per_image,
+                    self.config.audio_config.hidden_size]`
+                audio_mel_mask: a torch.BoolTensor of shape [batch, num_frames].
+            """
+            audio_encodings = self.subsample_conv_projection(audio_mel)  # audio_encodings: [B, T_sub, D]
+    
+            # Subsample the input audio_mel_mask to match the time dimension of audio_encodings (T_sub)
+            t_sub = audio_encodings.shape[1]
+    
+            time_stride_product = 1
+            for stride_pair_idx in range(len(self.config.sscp_conv_stride_size)):
+                time_stride_product *= self.config.sscp_conv_stride_size[stride_pair_idx][0]
+    
+            # Create indices for gathering from the original mask.
+            # These indices map to original time steps corresponding to the start of each
+            # receptive field in the subsampled output.
+            indices = torch.arange(t_sub, device=audio_mel_mask.device) * time_stride_product
+            indices = torch.clamp(indices, max=audio_mel_mask.shape[1] - 1)  # Ensure indices are valid
+    
+            # Expand indices for batch compatibility if B > 1 and indices is 1D.
+            if audio_mel_mask.ndim > 1 and indices.ndim == 1:
+                indices = indices.unsqueeze(0).expand(audio_mel_mask.shape[0], -1)  # [B, T_sub]
+            elif (
+                audio_mel_mask.ndim == indices.ndim
+                and audio_mel_mask.shape[0] == 1
+                and indices.shape[0] != 1
+                and t_sub == indices.shape[0]
+            ):
+                # Handle case where B=1 but indices became [T_sub] instead of [1, T_sub]
+                indices = indices.unsqueeze(0)
+    
+            current_mask = torch.gather(audio_mel_mask, 1, indices)  # [B, T_sub]
+    
+            for block in self.conformer:
+                audio_encodings = block(audio_encodings, current_mask)  # Pass the processed mask
+    
+            if self.config.conf_reduction_factor > 1:
+                audio_encodings = audio_encodings[:, :: self.config.conf_reduction_factor]
+                # Reduce the mask as well
+                current_mask = current_mask[:, :: self.config.conf_reduction_factor]
+    
+            audio_encodings = audio_encodings.masked_fill(current_mask.unsqueeze(-1), 0.0)
+            return audio_encodings, current_mask
+    
+    
+    class Gemma3nTextScaledWordEmbedding(nn.Embedding):
+        """
+        This module overrides nn.Embeddings' forward by multiplying with embeddings scale.
+        """
+    
+        def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float = 1.0):
+            super().__init__(num_embeddings, embedding_dim, padding_idx)
+            self.register_buffer("embed_scale", torch.tensor(embed_scale), persistent=False)
+    
+        def forward(self, input_ids: torch.Tensor):
+            return super().forward(input_ids) * self.embed_scale.to(self.weight.dtype)
+    
+    
+    class Gemma3nTextLaurelBlock(nn.Module):
+        """Learned Augmented Residual Layer"""
+    
+        def __init__(self, config: Gemma3nTextConfig):
+            super().__init__()
+            self.config = config
+    
+            self.linear_left = nn.Linear(self.config.hidden_size, self.config.laurel_rank, bias=False)
+            self.linear_right = nn.Linear(self.config.laurel_rank, self.config.hidden_size, bias=False)
+            self.post_laurel_norm = Gemma3nRMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+    
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            laurel_hidden_states: torch.Tensor = self.linear_left(hidden_states)
+            laurel_hidden_states: torch.Tensor = self.linear_right(laurel_hidden_states)
+            normed_laurel_hidden_states = self.post_laurel_norm(laurel_hidden_states)
+            return hidden_states + normed_laurel_hidden_states
+    
+    
+    class Gemma3nTextMLP(nn.Module):
+        def __init__(self, config: Gemma3nTextConfig, layer_idx: int = 0):
+            super().__init__()
+            self.config = config
+            self.hidden_size = config.hidden_size
+            self.intermediate_size = config.intermediate_size[layer_idx]
+            self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+            self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+            self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+            self.act_fn = ACT2FN[config.hidden_activation]
+            self.activation_sparsity = config.activation_sparsity_pattern[layer_idx]
+    
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            gate_proj = self.gate_proj(hidden_states)
+            if self.activation_sparsity > 0.0:
+                gate_proj = self._gaussian_topk(gate_proj)
+            activations = self.act_fn(gate_proj)
+            up_proj = self.up_proj(hidden_states)
+            down_proj = self.down_proj(activations * up_proj)
+            return down_proj
+    
+        def _gaussian_topk(self, inputs: torch.Tensor) -> torch.Tensor:
+            target_sparsity_tensor = torch.tensor(self.activation_sparsity, dtype=torch.float32, device=inputs.device)
+            # normal_dist and std_multiplier are adapted from jax.scipy.stats.norm.ppf().
+            #
+            # References:
+            #   *   https://docs.jax.dev/en/latest/_autosummary/jax.scipy.stats.norm.ppf.html
+            #   *   https://pytorch.org/docs/stable/distributions.html#torch.distributions.normal.Normal
+            #   *   https://pytorch.org/docs/stable/distributions.html#torch.distributions.transformed_distribution.TransformedDistribution.icdf
+            normal_dist = torch.distributions.normal.Normal(0, 1)
+            std_multiplier: torch.Tensor = normal_dist.icdf(target_sparsity_tensor)
+            std_multiplier = std_multiplier.type(inputs.dtype)
+            inputs_mean = torch.mean(inputs, dim=-1, keepdim=True)
+            inputs_std = torch.std(inputs, dim=-1, keepdim=True, unbiased=False)
+            cutoff_x = inputs_mean + inputs_std * std_multiplier
+            return nn.functional.relu(inputs - cutoff_x)
+    
+    
+    class Gemma3nTextAltUp(nn.Module):
+        """Alternating Updates (AltUp)
+    
+        The AltUp module wraps transformer layers. The `predict` step modifies the
+        input to the transformer layer, and the `correct` step propagates the output
+        of the transformer layer to the sparsely updated dimensions.
+    
+        See more in the research paper:
+    
+        https://proceedings.neurips.cc/paper_files/paper/2023/file/f2059277ac6ce66e7e5543001afa8bb5-Paper-Conference.pdf
+        """
+    
+        def __init__(self, config: Gemma3nTextConfig):
+            super().__init__()
+            self.config = config
+            self.correct_output_scale = nn.Parameter(torch.zeros(self.config.hidden_size))
+            self.correction_coefs = nn.Linear(self.config.altup_num_inputs, self.config.altup_num_inputs, bias=False)
+            self.prediction_coefs = nn.Linear(self.config.altup_num_inputs, self.config.altup_num_inputs**2, bias=False)
+            self.modality_router = nn.Linear(self.config.hidden_size, self.config.altup_num_inputs, bias=False)
+            self.router_norm = Gemma3nRMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
+            self.register_buffer("router_input_scale", torch.tensor(self.config.hidden_size**-1.0), persistent=False)
+    
+        def compute_router_modalities(self, x: torch.Tensor) -> torch.Tensor:
+            router_inputs = self.router_norm(x) * self.router_input_scale
+            routed = self.modality_router(router_inputs)
+            return torch.tanh(routed.float()).type_as(x)
+    
+        def predict(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            """Predicts the output of a layer using a trainable map.
+    
+            Args:
+                hidden_states: A 4D tensor of shape `[num_altup_inputs, batch_size, num_tokens, hidden_size]` derived by
+                    stacking the input embeddings and preprocessing the last `num_altup_inputs - 1` matrices.
+    
+            Returns:
+                A 4D tensor of shape `[num_altup_inputs, batch_size, num_tokens, hidden_size]` containing the predictions.
+            """
+            modalities = self.compute_router_modalities(hidden_states[self.config.altup_active_idx])
+    
+            if self.training and self.config.altup_coef_clip is not None:
+                self.prediction_coefs.weight.data.clamp_(-self.config.altup_coef_clip, self.config.altup_coef_clip)
+    
+            # Project and then transpose all 2D matrices contained so that mulmat gives the correct result
+            all_coefs: torch.Tensor = (
+                self.prediction_coefs(modalities)
+                .reshape(*modalities.shape[:-1], self.config.altup_num_inputs, self.config.altup_num_inputs)
+                .permute(0, 1, 3, 2)
+            )
+    
+            # permute hidden_states to [batch_size, num_tokens, hidden_size, altup_num_inputs]
+            predictions = torch.matmul(hidden_states.permute(1, 2, 3, 0), all_coefs)
+            predictions = predictions.permute(3, 0, 1, 2)  # undo the permute
+            predictions += hidden_states  # add the original input
+            return predictions.contiguous().type_as(hidden_states)
+    
+        def correct(self, predictions: torch.Tensor, activated: torch.Tensor) -> torch.Tensor:
+            """Corrects the predictions relative to the
+    
+            Args:
+                predictions: A 4D tensor of shape `[num_altup_inputs, batch_size, num_tokens, hidden_size]` derived by
+                    stacking the input embeddings and preprocessing the last `num_altup_inputs - 1` matrices.
+                activated: A 3D tensor of shape `[batch_size, num_tokens, hidden_size]` containing the activated inputs.
+    
+            Returns:
+                A 4D tensor of shape `[num_altup_inputs, batch_size, num_tokens, hidden_size]` correcting the original
+                    predictions relative to the activated input embeddings.
+            """
+            modalities = self.compute_router_modalities(activated)
+            innovation = activated - predictions[self.config.altup_active_idx]  # (batch, num_tokens, hidden_size)
+            innovation = innovation.repeat(self.config.altup_num_inputs, 1, 1, 1)  # Repeat on dim0 to match predictions
+    
+            if self.config.altup_coef_clip is not None:
+                self.correction_coefs.weight.data.clamp_(-self.config.altup_coef_clip, self.config.altup_coef_clip)
+    
+            # all_coefs adapted from jax.numpy.einsum("...p,pi->...i", ...)
+            # Permute to (altup_num_inputs, batch_size, num_tokens) as the last dim is a scalar applied to each altup input
+            # and expand on dim1 for broadcastability
+            all_coefs: torch.Tensor = self.correction_coefs(modalities) + 1.0
+            all_coefs = all_coefs.permute(2, 0, 1).unsqueeze(-1)
+    
+            corrected = torch.mul(innovation, all_coefs)
+            corrected += predictions  # add the original input
+            return corrected.contiguous().type_as(activated)
+    
+        def forward(self, corrected: torch.Tensor) -> torch.Tensor:
+            """
+            This is only defined as the `forward` so that accelerate hooks can move correctly `correct_output_scale`
+            (which is a nn.Parameter, not a Module) between devices when offloading. It is otherwise only used in
+            `scale_corrected_output`
+            """
+            return (corrected.type_as(self.correct_output_scale) * self.correct_output_scale).type_as(corrected)
+    
+        def scale_corrected_output(self, corrected: torch.Tensor) -> torch.Tensor:
+            """Scales the provided 3D tensor of shape [batch_size, num_tokens, hidden_size]."""
+            return self.forward(corrected)
+    
+    
+    class Gemma3nTextRotaryEmbedding(nn.Module):
+        def __init__(self, config: Gemma3nTextConfig, device=None):
+            super().__init__()
+            # BC: "rope_type" was originally "type"
+            if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+                self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+            else:
+                self.rope_type = "default"
+            self.max_seq_len_cached = config.max_position_embeddings
+            self.original_max_seq_len = config.max_position_embeddings
+    
+            self.config = config
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+    
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
+            self.original_inv_freq = self.inv_freq
+    
+        @torch.no_grad()
+        @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+        def forward(self, x, position_ids):
+            inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+            position_ids_expanded = position_ids[:, None, :].float()
+    
+            device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+            with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+                freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+                emb = torch.cat((freqs, freqs), dim=-1)
+                cos = emb.cos() * self.attention_scaling
+                sin = emb.sin() * self.attention_scaling
+    
+            return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+    
+    
+    def rotate_half(x):
+        """Rotates half the hidden dims of the input."""
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+    
+    
+    def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+        """
+        This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+        num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+        """
+        batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+        if n_rep == 1:
+            return hidden_states
+        hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+        return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+    
+    
+    def eager_attention_forward(
+        module: nn.Module,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        dropout: float = 0.0,
+        scaling: Optional[float] = None,
+        softcap: Optional[float] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if scaling is None:
+            scaling = module.head_dim**-0.5
+    
+        key_states = repeat_kv(key, module.num_key_value_groups)
+        value_states = repeat_kv(value, module.num_key_value_groups)
+    
+        attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    
+        if softcap is not None:
+            attn_weights = attn_weights / softcap
+            attn_weights = torch.tanh(attn_weights)
+            attn_weights = attn_weights * softcap
+        if attention_mask is not None:  # no matter the length, we just slice it
+            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+            attn_weights = attn_weights + causal_mask
+    
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        return attn_output, attn_weights
+    
+    
+    def apply_rotary_pos_emb(
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        position_ids: Optional[torch.Tensor] = None,
+        unsqueeze_dim: int = 1,
+    ):
+        """Applies Rotary Position Embedding to the query and key tensors.
+    
+        Args:
+            x (`torch.Tensor`): The tensor to embed.
+            cos (`torch.Tensor`): The cosine part of the rotary embedding.
+            sin (`torch.Tensor`): The sine part of the rotary embedding.
+            position_ids (`torch.Tensor`, *optional*):
+                Deprecated and unused.
+            unsqueeze_dim (`int`, *optional*, defaults to 1):
+                The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+                sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+                that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+                k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+                cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+                the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+        Returns:
+            `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+        """
+        cos = cos.unsqueeze(unsqueeze_dim)
+        sin = sin.unsqueeze(unsqueeze_dim)
+        return (x * cos) + (rotate_half(x) * sin)
+    
+    
+    class Gemma3nTextAttention(nn.Module):
+        """Multi-headed attention from 'Attention Is All You Need' paper"""
+    
+        def __init__(self, config: Gemma3nTextConfig, layer_idx: int):
+            super().__init__()
+            self.is_sliding = config.layer_types[layer_idx] == "sliding_attention"
+            self.config = config
+            self.layer_idx = layer_idx
+            self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+            self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+            self.attention_dropout = self.config.attention_dropout
+            self.is_causal = True
+    
+            self.q_proj = nn.Linear(
+                config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+            )
+            self.k_proj = nn.Linear(
+                config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+            )
+            self.v_proj = nn.Linear(
+                config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+            )
+            self.o_proj = nn.Linear(
+                config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+            )
+            self.sliding_window = config.sliding_window if self.is_sliding else None
+    
+            self.q_norm = Gemma3nRMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = Gemma3nRMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)
+            self.v_norm = Gemma3nRMSNorm(dim=config.head_dim, eps=config.rms_norm_eps, with_scale=False)
+    
+            first_kv_shared_layer_idx = self.config.num_hidden_layers - self.config.num_kv_shared_layers
+            self.is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0
+            # Find the index of the last sliding or full layer before sharing starts (or None if no sharing)
+            layer_type = config.layer_types[layer_idx]
+            self.kv_shared_layer_index = (
+                first_kv_shared_layer_idx - 1 - config.layer_types[first_kv_shared_layer_idx - 1 :: -1].index(layer_type)
+                if self.is_kv_shared_layer
+                else None
+            )
+    
+        def forward(
+            self,
+            hidden_states: torch.Tensor,
+            position_embeddings: torch.Tensor,
+            attention_mask: Optional[torch.Tensor],
+            past_key_value: Optional[Cache] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            **kwargs: Unpack[FlashAttentionKwargs],
+        ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, self.config.head_dim)
+    
+            cos, sin = position_embeddings
+    
+            query_states = self.q_proj(hidden_states).view(hidden_shape)
+            query_states = self.q_norm(query_states)
+            query_states = apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=2)
+            query_states = query_states.transpose(1, 2)
+    
+            if self.is_kv_shared_layer and self.kv_shared_layer_index is not None and past_key_value is not None:
+                # Device of past layer may be different from current one
+                indices = cache_position.to(past_key_value.key_cache[self.kv_shared_layer_index].device)
+                # In this case we need special handling of the slice as the layer is of fixed small size (for full layers, we never go beyond)
+                if isinstance(past_key_value, HybridCache) and self.is_sliding:
+                    max_length = past_key_value.sliding_window
+                    indices = (
+                        slice(0, max_length)
+                        if cache_position.shape[0] > max_length
+                        else cache_position.clamp(min=0, max=max_length - 1)
+                    )
+    
+                # Device of past layer may be different from current one
+                key_states = past_key_value.key_cache[self.kv_shared_layer_index][:, :, indices].to(query_states.device)
+                value_states = past_key_value.value_cache[self.kv_shared_layer_index][:, :, indices].to(
+                    query_states.device
+                )
+            else:
+                key_states = self.k_proj(hidden_states).view(hidden_shape)
+                key_states = self.k_norm(key_states)
+                key_states = apply_rotary_pos_emb(key_states, cos, sin, unsqueeze_dim=2)
+                key_states = key_states.transpose(1, 2)
+    
+                value_states = self.v_proj(hidden_states).view(hidden_shape)
+                value_states = self.v_norm(value_states)
+                value_states = value_states.transpose(1, 2)
+    
+            if past_key_value is not None:
+                # sin and cos are specific to RoPE models; cache_position needed for the static cache
+                cache_kwargs = {
+                    "sin": sin,
+                    "cos": cos,
+                    "cache_position": cache_position,
+                    "sliding_window": self.sliding_window,
+                }
+                key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+    
+            attention_interface: Callable = eager_attention_forward
+            if self.config._attn_implementation != "eager":
+                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+    
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=self.attention_dropout if self.training else 0.0,
+                scaling=1.0,
+                sliding_window=self.sliding_window,
+                **kwargs,
+            )
+    
+            attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+            attn_output = self.o_proj(attn_output)
+            return attn_output, attn_weights
+    
+    
+    class Gemma3nTextDecoderLayer(GradientCheckpointingLayer):
+        def __init__(self, config: Gemma3nTextConfig, layer_idx: int):
+            super().__init__()
+            self.config = config
+            self.hidden_size = config.hidden_size
+            self.layer_idx = layer_idx
+            self.attention_type = config.layer_types[layer_idx]
+            self.self_attn = Gemma3nTextAttention(config, layer_idx)
+            self.mlp = Gemma3nTextMLP(config, layer_idx=layer_idx)
+            self.input_layernorm = Gemma3nRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
+            self.post_attention_layernorm = Gemma3nRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
+            self.pre_feedforward_layernorm = Gemma3nRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
+            self.post_feedforward_layernorm = Gemma3nRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
+    
+            self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
+            self.act_fn = ACT2FN[config.hidden_activation]
+    
+            self.altup = Gemma3nTextAltUp(config)
+            self.laurel = Gemma3nTextLaurelBlock(config)
+            self.per_layer_input_gate = nn.Linear(self.hidden_size, self.hidden_size_per_layer_input, bias=False)
+            self.per_layer_projection = nn.Linear(self.hidden_size_per_layer_input, self.hidden_size, bias=False)
+            self.post_per_layer_input_norm = Gemma3nRMSNorm(self.hidden_size, eps=config.rms_norm_eps)
+    
+        @deprecate_kwarg("last_cache_position", version="4.53.0")
+        def forward(
+            self,
+            hidden_states: torch.Tensor,
+            position_embeddings_global: torch.Tensor,
+            position_embeddings_local: torch.Tensor,
+            per_layer_input: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_value: Optional[Cache] = None,
+            output_attentions: Optional[bool] = False,
+            use_cache: Optional[bool] = False,
+            cache_position: Optional[torch.LongTensor] = None,
+            **kwargs,
+        ) -> tuple[torch.Tensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+            predictions = self.altup.predict(hidden_states)
+            active_prediction = predictions[self.config.altup_active_idx]
+    
+            active_prediction_normed = self.input_layernorm(active_prediction)
+            laurel_output = self.laurel(active_prediction_normed)
+    
+            # apply global RoPE to non-sliding layer only
+            if self.self_attn.is_sliding:
+                position_embeddings = position_embeddings_local
+            else:
+                position_embeddings = position_embeddings_global
+    
+            attn, self_attn_weights = self.self_attn(
+                hidden_states=active_prediction_normed,
+                position_embeddings=position_embeddings,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                **kwargs,
+            )
+            attn = self.post_attention_layernorm(attn)
+    
+            attn_gated = active_prediction + attn
+            attn_laurel = (attn_gated + laurel_output) / math.sqrt(2)
+    
+            attn_norm = self.pre_feedforward_layernorm(attn_laurel)
+            attn_ffw = self.mlp(attn_norm)
+            attn_ffw_norm = self.post_feedforward_layernorm(attn_ffw)
+            attn_ffw_laurel_gated = attn_laurel + attn_ffw_norm
+            corrected_predictions = self.altup.correct(predictions, attn_ffw_laurel_gated)
+    
+            first_prediction = corrected_predictions[self.config.altup_active_idx].clone()
+            if self.config.altup_correct_scale:
+                first_prediction = self.altup.scale_corrected_output(first_prediction)
+    
+            # per_layer_input_gate adapted from jax.numpy.einsum("btd,dp->btp", ...)
+            first_prediction = self.per_layer_input_gate(first_prediction)
+            first_prediction = self.act_fn(first_prediction)
+            first_prediction = torch.multiply(first_prediction, per_layer_input)
+    
+            # per_layer_projection adapted from jax.numpy.einsum("btp,pd->btd", ...)
+            first_prediction = self.per_layer_projection(first_prediction)
+            first_prediction = self.post_per_layer_input_norm(first_prediction)
+            corrected_predictions[1:] += first_prediction
+    
+            outputs = (corrected_predictions,)
+    
+            if output_attentions:
+                outputs += (self_attn_weights,)
+    
+            return outputs
+    
+    
+    @auto_docstring
+    class Gemma3nPreTrainedModel(PreTrainedModel):
+        config_class = Gemma3nConfig
+        base_model_prefix = ""
+        supports_gradient_checkpointing = True
+        _no_split_modules = ["Gemma3nTextDecoderLayer"]
+        _skip_keys_device_placement = ["past_key_values"]
+        _supports_flash_attn_3 = True
+        _supports_flash_attn_2 = True
+        _supports_sdpa = True
+        _supports_flex_attn = True
+        _supports_cache_class = True
+        _supports_quantized_cache = True
+        _supports_static_cache = True
+        _supports_attention_backend = True
+    
+        def _init_weights(self, module):
+            # important: this ported version of Gemma2 isn't meant for training from scratch - only
+            # inference and fine-tuning - so the proper init weights code has been removed
+            std = getattr(self.config, "initializer_range", self.config.get_text_config().initializer_range)
+    
+            if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d)):
+                module.weight.data.normal_(mean=0.0, std=std)
+                if module.bias is not None:
+                    module.bias.data.zero_()
+            elif isinstance(module, nn.Embedding):
+                module.weight.data.normal_(mean=0.0, std=std)
+                if module.padding_idx is not None:
+                    module.weight.data[module.padding_idx].zero_()
+            elif isinstance(module, Gemma3nRMSNorm):
+                if module.with_scale:
+                    module.weight.data.fill_(1.0)
+            elif isinstance(module, Gemma3nAudioCumulativeGroupNorm):
+                module.weight.data.fill_(1.0)
+            elif isinstance(module, Gemma3nAudioAttention):
+                module.per_dim_scale.data.zero_()
+            elif isinstance(module, Gemma3nTextAltUp):
+                module.correct_output_scale.data.zero_()
+    
+    
+    @auto_docstring(custom_intro="The base Gemma 3n language model without a language modeling head.")
+    class Gemma3nTextModel(Gemma3nPreTrainedModel):
+        config_class = Gemma3nTextConfig
+    
+        def __init__(self, config: Gemma3nTextConfig):
+            super().__init__(config)
+            self.padding_idx = config.pad_token_id
+            self.vocab_size = config.vocab_size
+    
+            # Gemma3n downcasts the below to bfloat16, causing sqrt(3072)=55.4256 to become 55.5. See https://github.com/huggingface/transformers/pull/29402
+            self.embed_tokens = Gemma3nTextScaledWordEmbedding(
+                config.vocab_size, config.hidden_size, self.padding_idx, embed_scale=self.config.hidden_size**0.5
+            )
+            self.layers = nn.ModuleList(
+                [Gemma3nTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+            )
+    
+            self.norm = Gemma3nRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+            self.rotary_emb = Gemma3nTextRotaryEmbedding(config=config)
+            self.gradient_checkpointing = False
+    
+            # TODO (raushan): Fix this after RoPE refactor. For now we hack it by
+            # reassigning thetas when we want to create a local RoPE layer. Config
+            # defaults should hold values for global RoPE.
+            config = copy.deepcopy(config)
+            config.rope_theta = config.rope_local_base_freq
+            config.rope_scaling = {"rope_type": "default"}
+            self.rotary_emb_local = Gemma3nTextRotaryEmbedding(config=config)
+    
+            self.hidden_size = config.hidden_size
+            self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
+    
+            self.embed_tokens_per_layer = Gemma3nTextScaledWordEmbedding(
+                config.vocab_size_per_layer_input,
+                config.num_hidden_layers * config.hidden_size_per_layer_input,
+                self.padding_idx,
+                embed_scale=config.hidden_size_per_layer_input**0.5,
+            )
+    
+            self.per_layer_model_projection = nn.Linear(
+                self.hidden_size,
+                config.num_hidden_layers * config.hidden_size_per_layer_input,
+                bias=False,
+            )
+    
+            self.per_layer_projection_norm = Gemma3nRMSNorm(config.hidden_size_per_layer_input, eps=config.rms_norm_eps)
+    
+            self.altup_projections = nn.ModuleList(
+                [nn.Linear(self.hidden_size, self.hidden_size, bias=False) for _ in range(1, self.config.altup_num_inputs)]
+            )
+    
+            self.altup_unembed_projections = nn.ModuleList(
+                [nn.Linear(self.hidden_size, self.hidden_size, bias=False) for _ in range(1, self.config.altup_num_inputs)]
+            )
+    
+            self.register_buffer("per_layer_projection_scale", torch.tensor(self.hidden_size**-0.5), persistent=False)
+            self.register_buffer("per_layer_input_scale", torch.rsqrt(torch.tensor(2.0)), persistent=False)
+    
+            # Initialize weights and apply final processing
+            self.post_init()
+    
+        def get_input_embeddings(self):
+            return self.embed_tokens
+    
+        def set_input_embeddings(self, value):
+            self.embed_tokens = value
+    
+        @can_return_tuple
+        @auto_docstring
+        def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            per_layer_inputs: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Cache] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+        ) -> BaseModelOutputWithPast:
+            r"""
+            per_layer_inputs (torch.Tensor, *optional*, defaults to None):
+                Pre-computed per-layer embeddings. If None, they are derived from input_ids if provided.
+            """
+            output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+            output_hidden_states = (
+                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+            )
+            use_cache = use_cache if use_cache is not None else self.config.use_cache
+    
+            if (input_ids is None) ^ (inputs_embeds is not None):
+                raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+    
+            if self.gradient_checkpointing and self.training and use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+                )
+                use_cache = False
+    
+            if input_ids is not None:
+                inputs_embeds = self.embed_tokens(input_ids)
+                per_layer_inputs = self.get_per_layer_inputs(input_ids)
+    
+            per_layer_inputs = self.project_per_layer_inputs(inputs_embeds, per_layer_inputs)
+    
+            if use_cache and past_key_values is None and not self.training:
+                past_key_values = DynamicCache()
+    
+            if cache_position is None:
+                past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+                cache_position = torch.arange(
+                    past_seen_tokens,
+                    past_seen_tokens + inputs_embeds.shape[1],
+                    device=inputs_embeds.device,
+                )
+    
+            if position_ids is None:
+                position_ids = cache_position.unsqueeze(0)
+    
+            # It may already have been prepared by e.g. `generate`
+            if not isinstance(causal_mask_mapping := attention_mask, dict):
+                # Prepare mask arguments
+                mask_kwargs = {
+                    "config": self.config,
+                    "input_embeds": inputs_embeds,
+                    "attention_mask": attention_mask,
+                    "cache_position": cache_position,
+                    "past_key_values": past_key_values,
+                }
+                # Create the masks
+                causal_mask_mapping = {
+                    "full_attention": create_causal_mask(**mask_kwargs),
+                    "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
+                }
+    
+            # embed positions
+            hidden_states_0 = inputs_embeds
+    
+            # Initialize RoPE embeddings
+            position_embeddings_global = self.rotary_emb(hidden_states_0, position_ids)
+            position_embeddings_local = self.rotary_emb_local(hidden_states_0, position_ids)
+    
+            # Expand hidden_states to support per-layer inputs
+            target_magnitude = torch.mean(hidden_states_0**2, dim=-1, keepdim=True) ** 0.5
+            epsilon_tensor = torch.tensor(1e-5)
+    
+            temp_hidden_states = [hidden_states_0]
+            for i in range(1, self.config.altup_num_inputs):
+                # altup_proj adapted from jax.numpy.einsum("btp,pd->btd", ...)
+                altup_proj = self.altup_projections[i - 1](hidden_states_0)
+                current_hidden_state = altup_proj.to(dtype=hidden_states_0.dtype, device=target_magnitude.device)
+                new_magnitude = torch.mean(current_hidden_state**2, dim=-1, keepdim=True)
+                new_magnitude = torch.sqrt(torch.maximum(new_magnitude, epsilon_tensor.to(target_magnitude.device)))
+                current_hidden_state = current_hidden_state * target_magnitude / new_magnitude
+                temp_hidden_states.append(current_hidden_state)
+    
+            hidden_states = torch.stack(temp_hidden_states, dim=0)  # [num_altup_inputs, batch, seq_len, hidden_size]
+    
+            # decoder layers
+            all_hidden_states = () if output_hidden_states else None
+            all_self_attns = () if output_attentions else None
+    
+            for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+                if output_hidden_states:
+                    all_hidden_states += (hidden_states,)
+    
+                causal_mask = causal_mask_mapping[decoder_layer.attention_type]
+                per_layer_input = per_layer_inputs[:, :, decoder_layer.layer_idx, :]
+    
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    position_embeddings_global,
+                    position_embeddings_local,
+                    per_layer_input,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    **flash_attn_kwargs,
+                )
+    
+                hidden_states = layer_outputs[0]
+    
+                if output_attentions:
+                    all_self_attns += (layer_outputs[1],)
+    
+            # add hidden states from the last decoder layer (but before reprojecting to stay consistent with layer output)
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+    
+            # Per-layer inputs to single output
+            target_magnitude = torch.mean(hidden_states[0] ** 2, dim=-1, keepdim=True) ** 0.5
+            temp_hidden_states = [hidden_states[0]]
+            for i in range(1, self.config.altup_num_inputs):
+                # altup_unembed_projections adapted from jax.numpy.einsum("btp,pd->btd", ...)
+                altup_unemb_proj: torch.Tensor = self.altup_unembed_projections[i - 1](hidden_states[i])
+                current_hidden_state = altup_unemb_proj.to(dtype=hidden_states_0.dtype, device=target_magnitude.device)
+                new_magnitude = torch.mean(current_hidden_state**2, dim=-1, keepdim=True)
+                new_magnitude = torch.sqrt(torch.maximum(new_magnitude, epsilon_tensor.to(target_magnitude.device)))
+                current_hidden_state = current_hidden_state * target_magnitude / new_magnitude
+                temp_hidden_states.append(current_hidden_state)
+    
+            hidden_states = torch.stack(temp_hidden_states)
+            hidden_states = torch.mean(hidden_states, dim=0)
+            hidden_states = self.norm(hidden_states)
+    
+            return BaseModelOutputWithPast(
+                last_hidden_state=hidden_states,
+                past_key_values=past_key_values,
+                hidden_states=all_hidden_states,
+                attentions=all_self_attns,
+            )
+    
+        def get_per_layer_inputs(self, input_ids: torch.LongTensor) -> torch.Tensor:
+            return self.embed_tokens_per_layer(input_ids).reshape(
+                *input_ids.shape,
+                self.config.num_hidden_layers,
+                self.hidden_size_per_layer_input,
+            )
+    
+        def project_per_layer_inputs(
+            self,
+            inputs_embeds: torch.Tensor,
+            per_layer_inputs: Optional[torch.Tensor] = None,
+        ) -> torch.Tensor:
+            per_layer_projection: torch.Tensor = self.per_layer_model_projection(inputs_embeds)
+            per_layer_projection *= self.per_layer_projection_scale.to(
+                dtype=inputs_embeds.dtype, device=per_layer_projection.device
+            )
+            per_layer_projection = per_layer_projection.reshape(
+                *inputs_embeds.shape[:-1],
+                self.config.num_hidden_layers,
+                self.hidden_size_per_layer_input,
+            )
+            per_layer_projection = self.per_layer_projection_norm(per_layer_projection)
+    
+            if per_layer_inputs is None:
+                return per_layer_projection
+    
+            if per_layer_projection.shape != per_layer_inputs.shape:
+                # per-layer inputs are sometimes padded with zeros, slice the relevant embeddings.
+                per_layer_inputs = per_layer_inputs[..., : self.config.num_hidden_layers, :]
+    
+            return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale.to(
+                dtype=inputs_embeds.dtype, device=per_layer_projection.device
+            )
+    
+    
+    @auto_docstring(custom_intro="The base Gemma 3n language model with a language modeling head.")
+    class Gemma3nForCausalLM(Gemma3nPreTrainedModel, GenerationMixin):
+        _tied_weights_keys = ["lm_head.weight"]
+        _tp_plan = {"lm_head": "colwise_rep"}
+        _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+        config_class = Gemma3nTextConfig
+        base_model_prefix = "model"
+        _checkpoint_conversion_mapping = {"model.language_model": "model"}
+    
+        def __init__(self, config: Gemma3nTextConfig):
+            super().__init__(config)
+            self.model = Gemma3nTextModel(config)
+            self.vocab_size = config.vocab_size
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+    
+            # Initialize weights and apply final processing
+            self.post_init()
+    
+        def get_input_embeddings(self):
+            return self.model.embed_tokens
+    
+        def set_input_embeddings(self, value):
+            self.model.embed_tokens = value
+    
+        def get_output_embeddings(self):
+            return self.lm_head
+    
+        def set_output_embeddings(self, new_embeddings):
+            self.lm_head = new_embeddings
+    
+        def set_decoder(self, decoder):
+            self.model = decoder
+    
+        def get_decoder(self):
+            return self.model
+    
+        @can_return_tuple
+        @auto_docstring
+        def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Cache] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            logits_to_keep: Union[int, torch.Tensor] = 0,
+            **loss_kwargs,
+        ) -> CausalLMOutputWithPast:
+            r"""
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+    
+            Example:
+    
+            ```python
+            >>> from transformers import AutoTokenizer, Gemma3nForCausalLM
+    
+            >>> model = Gemma3nForCausalLM.from_pretrained("google/gemma-2-9b")
+            >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
+    
+            >>> prompt = "What is your favorite condiment?"
+            >>> inputs = tokenizer(prompt, return_tensors="pt")
+    
+            >>> # Generate
+            >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+            >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+            "What is your favorite condiment?"
+            ```"""
+    
+            if self.training and self.config._attn_implementation != "eager":
+                logger.warning_once(
+                    "It is strongly recommended to train Gemma3n models with the `eager` attention implementation "
+                    f"instead of `{self.config._attn_implementation}`. Use `eager` with `AutoModelForCausalLM.from_pretrained('<path-to-checkpoint>', attn_implementation='eager')`."
+                )
+            output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+            output_hidden_states = (
+                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+            )
+            # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+            outputs: BaseModelOutputWithPast = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                cache_position=cache_position,
+                **loss_kwargs,
+            )
+    
+            hidden_states = outputs.last_hidden_state
+            # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+            logits = self.lm_head(hidden_states[:, slice_indices, :])
+            if self.config.final_logit_softcapping is not None:
+                logits = logits / self.config.final_logit_softcapping
+                logits = torch.tanh(logits)
+                logits = logits * self.config.final_logit_softcapping
+    
+            loss = None
+            if labels is not None:
+                loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
+    
+            return CausalLMOutputWithPast(
+                loss=loss,
+                logits=logits,
+                past_key_values=outputs.past_key_values,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+            )
+    
+    
+    class Gemma3nMultimodalEmbedder(nn.Module):
+        """Embeds token ids or soft tokens for multimodal content into language model space."""
+    
+        def __init__(
+            self,
+            multimodal_config: Union[Gemma3nAudioConfig, Gemma3nVisionConfig],
+            text_config: Gemma3nTextConfig,
+        ):
+            super().__init__()
+    
+            self.multimodal_hidden_size = multimodal_config.hidden_size
+            self.eps = multimodal_config.rms_norm_eps
+            self.vocab_offset = multimodal_config.vocab_offset
+            self.vocab_size = multimodal_config.vocab_size
+            self.text_hidden_size = text_config.hidden_size
+    
+            self.embedding = nn.Embedding(self.vocab_size, self.multimodal_hidden_size)
+            self.hard_embedding_norm = Gemma3nRMSNorm(self.multimodal_hidden_size, eps=self.eps)
+            self.soft_embedding_norm = Gemma3nRMSNorm(self.multimodal_hidden_size, eps=self.eps)
+            self.embedding_projection = nn.Linear(self.multimodal_hidden_size, self.text_hidden_size, bias=False)
+            self.embedding_post_projection_norm = Gemma3nRMSNorm(self.text_hidden_size, eps=self.eps, with_scale=False)
+    
+        def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            inputs_embeds: Optional[torch.Tensor] = None,
+        ) -> torch.Tensor:
+            """Embeds token ids or soft tokens for multimodal content into language model space.
+    
+            Args:
+                input_ids: A torch.LongTensor containing the token ids to embed. Values should be in the range
+                    `[vocab_offset, vocab_offset + vocab_size)`.
+                inputs_embeds: A torch.Tensor containing the soft tokens to embed.
+    
+            Returns:
+                A torch.Tensor of embeddings with  shape `[batch_size, seq_len, self.config.text_config.hidden_size]`.
+            """
+            if (input_ids is None) ^ (inputs_embeds is not None):
+                raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+    
+            if inputs_embeds is not None:
+                emb_norm = self.soft_embedding_norm(inputs_embeds)
+            else:
+                hard_emb = self.embedding(input_ids - self.vocab_offset)
+                emb_norm = self.hard_embedding_norm(hard_emb)
+    
+            emb_norm_proj = self.embedding_projection(emb_norm)
+            return self.embedding_post_projection_norm(emb_norm_proj)
+    
+    
+    @auto_docstring(
+        custom_intro="""
+        The base Gemma 3n model comprising a vision backbone, an audio backbone, and a language model without a
+        language modeling head.
+        """
+    )
+    class Gemma3nModel(Gemma3nPreTrainedModel):
+        _checkpoint_conversion_mapping = {}
+        # we are filtering the logits/labels so we shouldn't divide the loss based on num_items_in_batch
+        accepts_loss_kwargs = False
+    
+        def __init__(self, config: Gemma3nConfig):
+            super().__init__(config)
+            self.vision_tower = AutoModel.from_config(config=config.vision_config)
+            self.vocab_size = config.text_config.vocab_size
+    
+            language_model = AutoModel.from_config(config=config.text_config)
+            self.language_model = language_model
+    
+            self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
+            self.vocab_size_per_layer_input = config.text_config.vocab_size_per_layer_input
+            self.audio_tower = AutoModel.from_config(config.audio_config)
+            self.embed_vision = Gemma3nMultimodalEmbedder(config.vision_config, config.text_config)
+            self.embed_audio = Gemma3nMultimodalEmbedder(config.audio_config, config.text_config)
+            self.post_init()
+    
+        def get_input_embeddings(self):
+            return self.language_model.get_input_embeddings()
+    
+        def set_input_embeddings(self, value):
+            self.language_model.set_input_embeddings(value)
+    
+        def set_decoder(self, decoder):
+            self.language_model = decoder
+    
+        def get_decoder(self):
+            return self.language_model
+    
+        def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+            """
+            Projects the last hidden state from the vision model into language model space.
+    
+            Args:
+                pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
+                   The tensors corresponding to the input images.
+    
+            Returns:
+                image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
+            """
+            vision_outputs = self.vision_tower(
+                pixel_values=pixel_values, do_pooling=False, return_dict=True
+            ).last_hidden_state
+            # Convert from (batch, channels, height, width) to (batch, height * width, channels) where:
+            # height == width and height * width == Gemma3nConfig.vision_soft_tokens_per_image.
+            vision_outputs = vision_outputs.reshape(
+                vision_outputs.shape[0],
+                self.config.vision_config.hidden_size,
+                self.config.vision_soft_tokens_per_image,
+            ).permute(0, 2, 1)
+            # Normalize and embed the soft tokens into language model space.
+            vision_outputs *= self.config.vision_config.hidden_size**0.5
+            return self.embed_vision(inputs_embeds=vision_outputs)
+    
+        @can_return_tuple
+        def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,  # text inputs
+            pixel_values: Optional[torch.FloatTensor] = None,  # vision inputs
+            input_features: Optional[torch.FloatTensor] = None,  # audio inputs
+            attention_mask: Optional[torch.Tensor] = None,
+            input_features_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None,
+            token_type_ids: Optional[torch.LongTensor] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            **lm_kwargs,
+        ) -> Gemma3nCausalLMOutputWithPast:
+            r"""
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`.
+    
+            Example:
+    
+            ```python
+            >>> from PIL import Image
+            >>> import requests
+            >>> from transformers import AutoProcessor, Gemma3nForConditionalGeneration
+    
+            >>> model = Gemma3nForConditionalGeneration.from_pretrained("google/gemma3n2-3b-mix-224")
+            >>> processor = AutoProcessor.from_pretrained("google/gemma3n2-3b-mix-224")
+    
+            >>> prompt = "Where is the cat standing?"
+            >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
+            >>> image = Image.open(requests.get(url, stream=True).raw)
+    
+            >>> inputs = processor(images=image, text=prompt,  return_tensors="pt")
+    
+            >>> # Generate
+            >>> generate_ids = model.generate(**inputs,)
+            >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+            "Where is the cat standing?\nsnow"
+            ```
+            """
+            if (input_ids is None) ^ (inputs_embeds is not None):
+                raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+    
+            output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+            output_hidden_states = (
+                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+            )
+    
+            if input_ids is not None:
+                inputs_embeds = self.get_input_embeddings()(input_ids)
+    
+                # Prepare per-layer inputs from inputs_ids
+                per_layer_inputs_mask = torch.logical_and(input_ids >= 0, input_ids < self.vocab_size_per_layer_input)
+                per_layer_inputs_tokens = torch.where(per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids))
+                per_layer_inputs = self.language_model.get_per_layer_inputs(per_layer_inputs_tokens)
+    
+                # Handle vision tokens (>= embed_vision.vocab_offset and < embed_audio.vocab_offset)
+                vision_mask = torch.logical_and(
+                    input_ids >= self.embed_vision.vocab_offset, input_ids < self.embed_audio.vocab_offset
+                )
+                dummy_vision_token_id = self.embed_vision.vocab_offset + self.embed_vision.vocab_size - 1
+                vision_input_ids = torch.where(vision_mask, input_ids, dummy_vision_token_id).to(inputs_embeds.device)
+                vision_embeds = self.embed_vision(input_ids=vision_input_ids)
+                expanded_vision_mask = vision_mask.unsqueeze(-1).expand_as(inputs_embeds)
+                inputs_embeds = torch.where(expanded_vision_mask, vision_embeds, inputs_embeds)
+    
+                # Handle audio tokens (>= embed_audio.vocab_offset)
+                audio_mask = input_ids >= self.embed_audio.vocab_offset
+                dummy_audio_token_id = self.embed_audio.vocab_offset + self.embed_audio.vocab_size - 1
+                audio_input_ids = torch.where(audio_mask, input_ids, dummy_audio_token_id).to(inputs_embeds.device)
+                audio_embeds = self.embed_audio(input_ids=audio_input_ids)
+                expanded_audio_mask = audio_mask.unsqueeze(-1).expand_as(inputs_embeds)
+                inputs_embeds = torch.where(expanded_audio_mask, audio_embeds, inputs_embeds)
+            else:
+                per_layer_inputs = None
+    
+            # Merge text and images
+            if pixel_values is not None:
+                image_features = self.get_image_features(pixel_values)
+    
+                if input_ids is None:
+                    special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                        torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+                    )
+                else:
+                    special_image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1)
+                    special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+    
+                if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel():
+                    image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0]
+                    raise ValueError(
+                        f"Number of images does not match number of special image tokens in the input text. "
+                        f"Got {image_tokens_in_text} image tokens in the text and "
+                        f"{image_features.shape[0] * image_features.shape[1]} tokens from image embeddings."
+                    )
+                image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+                inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+    
+            # Merge text and audio
+            if input_features is not None and input_features_mask is not None:
+                audio_features, audio_mask = self.get_audio_features(input_features, ~input_features_mask)
+    
+                # The Gemma3nProcessor expects all audio will be 30s in length and inserts 188 audio soft tokens into the
+                # text to account for this. However, the audio preprocessing and encoder do not gurarantee they will
+                # produce 188 soft tokens; they will produce at most that many tokens, but they may produce fewer tokens
+                # depending on the length of the longest audio input in the batch. When we encounter this situation, we pad
+                # the audio feature out to 188 soft tokens with the emebedding of the last token in the embed_audio vocab.
+                audio_padding_toks = torch.tensor([[self.vocab_size - 1]], dtype=torch.long, device=audio_features.device)
+                audio_padding_embs = self.embed_audio(input_ids=audio_padding_toks)
+                audio_features = torch.where(audio_mask.unsqueeze(-1), audio_padding_embs, audio_features)
+    
+                audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape
+                extra_padding_tokens = self.config.audio_soft_tokens_per_image - audio_seq_len
+                extra_padding_features = audio_padding_embs.expand(audio_batch_size, extra_padding_tokens, audio_embed_dim)
+    
+                audio_features = torch.cat((audio_features, extra_padding_features), dim=1)
+    
+                if input_ids is None:
+                    special_audio_mask = inputs_embeds == self.embed_audio(
+                        input_ids=torch.tensor(self.config.audio_token_id, dtype=torch.long, device=inputs_embeds.device)
+                    )
+                else:
+                    special_audio_mask = (input_ids == self.config.audio_token_id).unsqueeze(-1)
+                    special_audio_mask = special_audio_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+    
+                if not is_torchdynamo_compiling() and inputs_embeds[special_audio_mask].numel() != audio_features.numel():
+                    audio_tokens_in_text = (special_audio_mask).sum(dim=1).sum(dim=0)[0]
+                    raise ValueError(
+                        f"Number of audio input features does not match number of special audio tokens in the input text. "
+                        f"Got {audio_tokens_in_text} audio tokens in the text and "
+                        f"{audio_features.shape[0] * audio_features.shape[1]} tokens from audio embeddings."
+                    )
+                audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
+                inputs_embeds = inputs_embeds.masked_scatter(special_audio_mask, audio_features)
+    
+            outputs = self.language_model(
+                input_ids=None,
+                per_layer_inputs=per_layer_inputs,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=True,
+                cache_position=cache_position,
+                **lm_kwargs,
+            )
+    
+            return Gemma3nModelOutputWithPast(
+                last_hidden_state=outputs.last_hidden_state,
+                past_key_values=outputs.past_key_values if use_cache else None,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+                image_hidden_states=image_features if pixel_values is not None else None,
+                audio_hidden_states=audio_features if input_features is not None else None,
+            )
+    
+        def get_audio_features(
+            self, input_features: torch.Tensor, input_features_mask: torch.Tensor
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            """
+            Projects the last hidden state from the audio encoder into language model space.
+    
+            Args:
+                input_features (`torch.FloatTensor]` of shape `(num_images, seq_length, num_features)`):
+                   The tensors corresponding to the input audio.
+                input_features (`torch.FloatTensor]` of shape `(num_images, seq_length)`):
+                   The attention mask for the input audio.
+    
+            Returns:
+                audio_features (`torch.Tensor`): Audio feature tensor of shape `(num_images, audio_length, embed_dim)`).
+            """
+            audio_outputs, audio_mask = self.audio_tower(input_features, input_features_mask)
+            return self.embed_audio(inputs_embeds=audio_outputs), audio_mask
+    
+    
+    @auto_docstring(
+        custom_intro="""
+        The base Gemma 3n model comprising a vision backbone, an audio backbone, a language model, and a language modeling
+        head.
+        """
+    )
+    class Gemma3nForConditionalGeneration(Gemma3nPreTrainedModel, GenerationMixin):
+        _checkpoint_conversion_mapping = {}
+        _tied_weights_keys = ["lm_head.weight"]
+        base_model_prefix = "model"
+    
+        def __init__(self, config: Gemma3nConfig):
+            super().__init__(config)
+            self.model = Gemma3nModel(config)
+            self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+            self.post_init()
+    
+        def get_input_embeddings(self):
+            return self.model.get_input_embeddings()
+    
+        def set_input_embeddings(self, value):
+            self.model.set_input_embeddings(value)
+    
+        def get_output_embeddings(self):
+            return self.lm_head
+    
+        def set_output_embeddings(self, new_embeddings):
+            self.lm_head = new_embeddings
+    
+        def set_decoder(self, decoder):
+            self.model.set_decoder(decoder)
+    
+        def get_decoder(self):
+            return self.model.get_decoder()
+    
+        def get_image_features(self, pixel_values):
+            return self.model.get_image_features(pixel_values)
+    
+        # Make modules available throught conditional class for BC
+        @property
+        def language_model(self):
+            return self.model.language_model
+    
+        @property
+        def vision_tower(self):
+            return self.model.vision_tower
+    
+        @property
+        def multi_modal_projector(self):
+            raise AttributeError("Use embed_vision instead of multi_modal_projector.")
+    
+        @can_return_tuple
+        @auto_docstring
+        def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,  # text inputs
+            pixel_values: Optional[torch.FloatTensor] = None,  # vision inputs
+            input_features: Optional[torch.FloatTensor] = None,  # audio inputs
+            attention_mask: Optional[torch.Tensor] = None,
+            input_features_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None,
+            token_type_ids: Optional[torch.LongTensor] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            logits_to_keep: Union[int, torch.Tensor] = 0,
+            **lm_kwargs,
+        ) -> Gemma3nCausalLMOutputWithPast:
+            r"""
+            input_features (torch.Tensor, *optional*, defaults to None):
+                The audio inputs to be encoded.
+            input_features_mask (torch.Tensor, *optional*, defaults to None):
+                The attention mask for the input audio.
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are
+                ignored (masked), the loss is only computed for the tokens with labels in
+                `[0, ..., config.text_config.vocab_size]`.
+    
+            Example:
+    
+            ```python
+            >>> from PIL import Image
+            >>> import requests
+            >>> from transformers import AutoProcessor, Gemma3ForConditionalGeneration
+    
+            >>> model = Gemma3ForConditionalGeneration.from_pretrained("google/gemma-3-4b-it")
+            >>> processor = AutoProcessor.from_pretrained("google/gemma-3-4b-it")
+    
+            >>> messages = [
+            ...     {
+            ...         "role": "system",
+            ...         "content": [
+            ...             {"type": "text", "text": "You are a helpful assistant."}
+            ...         ]
+            ...     },
+            ...     {
+            ...         "role": "user", "content": [
+            ...             {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"},
+            ...             {"type": "text", "text": "Where is the cat standing?"},
+            ...         ]
+            ...     },
+            ... ]
+    
+            >>> inputs = processor.apply_chat_template(
+            ...     messages,
+            ...     tokenizer=True,
+            ...     return_dict=True,
+            ...     return_tensors="pt",
+            ...     add_generation_prompt=True
+            ... )
+            >>> # Generate
+            >>> generate_ids = model.generate(**inputs)
+            >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+            "user\nYou are a helpful assistant.\n\n\n\n\n\nWhere is the cat standing?\nmodel\nBased on the image, the cat is standing in a snowy area, likely outdoors. It appears to"
+            ```
+            """
+            output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+            output_hidden_states = (
+                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+            )
+    
+            outputs = self.model(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                input_features=input_features,
+                attention_mask=attention_mask,
+                input_features_mask=input_features_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                token_type_ids=token_type_ids,
+                cache_position=cache_position,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=True,
+                **lm_kwargs,
+            )
+    
+            hidden_states = outputs.last_hidden_state
+            # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+            logits = self.lm_head(hidden_states[:, slice_indices, :])
+            if (final_logit_softcapping := self.config.get_text_config().final_logit_softcapping) is not None:
+                logits = logits / final_logit_softcapping
+                logits = torch.tanh(logits)
+                logits = logits * final_logit_softcapping
+    
+            loss = None
+            if labels is not None:
+                # Upcast to float if we need to compute the loss to avoid potential precision issues
+                logits = logits.float()
+                shift_logits = logits[..., :-1, :]
+                shift_labels = labels[..., 1:]
+                if attention_mask is not None:
+                    # we use the input attention mask to shift the logits and labels, because it is 2D.
+                    # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                    shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(logits.device)
+                    shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous()
+                    shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
+                else:
+                    shift_logits = shift_logits.contiguous()
+                    shift_labels = shift_labels.contiguous()
+                # Flatten the tokens
+                loss_fct = nn.CrossEntropyLoss()
+    
+                flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
+                flat_labels = shift_labels.view(-1).to(shift_logits.device)
+                loss = loss_fct(flat_logits, flat_labels)
+    
+            return Gemma3nCausalLMOutputWithPast(
+                loss=loss,
+                logits=logits,
+                past_key_values=outputs.past_key_values,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+                image_hidden_states=outputs.image_hidden_states,
+                audio_hidden_states=outputs.audio_hidden_states,
+            )
+    
+        def prepare_inputs_for_generation(
+            self,
+            input_ids,
+            past_key_values=None,
+            inputs_embeds=None,
+            cache_position=None,
+            position_ids=None,
+            pixel_values=None,
+            input_features=None,
+            attention_mask=None,
+            input_features_mask=None,
+            token_type_ids=None,
+            use_cache=True,
+            logits_to_keep=None,
+            labels=None,
+            **kwargs,
+        ):
+            # Overwritten -- custom `position_ids` and `pixel_values` handling
+            model_inputs = super().prepare_inputs_for_generation(
+                input_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                cache_position=cache_position,
+                use_cache=use_cache,
+                logits_to_keep=logits_to_keep,
+                token_type_ids=token_type_ids,
+                **kwargs,
+            )
+    
+            # If we're in cached decoding stage, multimodal inputs should be None because input ids do not contain special
+            # tokens anymore. Otherwise multimodal inputs should be passed to model.
+            # NOTE: use_cache=False always needs pixel_values, input_features, and input_features_mask
+            if cache_position[0] == 0:
+                model_inputs["pixel_values"] = pixel_values
+                model_inputs["input_features"] = input_features
+                model_inputs["input_features_mask"] = input_features_mask
+    
+            return model_inputs
+    
+        @property
+        def audio_tower(self):
+            return self.model.audio_tower
+    
+    
+    __all__ = [
+        "Gemma3nAudioEncoder",
+        "Gemma3nForCausalLM",
+        "Gemma3nForConditionalGeneration",
+        "Gemma3nModel",
+        "Gemma3nPreTrainedModel",
+        "Gemma3nTextModel",
+    ]
+    
+    </script>
+<script id="dependencies" type="text/plain">
+    # coding=utf-8
+    # Copyright 2020 The Google AI Language Team Authors, Facebook AI Research authors and The HuggingFace Inc. team.
+    # Copyright (c) 2020, NVIDIA CORPORATION.  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.
+    import copy
+    import inspect
+    import os
+    import warnings
+    from dataclasses import dataclass
+    from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+    
+    import numpy as np
+    import torch
+    import torch.distributed as dist
+    from huggingface_hub import file_exists
+    from packaging import version
+    from torch import nn
+    from torch.nn import functional as F
+    
+    from ..cache_utils import (
+        Cache,
+        DynamicCache,
+        EncoderDecoderCache,
+        HybridChunkedCache,
+        OffloadedCache,
+        OffloadedHybridCache,
+        QuantizedCacheConfig,
+    )
+    from ..configuration_utils import PretrainedConfig
+    from ..dynamic_module_utils import (
+        check_python_requirements,
+        get_cached_module_file,
+        get_class_in_module,
+        resolve_trust_remote_code,
+    )
+    from ..integrations.deepspeed import is_deepspeed_zero3_enabled
+    from ..integrations.fsdp import is_fsdp_managed_module
+    from ..masking_utils import create_masks_for_generate
+    from ..modeling_outputs import CausalLMOutputWithPast, Seq2SeqLMOutput
+    from ..pytorch_utils import isin_mps_friendly
+    from ..tokenization_utils import ExtensionsTrie
+    from ..utils import (
+        ModelOutput,
+        is_accelerate_available,
+        is_hqq_available,
+        is_optimum_quanto_available,
+        is_torchdynamo_exporting,
+        logging,
+    )
+    from .beam_constraints import DisjunctiveConstraint, PhrasalConstraint
+    from .beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
+    from .candidate_generator import (
+        AssistantVocabTranslatorCache,
+        AssistedCandidateGenerator,
+        AssistedCandidateGeneratorDifferentTokenizers,
+        CandidateGenerator,
+        EarlyExitCandidateGenerator,
+        PromptLookupCandidateGenerator,
+        UniversalSpeculativeDecodingGenerator,
+        _crop_past_key_values,
+        _prepare_attention_mask,
+        _prepare_token_type_ids,
+    )
+    from .configuration_utils import (
+        NEED_SETUP_CACHE_CLASSES_MAPPING,
+        QUANT_BACKEND_CLASSES_MAPPING,
+        CompileConfig,
+        GenerationConfig,
+        GenerationMode,
+    )
+    from .continuous_batching import ContinuousMixin
+    from .logits_process import (
+        EncoderNoRepeatNGramLogitsProcessor,
+        EncoderRepetitionPenaltyLogitsProcessor,
+        EpsilonLogitsWarper,
+        EtaLogitsWarper,
+        ExponentialDecayLengthPenalty,
+        ForcedBOSTokenLogitsProcessor,
+        ForcedEOSTokenLogitsProcessor,
+        HammingDiversityLogitsProcessor,
+        InfNanRemoveLogitsProcessor,
+        LogitNormalization,
+        LogitsProcessorList,
+        MinLengthLogitsProcessor,
+        MinNewTokensLengthLogitsProcessor,
+        MinPLogitsWarper,
+        NoBadWordsLogitsProcessor,
+        NoRepeatNGramLogitsProcessor,
+        PrefixConstrainedLogitsProcessor,
+        RepetitionPenaltyLogitsProcessor,
+        SequenceBiasLogitsProcessor,
+        SuppressTokensAtBeginLogitsProcessor,
+        SuppressTokensLogitsProcessor,
+        TemperatureLogitsWarper,
+        TopKLogitsWarper,
+        TopPLogitsWarper,
+        TypicalLogitsWarper,
+        UnbatchedClassifierFreeGuidanceLogitsProcessor,
+    )
+    from .stopping_criteria import (
+        ConfidenceCriteria,
+        EosTokenCriteria,
+        MaxLengthCriteria,
+        MaxTimeCriteria,
+        StoppingCriteria,
+        StoppingCriteriaList,
+        StopStringCriteria,
+    )
+    
+    
+    if TYPE_CHECKING:
+        from ..modeling_utils import PreTrainedModel
+        from ..tokenization_utils_base import PreTrainedTokenizerBase
+        from .streamers import BaseStreamer
+    
+    logger = logging.get_logger(__name__)
+    
+    if is_accelerate_available():
+        from accelerate.hooks import AlignDevicesHook, add_hook_to_module
+    
+    
+    # Variable names used to hold the cache at generation time
+    ALL_CACHE_NAMES = [
+        "past_key_values",  # default
+        "cache_params",  # mamba-based models
+        "state",  # rwkv
+        "mems",  # xlnet
+        "past_buckets_states",  # reformer
+    ]
+    
+    
+    @dataclass
+    class GenerateDecoderOnlyOutput(ModelOutput):
+        """
+        Outputs of decoder-only generation models, when using non-beam methods.
+    
+        Args:
+            sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
+                if all batches finished early due to the `eos_token_id`.
+            scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`):
+                Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+                at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+                each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+            logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`):
+                Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+                at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+                each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+            attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
+            hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`.
+            past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True`):
+                Returns the model cache, used to speed up decoding. Different models have a different cache format, check
+                the model's documentation. Usually, a [`~cache_utils.Cache`] instance.
+        """
+    
+        sequences: torch.LongTensor
+        scores: Optional[tuple[torch.FloatTensor]] = None
+        logits: Optional[tuple[torch.FloatTensor]] = None
+        attentions: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        hidden_states: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        past_key_values: Optional[tuple[tuple[tuple[torch.FloatTensor]]]] = None
+    
+    
+    @dataclass
+    class GenerateEncoderDecoderOutput(ModelOutput):
+        """
+        Outputs of encoder-decoder generation models, when using non-beam methods.
+    
+        Args:
+            sequences (`torch.LongTensor` of shape `(batch_size*num_return_sequences, sequence_length)`):
+                The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
+                if all batches finished early due to the `eos_token_id`.
+            scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`):
+                Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+                at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+                each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+            logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`):
+                Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+                at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+                each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+            encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True`):
+                Tuple of `torch.FloatTensor` (one for each layer of the decoder) of shape `(batch_size, num_heads,
+                sequence_length, sequence_length)`.
+            encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True`):
+                Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
+                shape `(batch_size, sequence_length, hidden_size)`.
+            decoder_attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
+            cross_attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
+            decoder_hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`.
+            past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+                Returns the model cache, used to speed up decoding. Different models have a different cache format, check
+                the model's documentation. Usually, a [`~cache_utils.Cache`] instance.
+        """
+    
+        sequences: torch.LongTensor
+        scores: Optional[tuple[torch.FloatTensor]] = None
+        logits: Optional[tuple[torch.FloatTensor]] = None
+        encoder_attentions: Optional[tuple[torch.FloatTensor]] = None
+        encoder_hidden_states: Optional[tuple[torch.FloatTensor]] = None
+        decoder_attentions: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        cross_attentions: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        decoder_hidden_states: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        past_key_values: Optional[tuple[tuple[tuple[torch.FloatTensor]]]] = None
+    
+    
+    @dataclass
+    class GenerateBeamDecoderOnlyOutput(ModelOutput):
+        """
+        Outputs of decoder-only generation models, when using beam methods.
+    
+        Args:
+            sequences (`torch.LongTensor` of shape `(batch_size*num_return_sequences, sequence_length)`):
+                The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
+                if all batches finished early due to the `eos_token_id`.
+            sequences_scores (`torch.FloatTensor` of shape `(batch_size*num_return_sequences)`, *optional*, returned when `output_scores=True`):
+                Final beam scores of the generated `sequences`.
+            scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`):
+                Beam transition scores for each vocabulary token at each generation step. Beam transition scores consisting
+                of log probabilities of tokens conditioned on log softmax of previously generated tokens in this beam.
+                Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for each generated token),
+                with each tensor of shape `(batch_size*num_beams, config.vocab_size)`.
+            logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`):
+                Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+                at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+                each generated token), with each tensor of shape `(batch_size*num_beams, config.vocab_size)`.
+            beam_indices (`torch.LongTensor`, *optional*, returned when `output_scores=True`):
+                Beam indices of generated token id at each generation step. `torch.LongTensor` of shape
+                `(batch_size*num_return_sequences, sequence_length)`.
+            attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size*num_beams, num_heads, generated_length, sequence_length)`.
+            hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`.
+            past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True`):
+                Returns the model cache, used to speed up decoding. Different models have a different cache format, check
+                the model's documentation. Usually, a [`~cache_utils.Cache`] instance.
+        """
+    
+        sequences: torch.LongTensor
+        sequences_scores: Optional[torch.FloatTensor] = None
+        scores: Optional[tuple[torch.FloatTensor]] = None
+        logits: Optional[tuple[torch.FloatTensor]] = None
+        beam_indices: Optional[torch.LongTensor] = None
+        attentions: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        hidden_states: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        past_key_values: Optional[tuple[tuple[tuple[torch.FloatTensor]]]] = None
+    
+    
+    @dataclass
+    class GenerateBeamEncoderDecoderOutput(ModelOutput):
+        """
+        Outputs of encoder-decoder generation models, when using beam methods.
+    
+        Args:
+            sequences (`torch.LongTensor` of shape `(batch_size*num_return_sequences, sequence_length)`):
+                The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
+                if all batches finished early due to the `eos_token_id`.
+            sequences_scores (`torch.FloatTensor` of shape `(batch_size*num_return_sequences)`, *optional*, returned when `output_scores=True`):
+                Final beam scores of the generated `sequences`.
+            scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`):
+                Beam transition scores for each vocabulary token at each generation step. Beam transition scores consisting
+                of log probabilities of tokens conditioned on log softmax of previously generated tokens in this beam.
+                Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for each generated token),
+                with each tensor of shape `(batch_size*num_beams, config.vocab_size)`.
+            logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`):
+                Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+                at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+                each generated token), with each tensor of shape `(batch_size*num_beams, config.vocab_size)`.
+            beam_indices (`torch.LongTensor`, *optional*, returned when `output_scores=True`):
+                Beam indices of generated token id at each generation step. `torch.LongTensor` of shape
+                `(batch_size*num_return_sequences, sequence_length)`.
+            encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True`):
+                Tuple of `torch.FloatTensor` (one for each layer of the decoder) of shape `(batch_size, num_heads,
+                sequence_length, sequence_length)`.
+            encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True`):
+                Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
+                shape `(batch_size*num_beams*num_return_sequences, sequence_length, hidden_size)`.
+            decoder_attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size*num_beams*num_return_sequences, num_heads, generated_length,
+                sequence_length)`.
+            cross_attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
+            decoder_hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`):
+                Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+                `torch.FloatTensor` of shape `(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`.
+            past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True`):
+                Returns the model cache, used to speed up decoding. Different models have a different cache format, check
+                the model's documentation. Usually, a [`~cache_utils.Cache`] instance.
+        """
+    
+        sequences: torch.LongTensor
+        sequences_scores: Optional[torch.FloatTensor] = None
+        scores: Optional[tuple[torch.FloatTensor]] = None
+        logits: Optional[tuple[torch.FloatTensor]] = None
+        beam_indices: Optional[torch.LongTensor] = None
+        encoder_attentions: Optional[tuple[torch.FloatTensor]] = None
+        encoder_hidden_states: Optional[tuple[torch.FloatTensor]] = None
+        decoder_attentions: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        cross_attentions: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        decoder_hidden_states: Optional[tuple[tuple[torch.FloatTensor]]] = None
+        past_key_values: Optional[tuple[tuple[tuple[torch.FloatTensor]]]] = None
+    
+    
+    # TODO (joao): remove the equivalent classes and typing shortcuts below in v5
+    # Equivalent classes (kept for retrocompatibility purposes)
+    GreedySearchDecoderOnlyOutput = GenerateDecoderOnlyOutput
+    ContrastiveSearchDecoderOnlyOutput = GenerateDecoderOnlyOutput
+    SampleDecoderOnlyOutput = GenerateDecoderOnlyOutput
+    
+    ContrastiveSearchEncoderDecoderOutput = GenerateEncoderDecoderOutput
+    GreedySearchEncoderDecoderOutput = GenerateEncoderDecoderOutput
+    SampleEncoderDecoderOutput = GenerateEncoderDecoderOutput
+    
+    BeamSearchDecoderOnlyOutput = GenerateBeamDecoderOnlyOutput
+    BeamSampleDecoderOnlyOutput = GenerateBeamDecoderOnlyOutput
+    
+    BeamSearchEncoderDecoderOutput = GenerateBeamEncoderDecoderOutput
+    BeamSampleEncoderDecoderOutput = GenerateBeamEncoderDecoderOutput
+    
+    GreedySearchOutput = Union[GreedySearchEncoderDecoderOutput, GreedySearchDecoderOnlyOutput]
+    SampleOutput = Union[SampleEncoderDecoderOutput, SampleDecoderOnlyOutput]
+    BeamSearchOutput = Union[BeamSearchEncoderDecoderOutput, BeamSearchDecoderOnlyOutput]
+    BeamSampleOutput = Union[BeamSampleEncoderDecoderOutput, BeamSampleDecoderOnlyOutput]
+    ContrastiveSearchOutput = Union[ContrastiveSearchEncoderDecoderOutput, ContrastiveSearchDecoderOnlyOutput]
+    
+    # Typing shortcuts
+    GenerateNonBeamOutput = Union[GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput]
+    GenerateBeamOutput = Union[GenerateBeamDecoderOnlyOutput, GenerateBeamEncoderDecoderOutput]
+    GenerateOutput = Union[GenerateNonBeamOutput, GenerateBeamOutput]
+    
+    
+    class GenerationMixin(ContinuousMixin):
+        """
+        A class containing all functions for auto-regressive text generation, to be used as a mixin in model classes.
+        Inheriting from this class causes the model to have special generation-related behavior, such as loading a
+        `GenerationConfig` at initialization time or ensuring `generate`-related tests are run in `transformers` CI.
+    
+        A model class should inherit from `GenerationMixin` to enable calling methods like `generate`, or when it
+        has defined a custom `generate` method that relies on `GenerationMixin`, directly or indirectly, which
+        approximately shares the same interface to public methods like `generate`. Three examples:
+            - `LlamaForCausalLM` should inherit from `GenerationMixin` to enable calling `generate` and other public
+                methods in the mixin;
+            - `BlipForQuestionAnswering` has a custom `generate` method that approximately shares the same interface as
+                `GenerationMixin.generate` (it has a few extra arguments, and the same output). That function also calls
+                `GenerationMixin.generate` indirectly, through an inner model. As such, `BlipForQuestionAnswering` should
+                inherit from `GenerationMixin` to benefit from all generation-related automation in our codebase;
+            - `BarkModel` has a custom `generate` method and one of its inner models calls `GenerationMixin.generate`.
+                However, its `generate` does not share the same interface as `GenerationMixin.generate`. In this case,
+                `BarkModel` should NOT inherit from `GenerationMixin`, as it breaks the `generate` interface.
+    
+        The class exposes [`~generation.GenerationMixin.generate`], which can be used for:
+            - *greedy decoding* if `num_beams=1` and `do_sample=False`
+            - *contrastive search* if `penalty_alpha>0` and `top_k>1`
+            - *multinomial sampling* if `num_beams=1` and `do_sample=True`
+            - *beam-search decoding* if `num_beams>1` and `do_sample=False`
+            - *beam-search multinomial sampling* if `num_beams>1` and `do_sample=True`
+            - *diverse beam-search decoding* if `num_beams>1` and `num_beam_groups>1`
+            - *constrained beam-search decoding* if `constraints!=None` or `force_words_ids!=None`
+            - *assisted decoding* if `assistant_model` or `prompt_lookup_num_tokens` is passed to `.generate()`
+    
+        To learn more about decoding strategies refer to the [text generation strategies guide](../generation_strategies).
+        """
+    
+        def load_custom_generate(
+            self,
+            pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
+            trust_remote_code: Optional[bool] = None,
+            **kwargs,
+        ) -> Callable:
+            """
+            Loads and returns a custom generate function, given a model repo.
+    
+            Args:
+                pretrained_model_name_or_path (`str` or `os.PathLike`):
+                        Can be either:
+                        - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
+                        - A path to a *directory* containing model weights saved using
+                            [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+                trust_remote_code (`bool`, *optional*):
+                    Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+                    should only be set to `True` for repositories you trust and in which you have read the code, as it will
+                    execute code present on the Hub on your local machine.
+                **kwargs:
+                    Additional keyword arguments for remote code loading.
+    
+            Raises:
+                OSError: If `pretrained_model_name_or_path` does not contain a `custom_generate` subdirectory.
+    
+            Returns:
+                A callable that can be used to generate text.
+            """
+            # Does `pretrained_model_name_or_path` have a `custom_generate` subdirectory? If not -> OSError
+            is_local_code = os.path.exists(pretrained_model_name_or_path)
+            has_custom_generate_folder = True
+            if is_local_code:
+                if not os.path.exists(os.path.join(pretrained_model_name_or_path, "custom_generate/generate.py")):
+                    has_custom_generate_folder = False
+            else:
+                if not file_exists(pretrained_model_name_or_path, "custom_generate/generate.py"):
+                    has_custom_generate_folder = False
+    
+            if not has_custom_generate_folder:
+                raise OSError(
+                    f"`{pretrained_model_name_or_path}` does not contain a `custom_generate` subdirectory with a "
+                    "`generate.py` file, can't load the custom generate function."
+                )
+    
+            # Handle opt-in `trust_remote_code` and related exceptions
+            error_message = (
+                f"The repository `{pretrained_model_name_or_path}` contains custom generation code that will override "
+                "the default `generate` method."
+            )
+            resolve_trust_remote_code(
+                trust_remote_code,
+                pretrained_model_name_or_path,
+                has_local_code=is_local_code,
+                has_remote_code=not is_local_code,
+                error_message=error_message,
+            )
+    
+            # Load the custom generate function
+            check_python_requirements(
+                pretrained_model_name_or_path, requirements_file="custom_generate/requirements.txt", **kwargs
+            )
+            module = get_cached_module_file(
+                pretrained_model_name_or_path, module_file="custom_generate/generate.py", **kwargs
+            )
+            custom_generate_function = get_class_in_module("generate", module)
+            return custom_generate_function
+    
+        def _cache_dependant_input_preparation(
+            self,
+            input_ids: torch.LongTensor,
+            inputs_embeds: Optional[torch.FloatTensor],
+            cache_position: Optional[torch.LongTensor],
+        ) -> tuple[torch.FloatTensor, torch.LongTensor]:
+            """
+            Generic cache-dependent input preparation
+            The code is put in a separate function to allow granular unit testing
+            as it needs a different implementation to be exportable.
+    
+            If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
+            - Exception 1: when passing input_embeds, input_ids may be missing entries
+            - Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
+            - Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.
+            - Exception 4: If input_embeds are passed then slice it through `cache_position`, to keep only the unprocessed tokens and
+                generate the first token for each sequence. Later use the generated Input ids for continuation.
+    
+            The current implementation does not rely on ``self`` and could be
+            a class method. It is left as a standard method to be easily rewritten.
+            """
+            if is_torchdynamo_exporting():
+                return self._cache_dependant_input_preparation_exporting(input_ids, inputs_embeds, cache_position)
+            if inputs_embeds is not None and input_ids.shape[1] == 0:  # Exception 4
+                inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :]
+            elif (
+                inputs_embeds is not None  # Exception 1
+                or (cache_position[-1] >= input_ids.shape[1])  # Exception 3
+            ):
+                input_ids = input_ids[:, -cache_position.shape[0] :]
+            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
+                input_ids = input_ids[:, cache_position]
+            return inputs_embeds, input_ids
+    
+        def _cache_dependant_input_preparation_exporting(
+            self,
+            input_ids: torch.LongTensor,
+            inputs_embeds: Optional[torch.FloatTensor],
+            cache_position: Optional[torch.LongTensor],
+        ) -> tuple[torch.FloatTensor, torch.LongTensor]:
+            """
+            This method implements method ``_cache_dependant_input_preparation``
+            with :func:`torch.cond` to make it exportable with :func:`torch.export.export`.
+            The code is put in a separate function to allow granular unit testing.
+            """
+            if inputs_embeds is None:
+                input_ids = input_ids[:, cache_position]
+            else:
+                # This is the code we need to implemented with torch.cond.
+                # if input_ids.shape[1] == 0:
+                #     inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :]
+                # else:
+                #     if cache_position[-1] >= input_ids.shape[1]:
+                #         input_ids = input_ids[:, -cache_position.shape[0] :]
+                #     else:
+                #         if input_ids.shape[1] != cache_position.shape[0]:
+                #             input_ids = input_ids[:, cache_position]
+                def branch_1(inputs_embeds, cache_position):
+                    return inputs_embeds[:, -cache_position.shape[0] :]
+    
+                def branch_2(input_ids, cache_position):
+                    return input_ids[:, -cache_position.shape[0] :]
+    
+                def branch_3(input_ids, cache_position):
+                    return input_ids[:, cache_position]
+    
+                inputs_embeds, input_ids = torch.cond(
+                    input_ids.shape[1] == 0,
+                    (
+                        lambda input_ids, inputs_embeds, cache_position: (
+                            branch_1(inputs_embeds, cache_position),
+                            input_ids,
+                        )
+                    ),
+                    (
+                        lambda input_ids, inputs_embeds, cache_position: (
+                            inputs_embeds,
+                            torch.cond(
+                                cache_position[-1] >= input_ids.shape[1],
+                                branch_2,
+                                lambda input_ids, cache_position: (
+                                    torch.cond(
+                                        input_ids.shape[1] != cache_position.shape[0],
+                                        branch_3,
+                                        (lambda input_ids, cache_position: input_ids),
+                                        [input_ids, cache_position],
+                                    )
+                                ),
+                                [input_ids, cache_position],
+                            ),
+                        )
+                    ),
+                    [input_ids, inputs_embeds, cache_position],
+                )
+            return inputs_embeds, input_ids
+    
+        def prepare_inputs_for_generation(
+            self,
+            input_ids: torch.LongTensor,
+            past_key_values: Optional[Cache] = None,
+            attention_mask: Optional[torch.LongTensor] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            **kwargs,
+        ):
+            """
+            Prepare the model inputs for generation. It includes operations like computing the 4D attention mask or
+            slicing inputs given the existing cache.
+    
+            See the forward pass in the model documentation for expected arguments (different models might have different
+            requirements for e.g. `past_key_values`). This function should work as is for most LLMs.
+            """
+    
+            # 1. Handle BC:
+            model_inputs = {}
+            # - some models don't have `Cache` support (which implies they don't expect `cache_position` in `forward`)
+            if self._supports_cache_class:
+                model_inputs["cache_position"] = cache_position
+            # - `cache_position` was not a mandatory input in `prepare_inputs_for_generation` for those models, and this
+            #   function may be called outside of `generate`. Handle most use cases by creating `cache_position` on the fly
+            #   (this alternative is not as robust as calling `generate` and letting it create `cache_position`)
+            elif cache_position is None:
+                past_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+                cache_position = torch.arange(past_length, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
+    
+            # 2. Generic cache-dependent input preparation
+            if past_key_values is not None:
+                model_inputs["past_key_values"] = past_key_values
+                inputs_embeds, input_ids = self._cache_dependant_input_preparation(
+                    input_ids, inputs_embeds, cache_position
+                )
+    
+            # 3. Prepare base model inputs
+            input_ids_key = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
+            # if `inputs_embeds` are passed, we only want to use them in the 1st generation step for every prompt.
+            if not self.config.is_encoder_decoder:
+                if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]:
+                    model_inputs[input_ids_key] = None
+                    model_inputs["inputs_embeds"] = inputs_embeds
+                else:
+                    # `clone` calls in this function ensure a consistent stride. See #32227
+                    model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format)
+                    model_inputs["inputs_embeds"] = None
+            else:
+                model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format)
+    
+            # 4. Create missing `position_ids` on the fly
+            encoder_attention_mask = attention_mask if self.config.is_encoder_decoder else None
+            attention_mask = (
+                kwargs.pop("decoder_attention_mask", None) if self.config.is_encoder_decoder else attention_mask
+            )
+            attention_mask_key = "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask"
+            position_ids_key = "decoder_position_ids" if self.config.is_encoder_decoder else "position_ids"
+            if (
+                attention_mask is not None
+                and kwargs.get(position_ids_key) is None
+                and position_ids_key in set(inspect.signature(self.forward).parameters.keys())
+            ):
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids.masked_fill_(attention_mask == 0, 1)
+                kwargs[position_ids_key] = position_ids  # placed in kwargs for further processing (see below)
+    
+            # 5. Slice model inputs if it's an input that should have the same length as `input_ids`
+            for model_input_name in ["position_ids", "token_type_ids", "decoder_position_ids"]:
+                model_input = kwargs.get(model_input_name)
+                if model_input is not None:
+                    if past_key_values is not None:
+                        current_input_length = (
+                            model_inputs["inputs_embeds"].shape[1]
+                            if model_inputs.get("inputs_embeds") is not None
+                            else model_inputs[input_ids_key].shape[1]
+                        )
+                        model_input = model_input[:, -current_input_length:]
+                        model_input = model_input.clone(memory_format=torch.contiguous_format)
+                    model_inputs[model_input_name] = model_input
+    
+            # 6. Create 4D attention mask is we are using a compilable cache (important for performant compiled forward
+            # pass)
+            if (
+                isinstance(past_key_values, Cache)
+                and past_key_values.is_compileable
+                and attention_mask is not None
+                and attention_mask.ndim == 2
+            ):
+                if not self.config.is_encoder_decoder and model_inputs["inputs_embeds"] is not None:
+                    batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
+                else:
+                    batch_size, sequence_length = model_inputs[input_ids_key].shape[:2]
+    
+                # Create the causal mask with fixed shape in advance, to reduce recompilations. If the function to create
+                # the 4D causal mask exists, it should be present in the base model (XXXModel class) or in its decoder.
+                base_model = getattr(self, self.base_model_prefix, self)
+                decoder = base_model.get_decoder() if hasattr(base_model, "get_decoder") else None
+                causal_mask_creation_function = getattr(
+                    base_model, "_prepare_4d_causal_attention_mask_with_cache_position", None
+                )
+                if causal_mask_creation_function is None and decoder is not None:  # it may be in the decoder
+                    causal_mask_creation_function = getattr(
+                        decoder, "_prepare_4d_causal_attention_mask_with_cache_position", None
+                    )
+    
+                # If it's not defined, it means the model uses the new general mask API
+                if causal_mask_creation_function is None:  # can't be found
+                    token_type_ids = getattr(model_input, "token_type_ids", None)
+                    # Some models may overwrite the general one
+                    causal_mask_creation_function = getattr(self, "create_masks_for_generate", create_masks_for_generate)
+                    attention_mask = causal_mask_creation_function(
+                        config=self.config,
+                        # we only need batch size, seq_length and dtype here - we don't care about the values of the embeddings
+                        input_embeds=torch.empty((batch_size, sequence_length), dtype=self.dtype),
+                        attention_mask=attention_mask,
+                        cache_position=cache_position,
+                        past_key_values=past_key_values,
+                        token_type_ids=token_type_ids,
+                    )
+                else:
+                    attention_mask = causal_mask_creation_function(
+                        attention_mask,
+                        sequence_length=sequence_length,
+                        target_length=past_key_values.get_max_cache_shape(),
+                        dtype=self.dtype,
+                        cache_position=cache_position,
+                        batch_size=batch_size,
+                        config=self.config,
+                        past_key_values=past_key_values,
+                    )
+            if attention_mask is not None:
+                model_inputs[attention_mask_key] = attention_mask
+    
+            if encoder_attention_mask is not None:
+                model_inputs["attention_mask"] = encoder_attention_mask
+    
+            # 7. Forward ALL kwargs that are uninitialized (e.g. `use_cache`).
+            for key, value in kwargs.items():
+                if key not in model_inputs:
+                    model_inputs[key] = value
+    
+            # 8. Remove unexpected `generate` inputs (TODO @joao: fix trainer and examples)
+            model_inputs.pop("labels", None)
+            return model_inputs
+    
+        def _prepare_model_inputs(
+            self,
+            inputs: Optional[torch.Tensor] = None,
+            bos_token_id: Optional[torch.Tensor] = None,
+            model_kwargs: Optional[dict[str, torch.Tensor]] = None,
+        ) -> tuple[torch.Tensor, Optional[str], dict[str, torch.Tensor]]:
+            """
+            This function extracts the model-specific `inputs` for generation.
+            """
+            # 1. retrieve all kwargs that are non-None or non-model input related.
+            # some encoder-decoder models have different names for model and encoder
+            if (
+                self.config.is_encoder_decoder
+                and hasattr(self, "encoder")
+                and self.encoder.main_input_name != self.main_input_name
+            ):
+                input_name = self.encoder.main_input_name
+            else:
+                input_name = self.main_input_name
+    
+            model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None or k != input_name}
+    
+            # 2. check whether model_input_name is passed as kwarg
+            # if yes and `inputs` is None use kwarg inputs
+            inputs_kwarg = model_kwargs.pop(input_name, None)
+            if inputs_kwarg is not None and inputs is not None:
+                raise ValueError(
+                    f"`inputs`: {inputs}` were passed alongside {input_name} which is not allowed. "
+                    f"Make sure to either pass {inputs} or {input_name}=..."
+                )
+            elif inputs_kwarg is not None:
+                inputs = inputs_kwarg
+    
+            # 3. In the presence of `inputs_embeds` for text models:
+            # - decoder-only models should complain if the user attempts to pass `inputs_embeds`, but the model
+            # doesn't have its forwarding implemented. `inputs_embeds` is kept in `model_kwargs` and can coexist with
+            # input_ids (`inputs_embeds` will be used in the 1st generation step, as opposed to `input_ids`)
+            # - encoder-decoder models should complain if the user attempts to pass `inputs_embeds` and `input_ids`, and
+            # pull the former to inputs. It will be used in place of `input_ids` to get the encoder hidden states.
+            if input_name == "input_ids" and "inputs_embeds" in model_kwargs:
+                if model_kwargs["inputs_embeds"] is None:
+                    model_kwargs.pop("inputs_embeds")
+                elif not self.config.is_encoder_decoder:
+                    has_inputs_embeds_forwarding = "inputs_embeds" in set(
+                        inspect.signature(self.prepare_inputs_for_generation).parameters.keys()
+                    )
+                    if not has_inputs_embeds_forwarding:
+                        raise ValueError(
+                            f"You passed `inputs_embeds` to `.generate()`, but the model class {self.__class__.__name__} "
+                            "doesn't have its forwarding implemented. See the GPT2 implementation for an example "
+                            "(https://github.com/huggingface/transformers/pull/21405), and feel free to open a PR with it!"
+                        )
+                    # In this case, `input_ids` is moved to the `model_kwargs`, so a few automations (like the creation of
+                    # the attention mask) can rely on the actual model input.
+                    model_kwargs["input_ids"] = self._maybe_initialize_input_ids_for_generation(
+                        inputs, bos_token_id, model_kwargs=model_kwargs
+                    )
+                    inputs, input_name = model_kwargs["inputs_embeds"], "inputs_embeds"
+                else:
+                    if inputs is not None:
+                        raise ValueError("You passed `inputs_embeds` and `input_ids` to `.generate()`. Please pick one.")
+                    inputs, input_name = model_kwargs["inputs_embeds"], "inputs_embeds"
+    
+            # 4. if `inputs` is still None, try to create `input_ids` from BOS token
+            inputs = self._maybe_initialize_input_ids_for_generation(inputs, bos_token_id, model_kwargs)
+            return inputs, input_name, model_kwargs
+    
+        def _maybe_initialize_input_ids_for_generation(
+            self,
+            inputs: Optional[torch.Tensor] = None,
+            bos_token_id: Optional[torch.Tensor] = None,
+            model_kwargs: Optional[dict[str, torch.Tensor]] = None,
+        ) -> torch.LongTensor:
+            """Initializes input ids for generation, if necessary."""
+            if inputs is not None:
+                return inputs
+    
+            encoder_outputs = model_kwargs.get("encoder_outputs")
+            if self.config.is_encoder_decoder and encoder_outputs is not None:
+                # make dummy input_ids with value -100, as a sanity check ensuring that they won't be used for encoding
+                shape = encoder_outputs.last_hidden_state.size()[:-1]
+                return torch.ones(shape, dtype=torch.long, device=self.device) * -100
+    
+            # If there is some tensor in `model_kwargs`, we can infer the batch size from it. This is helpful with
+            # soft-prompting or in multimodal implementations built on top of decoder-only language models.
+            batch_size = 1
+            for value in model_kwargs.values():
+                if isinstance(value, torch.Tensor):
+                    batch_size = value.shape[0]
+                    break
+    
+            if "inputs_embeds" in model_kwargs:
+                return torch.ones((batch_size, 0), dtype=torch.long, device=self.device)
+    
+            if bos_token_id is None:
+                raise ValueError("`bos_token_id` has to be defined when no `input_ids` are provided.")
+    
+            return torch.ones((batch_size, 1), dtype=torch.long, device=self.device) * bos_token_id
+    
+        def _prepare_attention_mask_for_generation(
+            self,
+            inputs_tensor: torch.Tensor,
+            generation_config: GenerationConfig,
+            model_kwargs: dict[str, Any],
+        ) -> torch.LongTensor:
+            pad_token_id = generation_config._pad_token_tensor
+            eos_token_id = generation_config._eos_token_tensor
+    
+            # `input_ids` may be present in the model kwargs, instead of being the main input (e.g. multimodal model)
+            if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0:
+                inputs_tensor = model_kwargs["input_ids"]
+    
+            # No information for attention mask inference -> return default attention mask
+            default_attention_mask = torch.ones(inputs_tensor.shape[:2], dtype=torch.long, device=inputs_tensor.device)
+            if pad_token_id is None:
+                return default_attention_mask
+    
+            is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long]
+            if not is_input_ids:
+                return default_attention_mask
+    
+            is_pad_token_in_inputs = (pad_token_id is not None) and (
+                isin_mps_friendly(elements=inputs_tensor, test_elements=pad_token_id).any()
+            )
+            is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or ~(
+                isin_mps_friendly(elements=eos_token_id, test_elements=pad_token_id).any()
+            )
+            can_infer_attention_mask = is_pad_token_in_inputs * is_pad_token_not_equal_to_eos_token_id
+            attention_mask_from_padding = inputs_tensor.ne(pad_token_id).long()
+    
+            attention_mask = (
+                attention_mask_from_padding * can_infer_attention_mask + default_attention_mask * ~can_infer_attention_mask
+            )
+            return attention_mask
+    
+        def _prepare_encoder_decoder_kwargs_for_generation(
+            self,
+            inputs_tensor: torch.Tensor,
+            model_kwargs,
+            model_input_name: Optional[str],
+            generation_config: GenerationConfig,
+        ) -> dict[str, Any]:
+            # 1. get encoder
+            encoder = self.get_encoder()
+            # Compatibility with Accelerate big model inference: we need the encoder to outputs stuff on the same device
+            # as the inputs.
+            if hasattr(self, "hf_device_map"):
+                if hasattr(encoder, "_hf_hook"):
+                    encoder._hf_hook.io_same_device = True
+                else:
+                    add_hook_to_module(encoder, AlignDevicesHook(io_same_device=True))
+    
+            # 2. Prepare encoder args and encoder kwargs from model kwargs and generation config.
+            irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]
+            encoder_kwargs = {
+                argument: value
+                for argument, value in model_kwargs.items()
+                if not any(argument.startswith(p) for p in irrelevant_prefix)
+            }
+            encoder_signature = set(inspect.signature(encoder.forward).parameters)
+            encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature
+            if not encoder_accepts_wildcard:
+                encoder_kwargs = {
+                    argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature
+                }
+            encoder_kwargs["output_attentions"] = generation_config.output_attentions
+            encoder_kwargs["output_hidden_states"] = generation_config.output_hidden_states
+    
+            # 3. make sure that encoder returns `ModelOutput`
+            model_input_name = model_input_name if model_input_name is not None else self.main_input_name
+            encoder_kwargs["return_dict"] = True
+            encoder_kwargs[model_input_name] = inputs_tensor
+            model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)  # type: ignore
+    
+            return model_kwargs
+    
+        def _prepare_decoder_input_ids_for_generation(
+            self,
+            batch_size: int,
+            model_input_name: str,
+            model_kwargs: dict[str, torch.Tensor],
+            decoder_start_token_id: torch.Tensor,
+            device: Optional[torch.device] = None,
+        ) -> tuple[torch.LongTensor, dict[str, torch.Tensor]]:
+            """Prepares `decoder_input_ids` for generation with encoder-decoder models"""
+            # 1. Check whether the user has defined `decoder_input_ids` manually. To facilitate in terms of input naming,
+            # we also allow the user to pass it under `input_ids`, if the encoder does not use it as the main input.
+            if model_kwargs is not None and "decoder_input_ids" in model_kwargs:
+                decoder_input_ids = model_kwargs.pop("decoder_input_ids")
+            elif "input_ids" in model_kwargs and model_input_name != "input_ids":
+                decoder_input_ids = model_kwargs.pop("input_ids")
+            else:
+                decoder_input_ids = None
+    
+            # 2. `decoder_start_token_id` must have shape (batch_size, 1)
+            if device is None:
+                device = self.device
+            if decoder_start_token_id.ndim == 1:
+                if decoder_start_token_id.shape[0] != batch_size:
+                    raise ValueError(
+                        f"`decoder_start_token_id` expected to have length {batch_size} but got {decoder_start_token_id.shape[0]}"
+                    )
+                decoder_start_token_id = decoder_start_token_id.view(-1, 1)
+            else:
+                decoder_start_token_id = (
+                    torch.ones((batch_size, 1), dtype=torch.long, device=device) * decoder_start_token_id
+                )
+    
+            # 3. Encoder-decoder models expect the `decoder_input_ids` to start with a special token. Let's ensure that.
+            # no user input -> use decoder_start_token_id as decoder_input_ids
+            if decoder_input_ids is None:
+                decoder_input_ids = decoder_start_token_id
+            # exception: Donut checkpoints have task-specific decoder starts and don't expect a BOS token. Note that the
+            # original checkpoints can't be detected through `self.__class__.__name__.lower()`, needing custom logic.
+            # See: https://github.com/huggingface/transformers/pull/31470
+            elif "donut" in self.__class__.__name__.lower() or (
+                self.config.model_type == "vision-encoder-decoder" and "donut" in self.config.encoder.model_type.lower()
+            ):
+                pass
+            elif self.config.model_type in ["whisper"]:
+                pass
+            # user input but doesn't start with decoder_start_token_id -> prepend decoder_start_token_id (and adjust
+            # decoder_attention_mask if provided)
+            elif (decoder_input_ids[:, 0] != decoder_start_token_id[:, 0]).all().item():
+                decoder_input_ids = torch.cat([decoder_start_token_id, decoder_input_ids], dim=-1)
+                if "decoder_attention_mask" in model_kwargs:
+                    decoder_attention_mask = model_kwargs["decoder_attention_mask"]
+                    decoder_attention_mask = torch.cat(
+                        (torch.ones_like(decoder_attention_mask)[:, :1], decoder_attention_mask),
+                        dim=-1,
+                    )
+                    model_kwargs["decoder_attention_mask"] = decoder_attention_mask
+    
+            return decoder_input_ids, model_kwargs
+    
+        @staticmethod
+        def _expand_inputs_for_generation(
+            expand_size: int = 1,
+            is_encoder_decoder: bool = False,
+            input_ids: Optional[torch.LongTensor] = None,
+            **model_kwargs,
+        ) -> tuple[torch.LongTensor, dict[str, Any]]:
+            """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]"""
+            # Do not call torch.repeat_interleave if expand_size is 1 because it clones
+            # the input tensor and thus requires more memory although no change is applied
+            if expand_size == 1:
+                return input_ids, model_kwargs
+    
+            def _expand_dict_for_generation(dict_to_expand):
+                for key in dict_to_expand:
+                    if (
+                        key != "cache_position"
+                        and dict_to_expand[key] is not None
+                        and isinstance(dict_to_expand[key], torch.Tensor)
+                    ):
+                        dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+                return dict_to_expand
+    
+            if input_ids is not None:
+                input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+    
+            model_kwargs = _expand_dict_for_generation(model_kwargs)
+    
+            if is_encoder_decoder:
+                if model_kwargs.get("encoder_outputs") is None:
+                    raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+                model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+    
+            return input_ids, model_kwargs
+    
+        def _update_model_kwargs_for_generation(
+            self,
+            outputs: ModelOutput,
+            model_kwargs: dict[str, Any],
+            is_encoder_decoder: bool = False,
+            num_new_tokens: int = 1,
+        ) -> dict[str, Any]:
+            # update past_key_values keeping its naming used in model code
+            for possible_cache_name in ALL_CACHE_NAMES:
+                if possible_cache_name in outputs:
+                    # TODO (joao): remove output/input mismatch when these old models (xlnet, reformer) are deprecated
+                    if possible_cache_name in ("past_buckets_states", "mems"):
+                        cache_name = "past_key_values"
+                    else:
+                        cache_name = possible_cache_name
+                    model_kwargs[cache_name] = getattr(outputs, possible_cache_name)
+                    break
+    
+            # update token_type_ids with last value
+            if "token_type_ids" in model_kwargs:
+                token_type_ids = model_kwargs["token_type_ids"]
+                model_kwargs["token_type_ids"] = torch.cat([token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1)
+    
+            if not is_encoder_decoder:
+                # update attention mask
+                if "attention_mask" in model_kwargs:
+                    attention_mask = model_kwargs["attention_mask"]
+                    model_kwargs["attention_mask"] = torch.cat(
+                        [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
+                    )
+            else:
+                # update decoder attention mask
+                if "decoder_attention_mask" in model_kwargs:
+                    decoder_attention_mask = model_kwargs["decoder_attention_mask"]
+                    model_kwargs["decoder_attention_mask"] = torch.cat(
+                        [decoder_attention_mask, decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], 1))],
+                        dim=-1,
+                    )
+    
+            if model_kwargs.get("use_cache", True):
+                model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
+            else:
+                past_positions = model_kwargs.pop("cache_position")
+                new_positions = torch.arange(
+                    past_positions[-1] + 1, past_positions[-1] + num_new_tokens + 1, dtype=past_positions.dtype
+                ).to(past_positions.device)
+                model_kwargs["cache_position"] = torch.cat((past_positions, new_positions))
+            return model_kwargs
+    
+        def _reorder_cache(self, past_key_values, beam_idx):
+            raise NotImplementedError(
+                f"Make sure that a `_reorder_cache` function is correctly implemented in {self.__class__.__module__} to"
+                f" enable beam search for {self.__class__}"
+            )
+    
+        def _get_candidate_generator(
+            self,
+            generation_config: GenerationConfig,
+            input_ids: torch.LongTensor,
+            inputs_tensor: torch.Tensor,
+            assistant_model: "PreTrainedModel",
+            logits_processor: LogitsProcessorList,
+            target_tokenizer: "PreTrainedTokenizerBase",
+            assistant_tokenizer: "PreTrainedTokenizerBase",
+            model_kwargs: dict,
+        ) -> CandidateGenerator:
+            """
+            Returns the candidate generator to be used in `assisted_generation`
+            """
+            different_tokenizers = all(v is not None for v in (assistant_model, target_tokenizer, assistant_tokenizer))
+    
+            if generation_config.assistant_early_exit is not None:
+                candidate_generator = EarlyExitCandidateGenerator(
+                    input_ids=input_ids,
+                    assistant_model=self,
+                    generation_config=generation_config,
+                    model_kwargs=model_kwargs,
+                    inputs_tensor=inputs_tensor,
+                    logits_processor=logits_processor,
+                )
+            elif generation_config.prompt_lookup_num_tokens is not None:
+                candidate_generator = PromptLookupCandidateGenerator(
+                    eos_token_id=generation_config._eos_token_tensor,
+                    num_output_tokens=generation_config.prompt_lookup_num_tokens,
+                    max_matching_ngram_size=generation_config.max_matching_ngram_size,
+                    max_length=generation_config.max_length,
+                )
+            elif different_tokenizers:
+                if generation_config.do_sample is True:
+                    atm_translator = AssistantVocabTranslatorCache.get_translator(
+                        target_tokenizer,
+                        assistant_tokenizer,
+                        self.config.get_text_config().vocab_size,
+                        assistant_model=assistant_model,
+                        assistant_prune_lm_head=True,  # prune LM head of assistant model
+                    )
+                    # Since we prune the LM head, we cannot use the repetition penalty on the assistant model due to mismatches between token ids and logits index
+                    assistant_model.generation_config.repetition_penalty = None
+                    candidate_generator = UniversalSpeculativeDecodingGenerator(
+                        input_ids=input_ids,
+                        assistant_model=assistant_model,
+                        generation_config=generation_config,
+                        model_kwargs=model_kwargs,
+                        inputs_tensor=inputs_tensor,
+                        logits_processor=logits_processor,
+                        target_tokenizer=target_tokenizer,
+                        assistant_tokenizer=assistant_tokenizer,
+                        atm_translator=atm_translator,
+                    )
+                elif generation_config.do_sample is False:
+                    candidate_generator = AssistedCandidateGeneratorDifferentTokenizers(
+                        input_ids=input_ids,
+                        assistant_model=assistant_model,
+                        generation_config=generation_config,
+                        model_kwargs=model_kwargs,
+                        inputs_tensor=inputs_tensor,
+                        logits_processor=logits_processor,
+                        target_tokenizer=target_tokenizer,
+                        assistant_tokenizer=assistant_tokenizer,
+                    )
+                else:
+                    raise ValueError(
+                        f"Invalid value for `do_sample`: expected a boolean, got {type(generation_config.do_sample).__name__}"
+                    )
+            else:
+                candidate_generator = AssistedCandidateGenerator(
+                    input_ids=input_ids,
+                    assistant_model=assistant_model,
+                    generation_config=generation_config,
+                    model_kwargs=model_kwargs,
+                    inputs_tensor=inputs_tensor,
+                    logits_processor=logits_processor,
+                )
+            return candidate_generator
+    
+        def _get_logits_processor(
+            self,
+            generation_config: GenerationConfig,
+            input_ids_seq_length: Optional[int] = None,
+            encoder_input_ids: torch.LongTensor = None,
+            prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], list[int]]] = None,
+            logits_processor: Optional[LogitsProcessorList] = None,
+            device: Optional[str] = None,
+            model_kwargs: Optional[dict[str, Any]] = None,
+            negative_prompt_ids: Optional[torch.Tensor] = None,
+            negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        ) -> LogitsProcessorList:
+            """
+            This class returns a [`LogitsProcessorList`] list object that contains all relevant [`LogitsProcessor`]
+            instances used to modify the scores of the language model head.
+            """
+            # instantiate processors list
+            processors = LogitsProcessorList()
+            if logits_processor is None:
+                logits_processor = []
+    
+            if generation_config.guidance_scale is not None and generation_config.guidance_scale != 1:
+                processors.append(
+                    UnbatchedClassifierFreeGuidanceLogitsProcessor(
+                        generation_config.guidance_scale,
+                        self,
+                        unconditional_ids=negative_prompt_ids,
+                        unconditional_attention_mask=negative_prompt_attention_mask,
+                        use_cache=generation_config.use_cache,
+                    )
+                )
+            if generation_config.sequence_bias is not None:
+                processors.append(SequenceBiasLogitsProcessor(sequence_bias=generation_config.sequence_bias))
+    
+            if generation_config.diversity_penalty is not None and generation_config.diversity_penalty > 0.0:
+                processors.append(
+                    HammingDiversityLogitsProcessor(
+                        diversity_penalty=generation_config.diversity_penalty,
+                        num_beams=generation_config.num_beams,
+                        num_beam_groups=generation_config.num_beam_groups,
+                    )
+                )
+            if (
+                generation_config.encoder_repetition_penalty is not None
+                and generation_config.encoder_repetition_penalty != 1.0
+            ):
+                if len(encoder_input_ids.shape) == 2:
+                    processors.append(
+                        EncoderRepetitionPenaltyLogitsProcessor(
+                            penalty=generation_config.encoder_repetition_penalty,
+                            encoder_input_ids=encoder_input_ids,
+                        )
+                    )
+                else:
+                    warnings.warn(
+                        "Passing `encoder_repetition_penalty` requires some form of `input_ids` to be passed to "
+                        "`generate`, ignoring the argument.",
+                        UserWarning,
+                    )
+            if generation_config.repetition_penalty is not None and generation_config.repetition_penalty != 1.0:
+                processors.append(RepetitionPenaltyLogitsProcessor(penalty=generation_config.repetition_penalty))
+            if generation_config.no_repeat_ngram_size is not None and generation_config.no_repeat_ngram_size > 0:
+                processors.append(NoRepeatNGramLogitsProcessor(generation_config.no_repeat_ngram_size))
+            if (
+                generation_config.encoder_no_repeat_ngram_size is not None
+                and generation_config.encoder_no_repeat_ngram_size > 0
+            ):
+                if len(encoder_input_ids.shape) == 2:
+                    processors.append(
+                        EncoderNoRepeatNGramLogitsProcessor(
+                            generation_config.encoder_no_repeat_ngram_size,
+                            encoder_input_ids,
+                        )
+                    )
+                else:
+                    warnings.warn(
+                        "Passing `encoder_no_repeat_ngram_size` requires some form of `input_ids` to be passed to "
+                        "`generate`, ignoring the argument.",
+                        UserWarning,
+                    )
+            if generation_config.bad_words_ids is not None:
+                processors.append(
+                    NoBadWordsLogitsProcessor(
+                        generation_config.bad_words_ids,
+                        generation_config._eos_token_tensor,
+                    )
+                )
+            if (
+                generation_config.min_length is not None
+                and getattr(generation_config, "_eos_token_tensor", None) is not None
+                and generation_config.min_length > 0
+            ):
+                processors.append(
+                    MinLengthLogitsProcessor(
+                        generation_config.min_length,
+                        generation_config._eos_token_tensor,
+                        device=device,
+                    )
+                )
+            if (
+                generation_config.min_new_tokens is not None
+                and getattr(generation_config, "_eos_token_tensor", None) is not None
+                and generation_config.min_new_tokens > 0
+            ):
+                processors.append(
+                    MinNewTokensLengthLogitsProcessor(
+                        input_ids_seq_length,
+                        generation_config.min_new_tokens,
+                        generation_config._eos_token_tensor,
+                        device=device,
+                    )
+                )
+            if prefix_allowed_tokens_fn is not None:
+                processors.append(
+                    PrefixConstrainedLogitsProcessor(
+                        prefix_allowed_tokens_fn,
+                        generation_config.num_beams // generation_config.num_beam_groups,
+                    )
+                )
+            if generation_config.forced_bos_token_id is not None:
+                processors.append(
+                    ForcedBOSTokenLogitsProcessor(
+                        generation_config.forced_bos_token_id,
+                    )
+                )
+            if generation_config.forced_eos_token_id is not None:
+                processors.append(
+                    ForcedEOSTokenLogitsProcessor(
+                        generation_config.max_length,
+                        generation_config.forced_eos_token_id,
+                        device=device,
+                    )
+                )
+            if generation_config.remove_invalid_values is True:
+                processors.append(InfNanRemoveLogitsProcessor())
+            if generation_config.exponential_decay_length_penalty is not None:
+                processors.append(
+                    ExponentialDecayLengthPenalty(
+                        generation_config.exponential_decay_length_penalty,
+                        generation_config._eos_token_tensor,
+                        input_ids_seq_length,
+                    )
+                )
+            if generation_config.suppress_tokens is not None:
+                processors.append(
+                    SuppressTokensLogitsProcessor(
+                        generation_config.suppress_tokens,
+                        device=device,
+                    )
+                )
+            if generation_config.begin_suppress_tokens is not None:
+                begin_index = input_ids_seq_length
+                begin_index = (
+                    begin_index
+                    if (input_ids_seq_length > 1 or generation_config.forced_bos_token_id is None)
+                    else begin_index + 1
+                )
+                processors.append(
+                    SuppressTokensAtBeginLogitsProcessor(
+                        generation_config.begin_suppress_tokens,
+                        begin_index,
+                        device=device,
+                    )
+                )
+    
+            # TODO (joao): find a strategy to specify the order of the processors
+            processors = self._merge_criteria_processor_list(processors, logits_processor)
+    
+            # Processors previously known as `LogitsWarpers`, only applied with sampling strategies
+            if generation_config.do_sample:
+                # In beam methods, we need to keep at least one non-eos token to explore continuations that might have a
+                # better score (i.e. keep len(list(generation_config._eos_token_tensor)) + 1)
+                if generation_config.num_beams > 1:
+                    if isinstance(generation_config._eos_token_tensor, list):
+                        min_tokens_to_keep = len(generation_config._eos_token_tensor) + 1
+                    elif isinstance(generation_config._eos_token_tensor, torch.Tensor):
+                        min_tokens_to_keep = generation_config._eos_token_tensor.shape[0] + 1
+                    else:
+                        min_tokens_to_keep = 2
+                else:
+                    min_tokens_to_keep = 1
+    
+                # the following idea is largely copied from this PR: https://github.com/huggingface/transformers/pull/5420/files
+                # all samplers can be found in `generation_utils_samplers.py`
+                if generation_config.temperature is not None and generation_config.temperature != 1.0:
+                    processors.append(TemperatureLogitsWarper(generation_config.temperature))
+                if generation_config.top_k is not None and generation_config.top_k != 0:
+                    processors.append(
+                        TopKLogitsWarper(top_k=generation_config.top_k, min_tokens_to_keep=min_tokens_to_keep)
+                    )
+                if generation_config.top_p is not None and generation_config.top_p < 1.0:
+                    processors.append(
+                        TopPLogitsWarper(top_p=generation_config.top_p, min_tokens_to_keep=min_tokens_to_keep)
+                    )
+                if generation_config.min_p is not None:
+                    # Applied after temperature scaling (see https://github.com/ggerganov/llama.cpp/pull/3841#issuecomment-2073826084)
+                    processors.append(
+                        MinPLogitsWarper(min_p=generation_config.min_p, min_tokens_to_keep=min_tokens_to_keep)
+                    )
+                if generation_config.typical_p is not None and generation_config.typical_p < 1.0:
+                    processors.append(
+                        TypicalLogitsWarper(mass=generation_config.typical_p, min_tokens_to_keep=min_tokens_to_keep)
+                    )
+                if generation_config.epsilon_cutoff is not None and 0.0 < generation_config.epsilon_cutoff < 1.0:
+                    processors.append(
+                        EpsilonLogitsWarper(
+                            epsilon=generation_config.epsilon_cutoff, min_tokens_to_keep=min_tokens_to_keep
+                        )
+                    )
+                if generation_config.eta_cutoff is not None and 0.0 < generation_config.eta_cutoff < 1.0:
+                    processors.append(
+                        EtaLogitsWarper(
+                            epsilon=generation_config.eta_cutoff, min_tokens_to_keep=min_tokens_to_keep, device=device
+                        )
+                    )
+    
+            # Watermarking should be after all logits processing is finished (see #34630)
+            if generation_config.watermarking_config is not None:
+                processors.append(
+                    generation_config.watermarking_config.construct_processor(
+                        self.config.get_text_config().vocab_size, device
+                    )
+                )
+    
+            # `LogitNormalization` should always be the last logit processor, when present
+            if generation_config.renormalize_logits is True:
+                processors.append(LogitNormalization())
+            return processors
+    
+        def _get_stopping_criteria(
+            self,
+            generation_config: GenerationConfig,
+            stopping_criteria: Optional[StoppingCriteriaList],
+            tokenizer: Optional["PreTrainedTokenizerBase"] = None,
+            **kwargs,
+        ) -> StoppingCriteriaList:
+            criteria = StoppingCriteriaList()
+            if generation_config.max_length is not None:
+                max_position_embeddings = getattr(self.config, "max_position_embeddings", None)
+                criteria.append(
+                    MaxLengthCriteria(
+                        max_length=generation_config.max_length,
+                        max_position_embeddings=max_position_embeddings,
+                    )
+                )
+            if generation_config.max_time is not None:
+                criteria.append(MaxTimeCriteria(max_time=generation_config.max_time))
+            if generation_config.stop_strings is not None:
+                if tokenizer is None:
+                    raise ValueError(
+                        "There are one or more stop strings, either in the arguments to `generate` or in the "
+                        "model's generation config, but we could not locate a tokenizer. When generating with "
+                        "stop strings, you must pass the model's tokenizer to the `tokenizer` argument of `generate`."
+                    )
+                criteria.append(StopStringCriteria(stop_strings=generation_config.stop_strings, tokenizer=tokenizer))
+            if generation_config._eos_token_tensor is not None:
+                criteria.append(EosTokenCriteria(eos_token_id=generation_config._eos_token_tensor))
+            if (
+                generation_config.is_assistant
+                and generation_config.assistant_confidence_threshold is not None
+                and generation_config.assistant_confidence_threshold > 0
+            ):
+                criteria.append(
+                    ConfidenceCriteria(assistant_confidence_threshold=generation_config.assistant_confidence_threshold)
+                )
+            criteria = self._merge_criteria_processor_list(criteria, stopping_criteria)
+            return criteria
+    
+        def _merge_criteria_processor_list(
+            self,
+            default_list: Union[LogitsProcessorList, StoppingCriteriaList],
+            custom_list: Union[LogitsProcessorList, StoppingCriteriaList],
+        ) -> Union[LogitsProcessorList, StoppingCriteriaList]:
+            """
+            Merge user-defined processors/criteria with the ones instantiated inside `generate`. In case the same
+            processor/criteria is present on both lists, use the user-defined one.
+    
+            (Note: up to v4.49.0, this function threw an exception is the same logit processor was found twice.)
+            """
+            if len(custom_list) == 0:
+                return default_list
+    
+            final_list = type(default_list)()
+            for default in default_list:
+                using_custom = False
+                for custom in custom_list:
+                    if type(custom) is type(default):
+                        object_type = "stopping criteria" if isinstance(custom, StoppingCriteria) else "logits processor"
+                        logger.warning_once(
+                            f"A custom {object_type} of type {type(custom)} has been passed to `.generate()`, but it "
+                            f"was also created in `.generate()`, given its parameterization. The custom {type(custom)} "
+                            f"will take precedence. Please check the docstring of {type(custom)} to see related "
+                            "`.generate()` flags."
+                        )
+                        final_list.append(custom)
+                        using_custom = True
+                        break
+                if not using_custom:
+                    final_list.append(default)
+    
+            for custom in custom_list:
+                if custom not in final_list:
+                    final_list.append(custom)
+            return final_list
+    
+        def compute_transition_scores(
+            self,
+            sequences: torch.Tensor,
+            scores: tuple[torch.Tensor],
+            beam_indices: Optional[torch.Tensor] = None,
+            normalize_logits: bool = False,
+        ) -> torch.Tensor:
+            """
+            Computes the transition scores of sequences given the generation scores (and beam indices, if beam search was
+            used). This is a convenient method to quickly obtain the scores of the selected tokens at generation time.
+    
+            Parameters:
+                sequences (`torch.LongTensor`):
+                    The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or
+                    shorter if all batches finished early due to the `eos_token_id`.
+                scores (`tuple(torch.FloatTensor)`):
+                    Transition scores for each vocabulary token at each generation step. Beam transition scores consisting
+                    of log probabilities of tokens conditioned on log softmax of previously generated tokens in this beam.
+                    Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for each generated token),
+                    with each tensor of shape `(batch_size*num_beams, config.vocab_size)`.
+                beam_indices (`torch.LongTensor`, *optional*):
+                    Beam indices of generated token id at each generation step. `torch.LongTensor` of shape
+                    `(batch_size*num_return_sequences, sequence_length)`. Only required if a `num_beams>1` at
+                    generate-time.
+                normalize_logits (`bool`, *optional*, defaults to `False`):
+                    Whether to normalize the logits (which, for legacy reasons, may be unnormalized).
+    
+            Return:
+                `torch.Tensor`: A `torch.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)` containing
+                    the transition scores (logits)
+    
+            Examples:
+    
+            ```python
+            >>> from transformers import GPT2Tokenizer, AutoModelForCausalLM
+            >>> import numpy as np
+    
+            >>> tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+            >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
+            >>> tokenizer.pad_token_id = tokenizer.eos_token_id
+            >>> inputs = tokenizer(["Today is"], return_tensors="pt")
+    
+            >>> # Example 1: Print the scores for each token generated with Greedy Search
+            >>> outputs = model.generate(**inputs, max_new_tokens=5, return_dict_in_generate=True, output_scores=True)
+            >>> transition_scores = model.compute_transition_scores(
+            ...     outputs.sequences, outputs.scores, normalize_logits=True
+            ... )
+            >>> # input_length is the length of the input prompt for decoder-only models, like the GPT family, and 1 for
+            >>> # encoder-decoder models, like BART or T5.
+            >>> input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
+            >>> generated_tokens = outputs.sequences[:, input_length:]
+            >>> for tok, score in zip(generated_tokens[0], transition_scores[0]):
+            ...     # | token | token string | log probability | probability
+            ...     print(f"| {tok:5d} | {tokenizer.decode(tok):8s} | {score.numpy():.3f} | {np.exp(score.numpy()):.2%}")
+            |   262 |  the     | -1.414 | 24.33%
+            |  1110 |  day     | -2.609 | 7.36%
+            |   618 |  when    | -2.010 | 13.40%
+            |   356 |  we      | -1.859 | 15.58%
+            |   460 |  can     | -2.508 | 8.14%
+    
+            >>> # Example 2: Reconstruct the sequence scores from Beam Search
+            >>> outputs = model.generate(
+            ...     **inputs,
+            ...     max_new_tokens=5,
+            ...     num_beams=4,
+            ...     num_return_sequences=4,
+            ...     return_dict_in_generate=True,
+            ...     output_scores=True,
+            ... )
+            >>> transition_scores = model.compute_transition_scores(
+            ...     outputs.sequences, outputs.scores, outputs.beam_indices, normalize_logits=False
+            ... )
+            >>> # If you sum the generated tokens' scores and apply the length penalty, you'll get the sequence scores.
+            >>> # Tip 1: recomputing the scores is only guaranteed to match with `normalize_logits=False`. Depending on the
+            >>> # use case, you might want to recompute it with `normalize_logits=True`.
+            >>> # Tip 2: the output length does NOT include the input length
+            >>> output_length = np.sum(transition_scores.numpy() < 0, axis=1)
+            >>> length_penalty = model.generation_config.length_penalty
+            >>> reconstructed_scores = transition_scores.sum(axis=1) / (output_length**length_penalty)
+            >>> print(np.allclose(outputs.sequences_scores, reconstructed_scores))
+            True
+            ```"""
+            # 1. In absence of `beam_indices`, we can assume that we come from e.g. greedy search, which is equivalent
+            # to a beam search approach were the first (and only) beam is always selected
+            if beam_indices is None:
+                beam_indices = torch.arange(scores[0].shape[0]).view(-1, 1).to(sequences.device)
+                beam_indices = beam_indices.expand(-1, len(scores))
+    
+            # 2. reshape scores as [batch_size*vocab_size, # generation steps] with # generation steps being
+            # seq_len - input_length
+            scores = torch.stack(scores).reshape(len(scores), -1).transpose(0, 1)
+    
+            # 3. Optionally normalize the logits (across the vocab dimension)
+            if normalize_logits:
+                scores = scores.reshape(-1, self.config.get_text_config().vocab_size, scores.shape[-1])
+                scores = torch.nn.functional.log_softmax(scores, dim=1)
+                scores = scores.reshape(-1, scores.shape[-1])
+    
+            # 4. cut beam_indices to longest beam length
+            beam_indices_mask = beam_indices < 0
+            max_beam_length = (1 - beam_indices_mask.long()).sum(-1).max()
+            beam_indices = beam_indices.clone()[:, :max_beam_length]
+            beam_indices_mask = beam_indices_mask[:, :max_beam_length]
+    
+            # 5. Set indices of beams that finished early to 0; such indices will be masked correctly afterwards
+            beam_indices[beam_indices_mask] = 0
+    
+            # 6. multiply beam_indices with vocab size to gather correctly from scores
+            beam_sequence_indices = beam_indices * self.config.get_text_config().vocab_size
+    
+            # 7. Define which indices contributed to scores
+            cut_idx = sequences.shape[-1] - max_beam_length
+            indices = sequences[:, cut_idx:] + beam_sequence_indices
+    
+            # 8. Compute scores
+            transition_scores = scores.gather(0, indices)
+    
+            # 9. Mask out transition_scores of beams that stopped early
+            transition_scores[beam_indices_mask] = 0
+    
+            return transition_scores
+    
+        def _validate_assistant(self, assistant_model, tokenizer, assistant_tokenizer):
+            if assistant_model is None:
+                return
+    
+            if self.config.is_encoder_decoder and not assistant_model.config.is_encoder_decoder:
+                attributes_to_check = ["encoder_attention_heads", "encoder_ffn_dim", "encoder_layers"]
+                attributes_to_check = [attr for attr in dir(assistant_model.config) if attr in attributes_to_check]
+                are_equal = all(
+                    getattr(self.config, attr) == getattr(assistant_model.config, attr) for attr in attributes_to_check
+                )
+                if not are_equal:
+                    raise ValueError(
+                        "The main model and the assistant don't have compatible encoder-dependent input shapes. "
+                        "Ensure you load the assistant with the correct encoder-decoder class, e.g. `AutoModelForSpeechSeq2Seq` for Whisper."
+                    )
+    
+            doc_reference = (
+                "(see https://huggingface.co/docs/transformers/en/generation_strategies#universal-assisted-decoding)"
+            )
+            if self.config.get_text_config().vocab_size == assistant_model.config.get_text_config().vocab_size:
+                if assistant_tokenizer is not None:
+                    raise ValueError(
+                        f"`assistant_tokenizer` is not required when the main and assistant models use the same tokenizer. Please omit `assistant_tokenizer` from `generate()` {doc_reference}."
+                    )
+            else:
+                if tokenizer is None or assistant_tokenizer is None:
+                    raise ValueError(
+                        f"The main and assistant moedels have different tokenizers. Please provide `tokenizer` and `assistant_tokenizer` to `generate()` {doc_reference}."
+                    )
+    
+        def _validate_model_kwargs(self, model_kwargs: dict[str, Any]):
+            """Validates model kwargs for generation. Generate argument typos will also be caught here."""
+            # If a `Cache` instance is passed, checks whether the model is compatible with it
+            if isinstance(model_kwargs.get("past_key_values", None), Cache) and not self._supports_cache_class:
+                raise ValueError(
+                    f"{self.__class__.__name__} does not support an instance of `Cache` as `past_key_values`. Please "
+                    "check the model documentation for supported cache formats."
+                )
+    
+            # Excludes arguments that are handled before calling any model function
+            if self.config.is_encoder_decoder:
+                for key in ["decoder_input_ids"]:
+                    model_kwargs.pop(key, None)
+    
+            unused_model_args = []
+            model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters)
+            # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If
+            # `prepare_inputs_for_generation` doesn't accept them, then a stricter check can be made ;)
+            if "kwargs" in model_args or "model_kwargs" in model_args:
+                model_args |= set(inspect.signature(self.forward).parameters)
+    
+            # Encoder-Decoder models may also need Encoder arguments from `model_kwargs`
+            if self.config.is_encoder_decoder:
+                base_model = getattr(self, self.base_model_prefix, None)
+    
+                # allow encoder kwargs
+                encoder = getattr(self, "encoder", None)
+                # `MusicgenForConditionalGeneration` has `text_encoder` and `audio_encoder`.
+                # Also, it has `base_model_prefix = "encoder_decoder"` but there is no `self.encoder_decoder`
+                # TODO: A better way to handle this.
+                if encoder is None and base_model is not None:
+                    encoder = getattr(base_model, "encoder", None)
+    
+                if encoder is not None:
+                    encoder_model_args = set(inspect.signature(encoder.forward).parameters)
+                    model_args |= encoder_model_args
+    
+                # allow decoder kwargs
+                decoder = getattr(self, "decoder", None)
+                if decoder is None and base_model is not None:
+                    decoder = getattr(base_model, "decoder", None)
+    
+                if decoder is not None:
+                    decoder_model_args = set(inspect.signature(decoder.forward).parameters)
+                    model_args |= {f"decoder_{x}" for x in decoder_model_args}
+    
+            for key, value in model_kwargs.items():
+                if value is not None and key not in model_args:
+                    unused_model_args.append(key)
+    
+            if unused_model_args:
+                raise ValueError(
+                    f"The following `model_kwargs` are not used by the model: {unused_model_args} (note: typos in the"
+                    " generate arguments will also show up in this list)"
+                )
+    
+        def _validate_generated_length(self, generation_config, input_ids_length, has_default_max_length):
+            """Performs validation related to the resulting generated length"""
+            # 1. Max length warnings related to poor parameterization
+            if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20:
+                # 20 is the default max_length of the generation config
+                warnings.warn(
+                    f"Using the model-agnostic default `max_length` (={generation_config.max_length}) to control the "
+                    "generation length. We recommend setting `max_new_tokens` to control the maximum length of the "
+                    "generation.",
+                    UserWarning,
+                )
+            if input_ids_length >= generation_config.max_length:
+                input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
+                raise ValueError(
+                    f"Input length of {input_ids_string} is {input_ids_length}, but `max_length` is set to"
+                    f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
+                    " increasing `max_length` or, better yet, setting `max_new_tokens`."
+                )
+    
+            # 2. Min length warnings due to unfeasible parameter combinations
+            min_length_error_suffix = (
+                " Generation will stop at the defined maximum length. You should decrease the minimum length and/or "
+                "increase the maximum length."
+            )
+            if has_default_max_length:
+                min_length_error_suffix += (
+                    f" Note that `max_length` is set to {generation_config.max_length}, its default value."
+                )
+            if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length:
+                warnings.warn(
+                    f"Unfeasible length constraints: `min_length` ({generation_config.min_length}) is larger than"
+                    f" the maximum possible length ({generation_config.max_length})." + min_length_error_suffix,
+                    UserWarning,
+                )
+            if generation_config.min_new_tokens is not None:
+                min_length = generation_config.min_new_tokens + input_ids_length
+                if min_length > generation_config.max_length:
+                    warnings.warn(
+                        f"Unfeasible length constraints: `min_new_tokens` ({generation_config.min_new_tokens}), when "
+                        f"added to the prompt length ({input_ids_length}), is larger than"
+                        f" the maximum possible length ({generation_config.max_length})." + min_length_error_suffix,
+                        UserWarning,
+                    )
+    
+        def _prepare_generated_length(
+            self,
+            generation_config,
+            has_default_max_length,
+            has_default_min_length,
+            model_input_name,
+            input_ids_length,
+            inputs_tensor,
+        ):
+            """Prepared max and min length in generation configs to avoid clashes between similar attributes"""
+    
+            if generation_config.max_new_tokens is not None:
+                if not has_default_max_length and generation_config.max_length is not None:
+                    logger.warning(
+                        f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
+                        f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
+                        "Please refer to the documentation for more information. "
+                        "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
+                    )
+                generation_config.max_length = generation_config.max_new_tokens + input_ids_length
+    
+            # if both `inputs_embeds` and `input_ids` are passed, we do not correct the length
+            # otherwise we need total length [inputs-embeds-len + new-tokens-len] to not go beyond indicated `max_length``
+            elif (
+                model_input_name == "inputs_embeds"
+                and input_ids_length != inputs_tensor.shape[1]
+                and not self.config.is_encoder_decoder
+            ):
+                generation_config.max_length -= inputs_tensor.shape[1]
+            elif has_default_max_length:  # by default let's always generate 20 new tokens
+                if generation_config.max_length == GenerationConfig().max_length:
+                    generation_config.max_length = generation_config.max_length + input_ids_length
+                    max_position_embeddings = getattr(self.config, "max_position_embeddings", None)
+                    if max_position_embeddings is not None:
+                        generation_config.max_length = min(generation_config.max_length, max_position_embeddings)
+    
+            # same for min length
+            if generation_config.min_new_tokens is not None:
+                if not has_default_min_length:
+                    logger.warning(
+                        f"Both `min_new_tokens` (={generation_config.min_new_tokens}) and `min_length`(="
+                        f"{generation_config.min_length}) seem to have been set. `min_new_tokens` will take precedence. "
+                        "Please refer to the documentation for more information. "
+                        "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
+                    )
+                generation_config.min_length = generation_config.min_new_tokens + input_ids_length
+    
+            elif (
+                model_input_name == "inputs_embeds"
+                and input_ids_length != inputs_tensor.shape[1]
+                and not self.config.is_encoder_decoder
+            ):
+                generation_config.min_length = max(generation_config.min_length - inputs_tensor.shape[1], 0)
+    
+            return generation_config
+    
+        def _prepare_generation_config(
+            self, generation_config: Optional[GenerationConfig], use_model_defaults: Optional[bool] = None, **kwargs: dict
+        ) -> tuple[GenerationConfig, dict]:
+            """
+            Prepares the base generation config, then applies any generation configuration options from kwargs. This
+            function handles retrocompatibility with respect to configuration files.
+            """
+            # parameterization priority:
+            # kwargs > non-global default values in `generation_config` > `model.generation_config` > GenerationConfig()
+            # TODO (joao): per-model generation config classes.
+    
+            using_model_generation_config = False
+            if generation_config is None:
+                # legacy: users may modify the model configuration to control generation. To trigger this legacy behavior,
+                # the following conditions must be met
+                # 1) the generation config must have been created from the model config (`_from_model_config` field);
+                # 2) the generation config must have seen no modification since its creation (the hash is the same);
+                # 3) there are non-default generation parameters in the model config.
+                # 4) the user must have set new generation parameters in the model config.
+                if (
+                    self.generation_config._from_model_config  # 1)
+                    and self.generation_config._original_object_hash == hash(self.generation_config)  # 2)
+                    and len(self.config._get_non_default_generation_parameters()) > 0  # 3)
+                ):
+                    new_generation_config = GenerationConfig.from_model_config(self.config)
+                    if new_generation_config != self.generation_config:  # 4)
+                        warnings.warn(
+                            "You have modified the pretrained model configuration to control generation. This is a"
+                            " deprecated strategy to control generation and will be removed in v5."
+                            " Please use and modify the model generation configuration (see"
+                            " https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )",
+                            UserWarning,
+                        )
+                        self.generation_config = new_generation_config
+    
+                generation_config = self.generation_config
+                using_model_generation_config = True
+    
+            # `torch.export.export` usually raises an exception if it is called
+            # with ``strict=True``. deepcopy can only be processed if ``strict=False``.
+            generation_config = copy.deepcopy(generation_config)
+    
+            if not using_model_generation_config:
+                # If `generation_config` is provided:
+                # - `use_model_defaults`: let's fallback ALL default values to the model's generation config
+                # - otherwise: legacy behavior, let's just make sure we have the tokens defined
+                model_base_version = version.parse(version.parse(self.generation_config.transformers_version).base_version)
+                if use_model_defaults is True or (
+                    use_model_defaults is None and model_base_version >= version.parse("4.50.0")
+                ):
+                    modified_values = {}
+                    global_default_generation_config = GenerationConfig()
+                    model_generation_config = self.generation_config
+                    # we iterate over the model's generation config: it may hold custom keys, which we'll want to copy
+                    for key, model_gen_config_value in model_generation_config.__dict__.items():
+                        if key.startswith("_") or key == "transformers_version":  # metadata
+                            continue
+                        global_default_value = getattr(global_default_generation_config, key, None)
+                        custom_gen_config_value = getattr(generation_config, key, None)
+                        if (
+                            custom_gen_config_value == global_default_value
+                            and model_gen_config_value != global_default_value
+                        ):
+                            modified_values[key] = model_gen_config_value
+                            setattr(generation_config, key, model_gen_config_value)
+                    if use_model_defaults is None and len(modified_values) > 0:
+                        logger.warning_once(
+                            f"`generation_config` default values have been modified to match model-specific defaults: "
+                            f"{modified_values}. If this is not desired, please set these values explicitly."
+                        )
+                else:
+                    if generation_config.bos_token_id is None:
+                        generation_config.bos_token_id = self.generation_config.bos_token_id
+                    if generation_config.eos_token_id is None:
+                        generation_config.eos_token_id = self.generation_config.eos_token_id
+                    if generation_config.pad_token_id is None:
+                        generation_config.pad_token_id = self.generation_config.pad_token_id
+                    if generation_config.decoder_start_token_id is None:
+                        generation_config.decoder_start_token_id = self.generation_config.decoder_start_token_id
+    
+            # Finally, apply any passed kwargs
+            model_kwargs = generation_config.update(**kwargs)
+    
+            return generation_config, model_kwargs
+    
+        def _get_initial_cache_position(self, seq_length, device, model_kwargs):
+            """Calculates `cache_position` for the pre-fill stage based on `input_ids` and optionally past length"""
+            # `torch.compile`-friendly `torch.arange` from a shape -- the lines below are equivalent to `torch.arange`
+            if "cache_position" in model_kwargs and model_kwargs["cache_position"]:
+                return model_kwargs
+            if "inputs_embeds" in model_kwargs and not self.config.is_encoder_decoder:
+                cache_position = torch.ones_like(model_kwargs["inputs_embeds"][0, :, 0], dtype=torch.int64).cumsum(0) - 1
+            elif "decoder_inputs_embeds" in model_kwargs and self.config.is_encoder_decoder:
+                cache_position = (
+                    torch.ones_like(model_kwargs["decoder_inputs_embeds"][0, :, 0], dtype=torch.int64).cumsum(0) - 1
+                )
+            else:
+                cache_position = torch.ones(seq_length, dtype=torch.int64, device=device).cumsum(0) - 1
+    
+            past_length = 0
+            if model_kwargs.get("past_key_values") is not None:
+                cache = model_kwargs["past_key_values"]
+                past_length = 0
+                if not isinstance(cache, Cache):
+                    past_length = cache[0][0].shape[2]
+                elif hasattr(cache, "get_seq_length") and cache.get_seq_length() is not None:
+                    past_length = cache.get_seq_length()
+    
+                cache_position = cache_position[past_length:]
+    
+            model_kwargs["cache_position"] = cache_position
+            return model_kwargs
+    
+        def _get_layer_device_map_for_cache_init(self) -> Optional[dict[int, Union[str, int]]]:
+            """
+            Returns the device map for each decoder layer, to allocate the cache on the right device.
+            Inspired from `dispatch_model` in accelerate.
+            """
+            execution_device_map = None
+    
+            if hasattr(self, "hf_device_map"):
+                if set(self.hf_device_map.values()) == {"cpu"} or set(self.hf_device_map.values()) == {"cpu", "disk"}:
+                    main_device = "cpu"
+                else:
+                    main_device = [d for d in self.hf_device_map.values() if d not in ["cpu", "disk"]][0]
+                execution_device_map = {
+                    name: main_device if device in ["cpu", "disk"] else device
+                    for name, device in self.hf_device_map.items()
+                }
+    
+            # No `execution_device_map` -> rely on `self.device` to allocate the cache
+            if execution_device_map is None:
+                return None
+    
+            # Single device for all layers
+            num_hidden_layers = self.config.get_text_config().num_hidden_layers
+            if len(execution_device_map) == 1 and "" in execution_device_map:
+                return dict.fromkeys(range(num_hidden_layers), execution_device_map[""])
+    
+            # Multiple devices in `execution_device_map` -> we need to map decoder layers to the correct device.
+            layer_device_map = {}
+            # Case 1: The model has a `get_decoder` method, we can use it to find the decoder name.
+            if hasattr(self, "get_decoder"):
+                decoder_name = None
+                for name, module in self.named_modules():
+                    if module is self.get_decoder():
+                        decoder_name = name
+                        break
+                if decoder_name is None:
+                    raise RuntimeError(
+                        "`model.get_decoder()` is not returning a named module of the model. This is unexpected, please "
+                        "open an issue on GitHub."
+                    )
+    
+                decoder_mapped_modules = [
+                    module_name for module_name in execution_device_map.keys() if decoder_name in module_name
+                ]
+                # The decoder name may be present in `execution_device_map` in two forms:
+                # a) each layer has a device mapping
+                if len(decoder_mapped_modules) >= num_hidden_layers:
+                    for idx in range(num_hidden_layers):
+                        for module_name in decoder_mapped_modules:
+                            if f".{idx}." in f"{module_name}.":
+                                layer_device_map[idx] = execution_device_map[module_name]
+                                break
+    
+                # b) the whole module is mapped to a single device. If the decoder name is NOT present in the device map,
+                # then the mapping is done in a parent module
+                else:
+                    while True:
+                        if decoder_name in execution_device_map:
+                            layer_device_map = dict.fromkeys(range(num_hidden_layers), execution_device_map[decoder_name])
+                            break
+                        elif "." in decoder_name:
+                            decoder_name = decoder_name.rsplit(".", 1)[0]  # gets the name of the parent module
+                        else:
+                            raise RuntimeError(f"Decoder name {decoder_name} not found in execution device map")
+    
+            # Case 2: Legacy code path: assume the decoder layers are named as `(...).X` (X being the layer index)
+            else:
+                for layer in execution_device_map:
+                    for idx in range(num_hidden_layers):
+                        if f".{idx}." in f"{layer}.":
+                            layer_device_map[idx] = execution_device_map[layer]
+                            break
+    
+            for idx in range(num_hidden_layers):
+                if idx not in layer_device_map:
+                    raise RuntimeError(f"layer {idx} has not been mapped to a device.")
+            return layer_device_map
+    
+        def _get_cache(
+            self, cache_implementation: str, batch_size: int, max_cache_len: int, device: torch.device, model_kwargs
+        ) -> Cache:
+            """
+            Sets a cache for `generate`, that will persist across calls. A new cache will only be initialized a
+            new `generate` call requires a larger cache or uses a different batch size.
+    
+            Returns the resulting cache object.
+            """
+            if cache_implementation == "hybrid" and "llama4" in getattr(self.config, "model_type", ""):
+                cache_implementation = "hybrid_chunked"
+    
+            cache_cls: Cache = NEED_SETUP_CACHE_CLASSES_MAPPING[cache_implementation]
+            requires_cross_attention_cache = (
+                self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None
+            )
+    
+            if hasattr(self, "_cache"):
+                cache_to_check = self._cache.self_attention_cache if requires_cross_attention_cache else self._cache
+    
+            if cache_implementation == "sliding_window":
+                max_cache_len = min(self.config.sliding_window, max_cache_len)
+    
+            need_new_cache = (
+                not hasattr(self, "_cache")
+                or (not isinstance(cache_to_check, cache_cls))
+                or cache_to_check.max_batch_size != batch_size
+                or isinstance(
+                    cache_to_check, (HybridChunkedCache, OffloadedHybridCache)
+                )  # due to internal slicing, we always re-init
+            )
+            if cache_implementation != "mamba":
+                need_new_cache = need_new_cache or cache_to_check.max_cache_len < max_cache_len
+    
+            if requires_cross_attention_cache and hasattr(self, "_cache"):
+                need_new_cache = (
+                    need_new_cache
+                    or self._cache.cross_attention_cache.max_cache_len != model_kwargs["encoder_outputs"][0].shape[1]
+                )
+    
+            if need_new_cache:
+                if hasattr(self.config, "_pre_quantization_dtype"):
+                    cache_dtype = self.config._pre_quantization_dtype
+                else:
+                    cache_dtype = self.dtype
+    
+                layer_device_map = self._get_layer_device_map_for_cache_init()
+                cache_kwargs = {
+                    "config": self.config.get_text_config(),
+                    "max_batch_size": batch_size,
+                    "max_cache_len": max_cache_len,
+                    "dtype": cache_dtype,
+                    "device": device,
+                    "layer_device_map": layer_device_map,
+                }
+                self._cache = cache_cls(**cache_kwargs)
+                if requires_cross_attention_cache:
+                    encoder_kwargs = cache_kwargs.copy()
+                    encoder_kwargs["max_cache_len"] = model_kwargs["encoder_outputs"][0].shape[1]
+                    self._cache = EncoderDecoderCache(self._cache, cache_cls(**encoder_kwargs))
+            else:
+                self._cache.reset()
+            return self._cache
+    
+        def _supports_default_dynamic_cache(self) -> bool:
+            """
+            Return `True` if current model can use a `DynamicCache` instance when initializing the `past_key_values`.
+            This is mostly the same as `_supports_cache_class` attribute, but add exception for `Jamba` model which
+            uses its own `HybridMambaAttentionDynamicCache` and do not need to initialize the Cache in advance in
+            order to save memory (because no back and forth `to_legacy_cache` and `from_legacy_cache` will be performed
+            for `HybridMambaAttentionDynamicCache`).
+            """
+            return (
+                self._supports_cache_class
+                and "jamba" not in self.__class__.__name__.lower()
+                and "zamba" not in self.__class__.__name__.lower()
+                and "bamba" not in self.__class__.__name__.lower()
+                and "minimax" not in self.__class__.__name__.lower()
+            )
+    
+        def _prepare_cache_for_generation(
+            self,
+            generation_config: GenerationConfig,
+            model_kwargs: dict,
+            assistant_model: "PreTrainedModel",
+            batch_size: int,
+            max_cache_length: int,
+            device: torch.device,
+        ) -> bool:
+            """
+            Prepares the cache for generation (if applicable), given `generate`'s parameterization. If a cache is
+            instantiated, writes it to `model_kwargs`, under the name expected by the model.
+            """
+    
+            is_hybrid_cache = any(class_name in self.__class__.__name__.lower() for class_name in ["mamba", "falconh1"])
+            cache_name = "past_key_values" if not is_hybrid_cache else "cache_params"
+    
+            requires_cross_attention_cache = (
+                self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None
+            )
+    
+            # Quick escape route 1: if the user specifies a cache, we only need to:
+            # a) check for conflicting `generate` arguments
+            # b) convert to the new cache format (if the user passes a legacy cache and model supports it)
+            user_defined_cache = model_kwargs.get(cache_name)
+            if user_defined_cache is not None:
+                if generation_config.cache_implementation is not None:
+                    raise ValueError(
+                        f"Passing both `cache_implementation` (used to initialize certain caches) and `{cache_name}` (a "
+                        "Cache object) is unsupported. Please use only one of the two."
+                    )
+                if isinstance(user_defined_cache, tuple) and self._supports_default_dynamic_cache():
+                    model_kwargs[cache_name] = (
+                        DynamicCache.from_legacy_cache(user_defined_cache)
+                        if not requires_cross_attention_cache
+                        else EncoderDecoderCache.from_legacy_cache(user_defined_cache)
+                    )
+                return
+    
+            # Quick escape route 2: if the user specifies no cache is to be used. (conflicting arguments are handled in
+            # `generation_config.validate()`)
+            if generation_config.use_cache is False:
+                return
+    
+            # Quick escape route 3: model that only supports legacy caches = nothing to prepare
+            if not self._supports_default_dynamic_cache():
+                if generation_config.cache_implementation is not None:
+                    warnings.warn(
+                        "This model does not support `Cache` instances, it only supports the legacy cache format (tuple "
+                        f"of tuples). `cache_implementation` (set to {generation_config.cache_implementation}) will be "
+                        "ignored.",
+                        UserWarning,
+                    )
+                return
+    
+            # Otherwise we NEED to prepare a cache, based on `generation_config.cache_implementation`
+    
+            # TODO(joao): support static caches in assisted generation. assisted generation needs to roll back caches,
+            # which is only supported in dynamic caches atm
+            if assistant_model is not None and generation_config.cache_implementation is not None:
+                logger.warning_once(
+                    "An assistant model is provided, using a dynamic cache instead of a cache of type="
+                    f"'{generation_config.cache_implementation}'."
+                )
+                generation_config.cache_implementation = None
+    
+            generation_config.cache_implementation = generation_config.cache_implementation or getattr(
+                self.config.get_text_config(), "cache_implementation", None
+            )
+            if generation_config.cache_implementation is not None:
+                if generation_config.cache_implementation in NEED_SETUP_CACHE_CLASSES_MAPPING:
+                    if generation_config.cache_implementation == "static" and not self._supports_static_cache:
+                        raise ValueError(
+                            "This model does not support `cache_implementation='static'`. Please check the following "
+                            "issue: https://github.com/huggingface/transformers/issues/28981"
+                        )
+                    model_kwargs[cache_name] = self._get_cache(
+                        cache_implementation=generation_config.cache_implementation,
+                        batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size,
+                        max_cache_len=max_cache_length,
+                        device=device,
+                        model_kwargs=model_kwargs,
+                    )
+                elif generation_config.cache_implementation == "quantized":
+                    if not self._supports_quantized_cache:
+                        raise ValueError(
+                            "This model does not support the quantized cache. If you want your model to support quantized "
+                            "cache, please open an issue and tag @zucchini-nlp."
+                        )
+    
+                    cache_config = (
+                        generation_config.cache_config
+                        if generation_config.cache_config is not None
+                        else QuantizedCacheConfig()
+                    )
+                    cache_class = QUANT_BACKEND_CLASSES_MAPPING[cache_config.backend]
+    
+                    if cache_config.backend == "quanto" and not is_optimum_quanto_available():
+                        raise ImportError(
+                            "You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto backend. "
+                            "Please install it via  with `pip install optimum-quanto`"
+                        )
+                    elif cache_config.backend == "HQQ" and not is_hqq_available():
+                        raise ImportError(
+                            "You need to install `HQQ` in order to use KV cache quantization with HQQ backend. "
+                            "Please install it via  with `pip install hqq`"
+                        )
+    
+                    model_kwargs[cache_name] = cache_class(cache_config)
+                elif generation_config.cache_implementation == "offloaded":
+                    model_kwargs[cache_name] = OffloadedCache()
+                elif generation_config.cache_implementation == "dynamic":
+                    model_kwargs[cache_name] = DynamicCache()
+    
+            # Use DynamicCache() instance by default. This will avoid back and forth from legacy format that
+            # keeps copying the cache thus using much more memory
+            else:
+                model_kwargs[cache_name] = (
+                    DynamicCache()
+                    if not requires_cross_attention_cache
+                    else EncoderDecoderCache(DynamicCache(), DynamicCache())
+                )
+    
+        def _supports_logits_to_keep(self) -> bool:
+            """
+            Return True if the current model supports the keyword argument `logits_to_keep` in forward()
+            to save memory. Checking it in this way allows to avoid using a new model attribute.
+            """
+            return "logits_to_keep" in set(inspect.signature(self.forward).parameters.keys())
+    
+        def _prepare_special_tokens(
+            self,
+            generation_config: GenerationConfig,
+            kwargs_has_attention_mask: Optional[bool] = None,
+            device: Optional[Union[torch.device, str]] = None,
+        ):
+            """
+            Prepares the special tokens for generation, overwriting the generation config with their processed versions
+            converted to tensor.
+    
+            Note that `generation_config` is changed in place and stops being serializable after this method is called.
+            That is no problem if called within `generate` (`generation_config` is a local copy that doesn't leave the
+            function). However, if called outside `generate`, consider creating a copy of `generation_config` first.
+            """
+    
+            # Convert special tokens to tensors
+            def _tensor_or_none(token, device=None):
+                if token is None:
+                    return token
+    
+                device = device if device is not None else self.device
+                if isinstance(token, torch.Tensor):
+                    return token.to(device)
+                return torch.tensor(token, device=device, dtype=torch.long)
+    
+            bos_token_tensor = _tensor_or_none(generation_config.bos_token_id, device=device)
+            eos_token_tensor = _tensor_or_none(generation_config.eos_token_id, device=device)
+            pad_token_tensor = _tensor_or_none(generation_config.pad_token_id, device=device)
+            decoder_start_token_tensor = _tensor_or_none(generation_config.decoder_start_token_id, device=device)
+    
+            # for BC we also try to get `decoder_start_token_id` or `bos_token_id` (#30892)
+            if self.config.is_encoder_decoder:
+                decoder_start_token_tensor = (
+                    decoder_start_token_tensor if decoder_start_token_tensor is not None else bos_token_tensor
+                )
+    
+            # We can have more than one eos token. Always treat it as a 1D tensor (when it exists).
+            if eos_token_tensor is not None and eos_token_tensor.ndim == 0:
+                eos_token_tensor = eos_token_tensor.unsqueeze(0)
+    
+            # Set pad token if unset (and there are conditions to do so)
+            if pad_token_tensor is None and eos_token_tensor is not None:
+                if kwargs_has_attention_mask is not None and not kwargs_has_attention_mask:
+                    logger.warning(
+                        "The attention mask and the pad token id were not set. As a consequence, you may observe "
+                        "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results."
+                    )
+                pad_token_tensor = eos_token_tensor[0]
+                logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{pad_token_tensor} for open-end generation.")
+    
+            # Sanity checks/warnings
+            if self.config.is_encoder_decoder and decoder_start_token_tensor is None:
+                raise ValueError(
+                    "`decoder_start_token_id` or `bos_token_id` has to be defined for encoder-decoder generation."
+                )
+            if (
+                eos_token_tensor is not None
+                and isin_mps_friendly(elements=eos_token_tensor, test_elements=pad_token_tensor).any()
+            ):
+                if kwargs_has_attention_mask is not None and not kwargs_has_attention_mask:
+                    logger.warning_once(
+                        "The attention mask is not set and cannot be inferred from input because pad token is same as "
+                        "eos token. As a consequence, you may observe unexpected behavior. Please pass your input's "
+                        "`attention_mask` to obtain reliable results."
+                    )
+            if eos_token_tensor is not None and (
+                torch.is_floating_point(eos_token_tensor) or (eos_token_tensor < 0).any()
+            ):
+                logger.warning(
+                    f"`eos_token_id` should consist of positive integers, but is {eos_token_tensor}. Your generation "
+                    "will not stop until the maximum length is reached. Depending on other flags, it may even crash."
+                )
+    
+            # Update generation config with the updated special tokens tensors
+            # NOTE: this must be written into a different attribute name than the one holding the original special tokens
+            # (in their non-tensor form), in order to enable end-to-end compilation. See
+            # https://pytorch.org/docs/stable/torch.compiler_cudagraph_trees.html#limitations
+            generation_config._bos_token_tensor = bos_token_tensor
+            generation_config._eos_token_tensor = eos_token_tensor
+            generation_config._pad_token_tensor = pad_token_tensor
+            generation_config._decoder_start_token_tensor = decoder_start_token_tensor
+    
+        def _valid_auto_compile_criteria(self, model_kwargs: dict, generation_config: GenerationConfig) -> bool:
+            """
+            Determines whether to trigger auto-compilation of the model's forward pass at generation time.
+            """
+            # Override: honor `disable_compile` flag
+            if generation_config.disable_compile:
+                return False
+    
+            # Base logic
+            valid_hardware = self.device.type == "cuda" or bool(
+                generation_config.compile_config is not None and generation_config.compile_config._compile_all_devices
+            )
+            using_compilable_cache = (
+                isinstance(model_kwargs.get("past_key_values"), Cache) and model_kwargs["past_key_values"].is_compileable
+            )
+            can_compile = valid_hardware and using_compilable_cache and self._supports_static_cache
+    
+            # Exception 1: Some quantization methods do not support compilation
+            if getattr(self, "hf_quantizer", None) is not None:
+                can_compile &= self.hf_quantizer.is_compileable
+    
+            if hasattr(self, "hf_device_map"):
+                all_model_devices = set(self.hf_device_map.values())
+                # Exception 2: Don't compile if the model is using CPU offload (as of April 2025, this results in a crash)
+                has_cpu_offload = "cpu" in all_model_devices and len(all_model_devices) > 1
+                can_compile &= not has_cpu_offload
+    
+                # Exception 3: Disk offload is not supported for compilation
+                has_disk_offload = "disk" in all_model_devices
+                can_compile &= not has_disk_offload
+    
+            # Finally: if the user has manually specified compilation options, but compilation is not possible, let's warn
+            # them
+            if generation_config.compile_config is not None and not can_compile:
+                logger.warning_once(
+                    "You have set `compile_config`, but we are unable to meet the criteria for compilation. Compilation "
+                    "will be skipped."
+                )
+    
+            return can_compile
+    
+        @torch.no_grad()
+        def generate(
+            self,
+            inputs: Optional[torch.Tensor] = None,
+            generation_config: Optional[GenerationConfig] = None,
+            logits_processor: Optional[LogitsProcessorList] = None,
+            stopping_criteria: Optional[StoppingCriteriaList] = None,
+            prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], list[int]]] = None,
+            synced_gpus: Optional[bool] = None,
+            assistant_model: Optional["PreTrainedModel"] = None,
+            streamer: Optional["BaseStreamer"] = None,
+            negative_prompt_ids: Optional[torch.Tensor] = None,
+            negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+            use_model_defaults: Optional[bool] = None,
+            custom_generate: Optional[str] = None,
+            **kwargs,
+        ) -> Union[GenerateOutput, torch.LongTensor]:
+            r"""
+    
+            Generates sequences of token ids for models with a language modeling head.
+    
+            <Tip warning={true}>
+    
+            Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+            model's default generation configuration. You can override any `generation_config` by passing the corresponding
+            parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`.
+    
+            For an overview of generation strategies and code examples, check out the [following
+            guide](../generation_strategies).
+    
+            </Tip>
+    
+            Parameters:
+                inputs (`torch.Tensor` of varying shape depending on the modality, *optional*):
+                    The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the
+                    method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+                    should be in the format of `input_ids`. For encoder-decoder models *inputs* can represent any of
+                    `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+                generation_config ([`~generation.GenerationConfig`], *optional*):
+                    The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+                    passed to generate matching the attributes of `generation_config` will override them. If
+                    `generation_config` is not provided, the default will be used, which has the following loading
+                    priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+                    configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+                    default values, whose documentation should be checked to parameterize generation.
+                logits_processor (`LogitsProcessorList`, *optional*):
+                    Custom logits processors that complement the default logits processors built from arguments and
+                    generation config. If a logit processor is passed that is already created with the arguments or a
+                    generation config an error is thrown. This feature is intended for advanced users.
+                stopping_criteria (`StoppingCriteriaList`, *optional*):
+                    Custom stopping criteria that complements the default stopping criteria built from arguments and a
+                    generation config. If a stopping criteria is passed that is already created with the arguments or a
+                    generation config an error is thrown. If your stopping criteria depends on the `scores` input, make
+                    sure you pass `return_dict_in_generate=True, output_scores=True` to `generate`. This feature is
+                    intended for advanced users.
+                prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], list[int]]`, *optional*):
+                    If provided, this function constraints the beam search to allowed tokens only at each step. If not
+                    provided no constraint is applied. This function takes 2 arguments: the batch ID `batch_id` and
+                    `input_ids`. It has to return a list with the allowed tokens for the next generation step conditioned
+                    on the batch ID `batch_id` and the previously generated tokens `inputs_ids`. This argument is useful
+                    for constrained generation conditioned on the prefix, as described in [Autoregressive Entity
+                    Retrieval](https://huggingface.co/papers/2010.00904).
+                synced_gpus (`bool`, *optional*):
+                    Whether to continue running the while loop until max_length. Unless overridden, this flag will be set
+                    to `True` if using `FullyShardedDataParallel` or DeepSpeed ZeRO Stage 3 with multiple GPUs to avoid
+                    deadlocking if one GPU finishes generating before other GPUs. Otherwise, defaults to `False`.
+                assistant_model (`PreTrainedModel`, *optional*):
+                    An assistant model that can be used to accelerate generation. The assistant model must have the exact
+                    same tokenizer. The acceleration is achieved when forecasting candidate tokens with the assistant model
+                    is much faster than running generation with the model you're calling generate from. As such, the
+                    assistant model should be much smaller.
+                streamer (`BaseStreamer`, *optional*):
+                    Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+                    through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+                negative_prompt_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                    The negative prompt needed for some processors such as CFG. The batch size must match the input batch
+                    size. This is an experimental feature, subject to breaking API changes in future versions.
+                negative_prompt_attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                    Attention_mask for `negative_prompt_ids`.
+                use_model_defaults (`bool`, *optional*):
+                    When it is `True`, unset parameters in `generation_config` will be set to the model-specific default
+                    generation configuration (`model.generation_config`), as opposed to the global defaults
+                    (`GenerationConfig()`). If unset, models saved starting from `v4.50` will consider this flag to be
+                    `True`.
+                custom_generate (`str`, *optional*):
+                    A string containing the name of a huggingface.co repository. If provided, the custom `generate`
+                    function defined in that reposity's `custom_generate/generate.py` file will be executed instead of the
+                    standard `generate` method. Note that the logic is for generation is entirely defined in that
+                    repository, and the return type may be different from the standard `generate` method.
+                kwargs (`dict[str, Any]`, *optional*):
+                    Ad hoc parametrization of `generation_config` and/or additional model-specific kwargs that will be
+                    forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
+                    specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.
+    
+            Return:
+                [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True`
+                or when `config.return_dict_in_generate=True`) or a `torch.LongTensor`.
+    
+                    If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible
+                    [`~utils.ModelOutput`] types are:
+    
+                        - [`~generation.GenerateDecoderOnlyOutput`],
+                        - [`~generation.GenerateBeamDecoderOnlyOutput`]
+    
+                    If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible
+                    [`~utils.ModelOutput`] types are:
+    
+                        - [`~generation.GenerateEncoderDecoderOutput`],
+                        - [`~generation.GenerateBeamEncoderDecoderOutput`]
+            """
+            # 0. If requested, load an arbitrary generation recipe from the Hub and run it instead
+            trust_remote_code = kwargs.pop("trust_remote_code", None)
+            if custom_generate is not None:
+                # Get all `generate` arguments in a single variable. Custom functions are responsible for handling them:
+                # they receive the same inputs as `generate`, with `model` instead of `self` and excluding the arguments to
+                # trigger the custom generation. They can access to methods from `GenerationMixin` through `model`.
+                global_keys_to_exclude = {
+                    "self",
+                    "kwargs",
+                    "global_keys_to_exclude",
+                    "trust_remote_code",
+                    "custom_generate",
+                }
+                generate_arguments = {key: value for key, value in locals().items() if key not in global_keys_to_exclude}
+                generate_arguments.update(kwargs)
+    
+                custom_generate_function = self.load_custom_generate(
+                    custom_generate, trust_remote_code=trust_remote_code, **kwargs
+                )
+                return custom_generate_function(model=self, **generate_arguments)
+    
+            # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
+            tokenizer = kwargs.pop("tokenizer", None)  # Pull this out first, we only use it for stopping criteria
+            assistant_tokenizer = kwargs.pop("assistant_tokenizer", None)  # only used for assisted generation
+    
+            generation_config, model_kwargs = self._prepare_generation_config(
+                generation_config, use_model_defaults, **kwargs
+            )
+            self._validate_model_kwargs(model_kwargs.copy())
+            self._validate_assistant(assistant_model, tokenizer, assistant_tokenizer)
+    
+            # 2. Set generation parameters if not already defined
+            if synced_gpus is None:
+                synced_gpus = (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and dist.get_world_size() > 1
+    
+            logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+            stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+    
+            accepts_attention_mask = "attention_mask" in set(inspect.signature(self.forward).parameters.keys())
+            requires_attention_mask = "encoder_outputs" not in model_kwargs
+            kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
+    
+            # 3. Define model inputs
+            inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
+                inputs, generation_config.bos_token_id, model_kwargs
+            )
+            batch_size = inputs_tensor.shape[0]
+    
+            device = inputs_tensor.device
+            self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=device)
+    
+            # decoder-only models must use left-padding for batched generation.
+            if not self.config.is_encoder_decoder:
+                # If `input_ids` was given, check if the last id in any sequence is `pad_token_id`
+                # Note: If using, `inputs_embeds` this check does not work, because we want to be more hands-off.
+                if (
+                    generation_config._pad_token_tensor is not None
+                    and batch_size > 1
+                    and len(inputs_tensor.shape) == 2
+                    and torch.sum(inputs_tensor[:, -1] == generation_config._pad_token_tensor) > 0
+                ):
+                    logger.warning(
+                        "A decoder-only architecture is being used, but right-padding was detected! For correct "
+                        "generation results, please set `padding_side='left'` when initializing the tokenizer."
+                    )
+    
+            # 4. Define other model kwargs
+            # decoder-only models with inputs_embeds forwarding must use caching (otherwise we can't detect whether we are
+            # generating the first new token or not, and we only want to use the embeddings for the first new token)
+            if not self.config.is_encoder_decoder and model_input_name == "inputs_embeds":
+                generation_config.use_cache = True
+    
+            if not kwargs_has_attention_mask and requires_attention_mask and accepts_attention_mask:
+                model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
+                    inputs_tensor, generation_config, model_kwargs
+                )
+            elif kwargs_has_attention_mask:
+                # TODO (joao): generalize this check with other types of inputs
+                if model_input_name == "input_ids" and len(model_kwargs["attention_mask"].shape) > 2:
+                    raise ValueError("`attention_mask` passed to `generate` must be 2D.")
+    
+            if self.config.is_encoder_decoder and "encoder_outputs" not in model_kwargs:
+                # if model is encoder decoder encoder_outputs are created and added to `model_kwargs`
+                model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
+                    inputs_tensor, model_kwargs, model_input_name, generation_config
+                )
+    
+            # 5. Prepare `input_ids` which will be used for auto-regressive generation
+            if self.config.is_encoder_decoder:
+                input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation(
+                    batch_size=batch_size,
+                    model_input_name=model_input_name,
+                    model_kwargs=model_kwargs,
+                    decoder_start_token_id=generation_config._decoder_start_token_tensor,
+                    device=inputs_tensor.device,
+                )
+            else:
+                input_ids = inputs_tensor if model_input_name == "input_ids" else model_kwargs.pop("input_ids")
+    
+            if generation_config.token_healing:
+                input_ids = self.heal_tokens(input_ids, tokenizer)
+    
+            if streamer is not None:
+                streamer.put(input_ids.cpu())
+    
+            # 6. Prepare `max_length` depending on other stopping criteria.
+            input_ids_length = input_ids.shape[1]
+            has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+            has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
+            generation_config = self._prepare_generated_length(
+                generation_config=generation_config,
+                has_default_max_length=has_default_max_length,
+                has_default_min_length=has_default_min_length,
+                model_input_name=model_input_name,
+                inputs_tensor=inputs_tensor,
+                input_ids_length=input_ids_length,
+            )
+    
+            # If the model supports `logits_to_keep` in forward(), set it to 1 to avoid computing the whole
+            # logit matrix. This can save a lot of memory during the first forward pass. Note that assisted decoding
+            # dynamically overrides this value as it can need more than the last token logits
+            if self._supports_logits_to_keep() and "logits_to_keep" not in model_kwargs:
+                model_kwargs["logits_to_keep"] = 1
+    
+            self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)
+    
+            # 7. Prepare the cache.
+            # - `model_kwargs` may be updated in place with a cache as defined by the parameters in `generation_config`.
+            # - different models have a different cache name expected by the model (default = "past_key_values")
+            # - `max_length`, prepared above, is used to determine the maximum cache length
+            max_cache_length = generation_config.max_length - 1
+            if (
+                inputs_tensor.shape[1] != input_ids_length
+                and model_input_name == "inputs_embeds"
+                and not self.config.is_encoder_decoder
+            ):
+                max_cache_length += inputs_tensor.shape[1]
+            self._prepare_cache_for_generation(
+                generation_config, model_kwargs, assistant_model, batch_size, max_cache_length, device
+            )
+    
+            # 8. determine generation mode
+            generation_mode = generation_config.get_generation_mode(assistant_model)
+    
+            if streamer is not None and (generation_config.num_beams > 1):
+                raise ValueError(
+                    "`streamer` cannot be used with beam search (yet!). Make sure that `num_beams` is set to 1."
+                )
+    
+            if self.device.type != input_ids.device.type:
+                warnings.warn(
+                    "You are calling .generate() with the `input_ids` being on a device type different"
+                    f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
+                    f" is on {self.device.type}. You may experience unexpected behaviors or slower generation."
+                    " Please make sure that you have put `input_ids` to the"
+                    f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"
+                    " running `.generate()`.",
+                    UserWarning,
+                )
+    
+            # 9. prepare logits processors and stopping criteria
+            prepared_logits_processor = self._get_logits_processor(
+                generation_config=generation_config,
+                input_ids_seq_length=input_ids_length,
+                encoder_input_ids=inputs_tensor,
+                prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+                logits_processor=logits_processor,
+                device=inputs_tensor.device,
+                model_kwargs=model_kwargs,
+                negative_prompt_ids=negative_prompt_ids,
+                negative_prompt_attention_mask=negative_prompt_attention_mask,
+            )
+            prepared_stopping_criteria = self._get_stopping_criteria(
+                generation_config=generation_config, stopping_criteria=stopping_criteria, tokenizer=tokenizer, **kwargs
+            )
+    
+            # Set model_kwargs `use_cache` so we can use it later in forward runs
+            model_kwargs["use_cache"] = generation_config.use_cache
+    
+            # 10. go into different generation modes
+            if generation_mode == GenerationMode.ASSISTED_GENERATION:
+                if generation_config.num_return_sequences > 1:
+                    raise ValueError(
+                        "num_return_sequences has to be 1 when doing assisted generate, "
+                        f"but is {generation_config.num_return_sequences}."
+                    )
+                if batch_size > 1:
+                    raise ValueError("assisted generate is only supported for batch_size = 1")
+                if not model_kwargs["use_cache"]:
+                    raise ValueError("assisted generate requires `use_cache=True`")
+                if generation_config.cache_implementation in ["static", "hybrid", "sliding_window"]:
+                    raise ValueError("assisted generate is not supported with Static cache classes`")
+                if self._is_stateful:
+                    # In assisted generation we need the ability to confirm whether the model would pick certain tokens,
+                    # which is not possible with stateful models (they can't reset to a previous subset of generated text)
+                    raise ValueError(
+                        f"assisted generation is not supported with stateful models, such as {self.__class__.__name__}"
+                    )
+    
+                # 11. Get the candidate generator, given the parameterization
+                candidate_generator = self._get_candidate_generator(
+                    generation_config=generation_config,
+                    input_ids=input_ids,
+                    inputs_tensor=inputs_tensor,
+                    assistant_model=assistant_model,
+                    logits_processor=logits_processor,
+                    target_tokenizer=tokenizer,
+                    assistant_tokenizer=assistant_tokenizer,
+                    model_kwargs=model_kwargs,
+                )
+    
+                # 12. run assisted generate
+                result = self._assisted_decoding(
+                    input_ids,
+                    candidate_generator=candidate_generator,
+                    logits_processor=prepared_logits_processor,
+                    stopping_criteria=prepared_stopping_criteria,
+                    generation_config=generation_config,
+                    synced_gpus=synced_gpus,
+                    streamer=streamer,
+                    **model_kwargs,
+                )
+            elif generation_mode == GenerationMode.DOLA_GENERATION:
+                if not trust_remote_code:
+                    logger.warning_once(
+                        "DoLa Decoding is scheduled to be moved to a `custom_generate` repository in v4.55.0. "
+                        "To prevent loss of backward compatibility, add `trust_remote_code=True` to your `generate` call."
+                    )
+                if self._is_stateful:
+                    # DoLa decoding was not designed for stateful models, and would require some changes
+                    raise ValueError(
+                        f"dola decoding is not supported with stateful models, such as {self.__class__.__name__}"
+                    )
+                result = self._dola_decoding(
+                    input_ids,
+                    dola_layers=generation_config.dola_layers,
+                    logits_processor=prepared_logits_processor,
+                    stopping_criteria=prepared_stopping_criteria,
+                    generation_config=generation_config,
+                    synced_gpus=synced_gpus,
+                    streamer=streamer,
+                    **model_kwargs,
+                )
+    
+            elif generation_mode == GenerationMode.CONTRASTIVE_SEARCH:
+                if not trust_remote_code:
+                    logger.warning_once(
+                        "Contrastive Search is scheduled to be moved to a `custom_generate` repository in v4.55.0. "
+                        "To prevent loss of backward compatibility, add `trust_remote_code=True` to your `generate` call."
+                    )
+                if not model_kwargs["use_cache"]:
+                    raise ValueError("Contrastive search requires `use_cache=True`")
+                if self._is_stateful:
+                    # Just like assisted generation, we need to be able to rollback to a previous state (see comment above)
+                    raise ValueError(
+                        f"contrastive search is not supported with stateful models, such as {self.__class__.__name__}"
+                    )
+    
+                result = self._contrastive_search(
+                    input_ids,
+                    logits_processor=prepared_logits_processor,
+                    stopping_criteria=prepared_stopping_criteria,
+                    generation_config=generation_config,
+                    synced_gpus=synced_gpus,
+                    streamer=streamer,
+                    **model_kwargs,
+                )
+    
+            elif generation_mode in (GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH):
+                # 11. expand input_ids with `num_return_sequences` additional sequences per batch
+                input_ids, model_kwargs = self._expand_inputs_for_generation(
+                    input_ids=input_ids,
+                    expand_size=generation_config.num_return_sequences,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                    **model_kwargs,
+                )
+    
+                # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
+                result = self._sample(
+                    input_ids,
+                    logits_processor=prepared_logits_processor,
+                    stopping_criteria=prepared_stopping_criteria,
+                    generation_config=generation_config,
+                    synced_gpus=synced_gpus,
+                    streamer=streamer,
+                    **model_kwargs,
+                )
+    
+            elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH):
+                # 11. interleave input_ids with `num_beams` additional sequences per batch
+                input_ids, model_kwargs = self._expand_inputs_for_generation(
+                    input_ids=input_ids,
+                    expand_size=generation_config.num_beams,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                    **model_kwargs,
+                )
+                # 12. run beam sample
+                result = self._beam_search(
+                    input_ids,
+                    logits_processor=prepared_logits_processor,
+                    stopping_criteria=prepared_stopping_criteria,
+                    generation_config=generation_config,
+                    synced_gpus=synced_gpus,
+                    **model_kwargs,
+                )
+    
+            elif generation_mode == GenerationMode.GROUP_BEAM_SEARCH:
+                logger.warning_once(
+                    "Group Beam Search is scheduled to be moved to a `custom_generate` repository in v4.55.0. "
+                    "To prevent loss of backward compatibility, add `trust_remote_code=True` to your `generate` call."
+                )
+                # 11. prepare beam search scorer
+                beam_scorer = BeamSearchScorer(
+                    batch_size=batch_size,
+                    num_beams=generation_config.num_beams,
+                    device=inputs_tensor.device,
+                    length_penalty=generation_config.length_penalty,
+                    do_early_stopping=generation_config.early_stopping,
+                    num_beam_hyps_to_keep=generation_config.num_return_sequences,
+                    num_beam_groups=generation_config.num_beam_groups,
+                    max_length=generation_config.max_length,
+                )
+                # 12. interleave input_ids with `num_beams` additional sequences per batch
+                input_ids, model_kwargs = self._expand_inputs_for_generation(
+                    input_ids=input_ids,
+                    expand_size=generation_config.num_beams,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                    **model_kwargs,
+                )
+                # 13. run beam search
+                result = self._group_beam_search(
+                    input_ids,
+                    beam_scorer,
+                    logits_processor=prepared_logits_processor,
+                    stopping_criteria=prepared_stopping_criteria,
+                    generation_config=generation_config,
+                    synced_gpus=synced_gpus,
+                    **model_kwargs,
+                )
+    
+            elif generation_mode == GenerationMode.CONSTRAINED_BEAM_SEARCH:
+                logger.warning_once(
+                    "Constrained Beam Search is scheduled to be moved to a `custom_generate` repository in v4.55.0. "
+                    "To prevent loss of backward compatibility, add `trust_remote_code=True` to your `generate` call."
+                )
+                final_constraints = []
+                if generation_config.constraints is not None:
+                    final_constraints = generation_config.constraints
+    
+                if generation_config.force_words_ids is not None:
+    
+                    def typeerror():
+                        raise ValueError(
+                            "`force_words_ids` has to either be a `list[list[list[int]]]` or `list[list[int]]` "
+                            f"of positive integers, but is {generation_config.force_words_ids}."
+                        )
+    
+                    if (
+                        not isinstance(generation_config.force_words_ids, list)
+                        or len(generation_config.force_words_ids) == 0
+                    ):
+                        typeerror()
+    
+                    for word_ids in generation_config.force_words_ids:
+                        if isinstance(word_ids[0], list):
+                            if not isinstance(word_ids, list) or len(word_ids) == 0:
+                                typeerror()
+                            if any(not isinstance(token_ids, list) for token_ids in word_ids):
+                                typeerror()
+                            if any(
+                                any((not isinstance(token_id, int) or token_id < 0) for token_id in token_ids)
+                                for token_ids in word_ids
+                            ):
+                                typeerror()
+    
+                            constraint = DisjunctiveConstraint(word_ids)
+                        else:
+                            if not isinstance(word_ids, list) or len(word_ids) == 0:
+                                typeerror()
+                            if any((not isinstance(token_id, int) or token_id < 0) for token_id in word_ids):
+                                typeerror()
+    
+                            constraint = PhrasalConstraint(word_ids)
+                        final_constraints.append(constraint)
+    
+                # 11. prepare beam search scorer
+                constrained_beam_scorer = ConstrainedBeamSearchScorer(
+                    constraints=final_constraints,
+                    batch_size=batch_size,
+                    num_beams=generation_config.num_beams,
+                    device=inputs_tensor.device,
+                    length_penalty=generation_config.length_penalty,
+                    do_early_stopping=generation_config.early_stopping,
+                    num_beam_hyps_to_keep=generation_config.num_return_sequences,
+                    max_length=generation_config.max_length,
+                )
+                # 12. interleave input_ids with `num_beams` additional sequences per batch
+                input_ids, model_kwargs = self._expand_inputs_for_generation(
+                    input_ids=input_ids,
+                    expand_size=generation_config.num_beams,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                    **model_kwargs,
+                )
+                # 13. run beam search
+                result = self._constrained_beam_search(
+                    input_ids,
+                    constrained_beam_scorer=constrained_beam_scorer,
+                    logits_processor=prepared_logits_processor,
+                    stopping_criteria=prepared_stopping_criteria,
+                    generation_config=generation_config,
+                    synced_gpus=synced_gpus,
+                    **model_kwargs,
+                )
+    
+            # Convert to legacy cache format if requested
+            if (
+                generation_config.return_legacy_cache is True
+                and hasattr(result, "past_key_values")
+                and getattr(result.past_key_values, "to_legacy_cache") is not None
+            ):
+                result.past_key_values = result.past_key_values.to_legacy_cache()
+            return result
+    
+        def _has_unfinished_sequences(self, this_peer_finished: bool, synced_gpus: bool, device: torch.device) -> bool:
+            """
+            Returns whether there are still unfinished sequences in the device. The existence of unfinished sequences is
+            fed through `this_peer_finished`. ZeRO stage 3-friendly.
+            """
+            if synced_gpus:
+                # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
+                # The following logic allows an early break if all peers finished generating their sequence
+                this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0, device=device)
+                # send 0.0 if we finished, 1.0 otherwise
+                dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
+                # did all peers finish? the reduced sum will be 0.0 then
+                if this_peer_finished_flag.item() == 0.0:
+                    return False
+            elif this_peer_finished:
+                return False
+            return True
+    
+        def heal_tokens(
+            self, input_ids: torch.LongTensor, tokenizer: Optional["PreTrainedTokenizerBase"] = None
+        ) -> torch.LongTensor:
+            r"""
+            Generates sequences of token ids for models with a language modeling head.
+            Parameters:
+                input_ids (`torch.LongTensor`): The sequence used as a prompt for the generation.
+                tokenizer (`PreTrainedTokenizerBase`, *optional*): The tokenizer used to decode the input ids.
+            Return:
+                `torch.LongTensor` where each sequence has its tail token replaced with its appropriate extension.
+            """
+            if tokenizer is None:
+                raise ValueError(
+                    " When generating with token healing, you must pass the model's tokenizer to the `tokenizer` "
+                    "argument of `generate`."
+                )
+    
+            bos_token_id, pad_token_id = tokenizer.bos_token_id, tokenizer.pad_token_id
+            vocab_trie = ExtensionsTrie(tokenizer.get_vocab())
+            generation_config = GenerationConfig(max_new_tokens=1, pad_token_id=pad_token_id)
+    
+            # assumption: leading/trailing whitespace is not meaningful, so the prompts are
+            # stripped before re-tokenizing to desensitize generation to whitespace artefacts
+            prompts = [p.strip() for p in tokenizer.batch_decode(input_ids, skip_special_tokens=True)]
+            input_ids = tokenizer(
+                prompts,
+                return_tensors="pt",
+                padding=True,
+            ).input_ids.to(input_ids.device)
+    
+            # replace bos with pad to not condition healing on it
+            input_ids = torch.where(input_ids == bos_token_id, pad_token_id, input_ids)
+    
+            """
+            the latter code assumes the input_ids is not empty,
+            input_id has to be checked if contains elements
+            """
+            if input_ids.numel() == 0:
+                return input_ids
+    
+            tail_ids = input_ids[:, -1].tolist()
+    
+            space_tok = tokenizer.convert_ids_to_tokens(tokenizer.convert_tokens_to_ids(" "))[0]
+            # tail tokens are used for a prefix search, thus, whitespaces are replaced with
+            # their tokenization (e.g. 'Ġ') to enable search for tokens prefixed with a whitespace
+            tail_toks = (tokenizer.decode(t).replace(" ", space_tok) for t in tail_ids)
+    
+            for batch_idx, (tail_id, tail_tok) in enumerate(zip(tail_ids, tail_toks)):
+                batch_ids = input_ids[batch_idx]
+                if torch.all(batch_ids == pad_token_id).item():
+                    continue  # skip empty sequences (all pad ids)
+    
+                # apply bias for alternatives (extensions) to the tail token
+                """
+                seq_bias key has to be tuple with int so have to use
+                tokenizer function to convert str to int
+                """
+                seq_bias = {
+                    (tokenizer.convert_tokens_to_ids(alt_tok),): 10.0 for alt_tok in vocab_trie.extensions(prefix=tail_tok)
+                }
+    
+                if len(seq_bias) == 1:
+                    continue  # skip if there are no token alternatives to heal with
+    
+                # slightly favor original token to limit aggressive healing e.g. 'http' -> 'https'
+                seq_bias[(tail_id,)] += 1.0
+                generation_config.update(sequence_bias=seq_bias)
+    
+                trimmed_ids = batch_ids[:-1]
+    
+                """
+                the latter code assumes trimmed_ids is not empty
+                so have to check the its element count
+                """
+                if trimmed_ids.numel() == 0:
+                    continue
+    
+                # if the prompt is a single (non-pad) token, regenerate from bos
+                if len(batch_ids[batch_ids != pad_token_id]) == 1:
+                    trimmed_ids[-1] = bos_token_id
+    
+                input_ids[batch_idx] = self.generate(trimmed_ids.unsqueeze(0), generation_config=generation_config)
+    
+            return input_ids
+    
+        def _dola_decoding(
+            self,
+            input_ids: torch.LongTensor,
+            dola_layers: Union[str, list[int]],
+            logits_processor: LogitsProcessorList,
+            stopping_criteria: StoppingCriteriaList,
+            generation_config: GenerationConfig,
+            synced_gpus: bool,
+            streamer: "BaseStreamer",
+            **model_kwargs,
+        ) -> Union[GenerateNonBeamOutput, torch.LongTensor]:
+            r"""
+            Generates sequences of token ids for models with a language modeling head using **dola decoding** and can be
+            used for decoder-only text models.
+            The method is based on the paper "DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language
+            Models" (https://huggingface.co/papers/2309.03883) in ICLR 2024.
+    
+            Parameters:
+                input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                    The sequence used as a prompt for the generation.
+                dola_layers (`Union[str, list[int]]`):
+                    The candidate layers used in contrasting layers of DoLa. It can be either 1) 'low' or 'high', which
+                    means the lower part or higher part of the model layers, respectively, or 2) a list of layer indices
+                    to be used for candidate layers. The 0-th layer is the word embedding layer of the model.
+                logits_processor (`LogitsProcessorList`):
+                    An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
+                    used to modify the prediction scores of the language modeling head applied at each generation step.
+                stopping_criteria (`StoppingCriteriaList`, *optional*):
+                    An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
+                    used to tell if the generation loop should stop.
+                generation_config ([`~generation.GenerationConfig`]):
+                    The generation configuration to be used as parametrization of the decoding method.
+                synced_gpus (`bool`):
+                    Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+                    `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+                streamer (`BaseStreamer`, *optional*):
+                    Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+                    through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+                model_kwargs:
+                    Additional model specific keyword arguments will be forwarded to the `forward` function of the model.
+                    If model is an encoder-decoder model the kwargs should include `encoder_outputs`.
+    
+            Return:
+                [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`]
+                or `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a
+                [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
+                `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if
+                `model.config.is_encoder_decoder=True`.
+            """
+    
+            if self.config.is_encoder_decoder:
+                raise ValueError("DoLa decoding is only available for decoder-only models.")
+            # init values
+    
+            pad_token_id = generation_config._pad_token_tensor
+            output_attentions = generation_config.output_attentions
+            output_hidden_states = generation_config.output_hidden_states
+            output_scores = generation_config.output_scores
+            output_logits = generation_config.output_logits
+            return_dict_in_generate = generation_config.return_dict_in_generate
+            has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria)
+            do_sample = generation_config.do_sample
+    
+            # init attention / hidden states / scores tuples
+            scores = () if (return_dict_in_generate and output_scores) else None
+            raw_logits = () if (return_dict_in_generate and output_logits) else None
+            decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+            cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+            decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+    
+            # keep track of which sequences are already finished
+            batch_size, cur_length = input_ids.shape[:2]
+            unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+            model_kwargs = self._get_initial_cache_position(cur_length, input_ids.device, model_kwargs)
+    
+            this_peer_finished = False
+    
+            # prepare layers for DoLa decoding
+            final_layer = self.config.get_text_config().num_hidden_layers
+            # if the model has tied word embeddings, we skip the word embeddings (0-th) layer and start from the 2nd layer,
+            # as the early exit from word embeddings will become identity function
+            # if the model is really shallow (<=2 layers), we use the 1st layer if it's not the final layer and the 0-th
+            # layer otherwise. Notice that DoLa does not help shallow models much.
+            if not self.config.tie_word_embeddings:
+                start_layer = 0
+            elif final_layer > 2:
+                start_layer = 2
+            elif final_layer == 2:
+                start_layer = 1
+            else:
+                start_layer = 0
+    
+            # For `N`-layer models with `N <= 40` layers, the layers of `range(0, N // 2, 2)` and `range(N // 2, N, 2)`
+            # are used for `'low'` and `'high'` layers, respectively.
+            # For models with `N > 40` layers, the layers of `range(0, 20, 2)` and `range(N - 20, N, 2)` are used for
+            # `'low'` and `'high'` layers, respectively.
+            if isinstance(dola_layers, str) and dola_layers == "low":
+                if start_layer == final_layer // 2:
+                    candidate_premature_layers = [start_layer]
+                else:
+                    candidate_premature_layers = (
+                        list(range(start_layer, final_layer // 2, 2))
+                        if final_layer <= 40
+                        else list(range(start_layer, 20, 2))
+                    )
+            elif isinstance(dola_layers, str) and dola_layers == "high":
+                candidate_premature_layers = (
+                    list(range(final_layer // 2, final_layer, 2))
+                    if final_layer <= 40
+                    else list(range(final_layer - 20, final_layer, 2))
+                )
+            # Set the `dola_layers` to a list of integers for layer indices to contrast manually specified layers.
+            elif isinstance(dola_layers, list):
+                candidate_premature_layers = [i for i in dola_layers if i < final_layer]
+            else:
+                raise ValueError("dola_layers must be either 'low', 'high' or a list of integers.")
+    
+            lm_head = self.get_output_embeddings()
+            if lm_head is None:
+                raise ValueError("DoLa is not supported for models that don't have output embeddings.")
+    
+            while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+                # prepare model inputs
+                model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+    
+                # forward pass to get next token
+                outputs = self(
+                    **model_inputs,
+                    return_dict=True,
+                    output_attentions=output_attentions,
+                    output_hidden_states=True,
+                )
+    
+                # .float() is needed to retain precision for later logits manipulations
+                final_layer_next_token_logits = outputs.logits[:, -1, :].detach().to(copy=True, dtype=torch.float32)
+                final_logits = outputs.logits[:, -1, :].float()
+                candidate_premature_logits = {}
+                for candidate_premature_layer in candidate_premature_layers:
+                    candidate_premature_logits[candidate_premature_layer] = lm_head(
+                        outputs.hidden_states[candidate_premature_layer][:, -1, :]
+                    ).to(final_logits.device)
+    
+                # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+                model_kwargs = self._update_model_kwargs_for_generation(
+                    outputs,
+                    model_kwargs,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                )
+                if synced_gpus and this_peer_finished:
+                    continue
+    
+                next_token_logits = _dola_select_contrast(
+                    candidate_premature_layers, candidate_premature_logits, final_logits
+                )
+                next_token_logits = next_token_logits.to(input_ids.device)
+                # pre-process distribution
+                next_token_scores = logits_processor(input_ids, next_token_logits)
+    
+                # Store scores, attentions and hidden_states when required
+                if return_dict_in_generate:
+                    if output_scores:
+                        scores += (next_token_scores,)
+                    if output_logits:
+                        raw_logits += (final_layer_next_token_logits,)
+                    if output_attentions:
+                        decoder_attentions += (
+                            (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
+                        )
+                        if self.config.is_encoder_decoder:
+                            cross_attentions += (outputs.cross_attentions,)
+    
+                    if output_hidden_states:
+                        decoder_hidden_states += (
+                            (outputs.decoder_hidden_states,)
+                            if self.config.is_encoder_decoder
+                            else (outputs.hidden_states,)
+                        )
+    
+                if do_sample:  # sample
+                    probs = nn.functional.softmax(next_token_scores, dim=-1)
+                    next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+                else:  # argmax
+                    next_tokens = torch.argmax(next_token_scores, dim=-1)
+    
+                # finished sentences should have their next token be a padding token
+                if has_eos_stopping_criteria:
+                    next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+    
+                # update generated ids, model inputs, and length for next step
+                input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
+                if streamer is not None:
+                    streamer.put(next_tokens.cpu())
+    
+                # stop when each sentence is finished
+                unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores)
+                this_peer_finished = unfinished_sequences.max() == 0
+    
+            if streamer is not None:
+                streamer.end()
+    
+            if return_dict_in_generate:
+                return GenerateDecoderOnlyOutput(
+                    sequences=input_ids,
+                    scores=scores,
+                    logits=raw_logits,
+                    attentions=decoder_attentions,
+                    hidden_states=decoder_hidden_states,
+                    past_key_values=model_kwargs.get("past_key_values"),
+                )
+            else:
+                return input_ids
+    
+        @torch.no_grad()
+        def _contrastive_search(
+            self,
+            input_ids: torch.LongTensor,
+            logits_processor: LogitsProcessorList,
+            stopping_criteria: StoppingCriteriaList,
+            generation_config: GenerationConfig,
+            synced_gpus: bool,
+            streamer: Optional["BaseStreamer"],
+            **model_kwargs,
+        ) -> Union[GenerateNonBeamOutput, torch.LongTensor]:
+            r"""
+            Generates sequences of token ids for models with a language modeling head using **contrastive search** and can
+            be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.
+    
+            Parameters:
+                input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                    The sequence used as a prompt for the generation.
+                logits_processor (`LogitsProcessorList`):
+                    An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
+                    used to modify the prediction scores of the language modeling head applied at each generation step.
+                stopping_criteria (`StoppingCriteriaList`):
+                    An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
+                    used to tell if the generation loop should stop.
+                generation_config ([`~generation.GenerationConfig`]):
+                    The generation configuration to be used as parametrization of the decoding method.
+                synced_gpus (`bool`):
+                    Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+                    `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+                streamer (`BaseStreamer`, *optional*):
+                    Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+                    through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+                model_kwargs:
+                    Additional model specific keyword arguments will be forwarded to the `forward` function of the model.
+                    If model is an encoder-decoder model the kwargs should include `encoder_outputs`.
+    
+            Return:
+                [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`]
+                or `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a
+                [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
+                `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if
+                `model.config.is_encoder_decoder=True`.
+            """
+            # init values
+            has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria)
+            top_k = generation_config.top_k
+            penalty_alpha = generation_config.penalty_alpha
+            pad_token_id = generation_config._pad_token_tensor
+            output_attentions = generation_config.output_attentions
+            output_hidden_states = generation_config.output_hidden_states
+            output_scores = generation_config.output_scores
+            output_logits = generation_config.output_logits
+            return_dict_in_generate = generation_config.return_dict_in_generate
+            sequential = generation_config.low_memory
+    
+            # init attention / hidden states / scores tuples
+            raw_logits = () if (return_dict_in_generate and output_logits) else None
+            scores = () if (return_dict_in_generate and output_scores) else None
+            decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+            cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+            decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+    
+            # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+            if return_dict_in_generate and self.config.is_encoder_decoder:
+                encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
+                encoder_hidden_states = (
+                    model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+                )
+    
+            # keep track of which sequences are already finished
+            batch_size, cur_len = input_ids.shape[:2]
+            unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+            model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+    
+            # Create cosine_matrix_mask based on the attention_mask
+            cosine_matrix_mask = torch.ones_like(input_ids, dtype=torch.long)
+            if self.config.is_encoder_decoder:
+                if "decoder_attention_mask" in model_kwargs and model_kwargs["decoder_attention_mask"] is not None:
+                    cosine_matrix_mask = model_kwargs["decoder_attention_mask"]
+            else:
+                cosine_matrix_mask = model_kwargs["attention_mask"]
+            cosine_matrix_mask = cosine_matrix_mask.repeat_interleave(top_k, dim=0)
+    
+            this_peer_finished = False
+    
+            while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+                # if the first step in the loop, encode all the prefix and obtain: (1) past_key_values;
+                # (2) last_hidden_states; (3) logit_for_next_step; (4) update model kwargs for the next step
+                if model_kwargs.get("past_key_values") is None or (
+                    isinstance(model_kwargs["past_key_values"], (Cache, EncoderDecoderCache))
+                    and model_kwargs["past_key_values"].get_seq_length() == 0
+                ):
+                    # prepare inputs
+                    model_kwargs["use_cache"] = True
+                    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+    
+                    # encode the given prefix and prepare model inputs; encoder-decoder model process the prefix and save
+                    # the `encoder_outputs`
+                    outputs = self(
+                        **model_inputs, return_dict=True, output_hidden_states=True, output_attentions=output_attentions
+                    )
+    
+                    # last decoder hidden states will be used to compute the degeneration penalty (cosine similarity with
+                    # previous tokens)
+                    if self.config.is_encoder_decoder:
+                        last_hidden_states = outputs.decoder_hidden_states[-1]
+                    else:
+                        last_hidden_states = outputs.hidden_states[-1]
+    
+                    # next logit for contrastive search to select top-k candidate tokens
+                    # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for this first iteration
+                    # (the clone itself is always small)
+                    # torch.float32 is needed to retain precision for later logits manipulations
+                    logit_for_next_step = outputs.logits[:, -1, :].to(
+                        copy=True, dtype=torch.float32, device=input_ids.device
+                    )
+    
+                    model_kwargs = self._update_model_kwargs_for_generation(
+                        outputs,
+                        model_kwargs,
+                        is_encoder_decoder=self.config.is_encoder_decoder,
+                    )
+    
+                    if not sequential:
+                        # Expands model inputs top_k times, for batched forward passes (akin to beam search).
+                        # input_ids is required for expanding visual inputs in qwen2vl
+                        _, model_kwargs = self._expand_inputs_for_generation(
+                            input_ids=input_ids,
+                            expand_size=top_k,
+                            is_encoder_decoder=self.config.is_encoder_decoder,
+                            **model_kwargs,
+                        )
+    
+                    past_key_values = model_kwargs.get("past_key_values")
+                    if past_key_values is None:
+                        raise ValueError(
+                            f"{self.__class__.__name__} does not support caching and therefore **can't** be used "
+                            "for contrastive search."
+                        )
+                    elif (
+                        not isinstance(past_key_values[0], (tuple, torch.Tensor))
+                        or past_key_values[0][0].shape[0] != batch_size
+                    ):
+                        raise ValueError(
+                            f"{self.__class__.__name__} does not have a standard cache format and therefore **can't** be "
+                            "used for contrastive search without further modifications."
+                        )
+    
+                # contrastive_search main logic start:
+                # contrastive search decoding consists of two steps: (1) candidate tokens recall; (2) candidate re-rank by
+                # degeneration penalty
+                processed_logit_for_next_step = logits_processor(input_ids, logit_for_next_step)
+                next_probs = nn.functional.softmax(processed_logit_for_next_step, dim=-1)
+    
+                top_k_probs, top_k_ids = torch.topk(next_probs, dim=-1, k=top_k)
+    
+                # Store scores, attentions and hidden_states when required
+                if return_dict_in_generate:
+                    if output_logits:
+                        raw_logits += (logit_for_next_step,)
+                    if output_scores:
+                        scores += (processed_logit_for_next_step,)
+                    if output_attentions:
+                        decoder_attentions += (
+                            (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
+                        )
+                        if self.config.is_encoder_decoder:
+                            cross_attentions += (outputs.cross_attentions,)
+    
+                    if output_hidden_states:
+                        decoder_hidden_states += (
+                            (outputs.decoder_hidden_states,)
+                            if self.config.is_encoder_decoder
+                            else (outputs.hidden_states,)
+                        )
+    
+                # This is needed to properly delete outputs.logits which may be very large for this first iteration
+                # Otherwise a reference to outputs.logits is kept all along until after the next call to self.forward()
+                del outputs
+    
+                if not sequential:
+                    # Replicates the new past_key_values to match the `top_k` candidates
+                    past = model_kwargs["past_key_values"]
+                    # If it is a static cache, modify it in-place layer after layer to save memory
+                    if isinstance(past, DynamicCache) or (
+                        isinstance(past, EncoderDecoderCache) and isinstance(past.self_attention_cache, DynamicCache)
+                    ):
+                        past.batch_repeat_interleave(top_k)
+                    else:
+                        new_key_values = []
+                        for layer in past:
+                            items = []
+                            # item is either the key or the value matrix
+                            for item in layer:
+                                items.append(item.repeat_interleave(top_k, dim=0))
+                            new_key_values.append(tuple(items))
+    
+                        past = tuple(new_key_values)
+    
+                    model_kwargs["past_key_values"] = past
+    
+                if sequential:
+                    all_outputs = []
+                    for i in range(top_k):
+                        # compute the candidate tokens by the language model and collect their hidden_states
+                        next_model_inputs = self.prepare_inputs_for_generation(top_k_ids[:, i].view(-1, 1), **model_kwargs)
+    
+                        outputs = self(
+                            **next_model_inputs,
+                            return_dict=True,
+                            output_hidden_states=True,
+                            output_attentions=output_attentions,
+                        )
+                        if isinstance(outputs["past_key_values"], DynamicCache) or (
+                            isinstance(outputs["past_key_values"], EncoderDecoderCache)
+                            and isinstance(outputs["past_key_values"].self_attention_cache, DynamicCache)
+                        ):
+                            # Remove past K-V from output since we don't need to stack later
+                            outputs["past_key_values"] = None
+                            # Remove last token from past K-V since we don't want to append it at this point
+                            model_kwargs["past_key_values"].crop(-1)
+    
+                        all_outputs.append(outputs)
+                    outputs = stack_model_outputs(all_outputs, self.config.get_text_config())
+    
+                else:
+                    # compute the candidate tokens by the language model and collect their hidden_states
+                    # assembles top_k_ids into batch of size k
+                    next_model_inputs = self.prepare_inputs_for_generation(top_k_ids.view(-1, 1), **model_kwargs)
+    
+                    outputs = self(
+                        **next_model_inputs,
+                        return_dict=True,
+                        output_hidden_states=True,
+                        output_attentions=output_attentions,
+                    )
+    
+                # This is essential to avoid having a last reference to the big past K-V and double the necessary memory
+                # in the next loop
+                del next_model_inputs
+    
+                # name is different for encoder-decoder and decoder-only models
+                if self.config.is_encoder_decoder:
+                    next_hidden = outputs.decoder_hidden_states[-1]
+                    full_hidden_states = outputs.decoder_hidden_states
+                else:
+                    next_hidden = outputs.hidden_states[-1]
+                    full_hidden_states = outputs.hidden_states
+    
+                # .float() is needed to retain precision for later logits manipulations
+                logits = outputs.logits[:, -1, :].float()
+                context_hidden = last_hidden_states.repeat_interleave(top_k, dim=0)
+    
+                # compute the degeneration penalty and re-rank the candidates based on the degeneration penalty and the
+                # model confidence. Keeping `selected_idx` on CPU enables multi-device contrastive search and doesn't
+                # introduce (noticeable) slowdowns on single-device runs.
+                selected_idx = _ranking_fast(
+                    context_hidden, next_hidden, top_k_probs, cosine_matrix_mask, penalty_alpha, top_k
+                )
+                cosine_matrix_mask = torch.cat(
+                    [cosine_matrix_mask, cosine_matrix_mask.new_ones((cosine_matrix_mask.shape[0], 1))], dim=-1
+                )
+                selected_idx = selected_idx.to("cpu")
+    
+                # This will be used instead of the previous inneficient torch.stack(torch.split())
+                augmented_idx = torch.tensor([x + i * top_k for i, x in enumerate(selected_idx)])
+    
+                # prepare for the next step: (1) next token_id; (2) past_key_values; (3) last_hidden_states for computing
+                # the degeneration penalty; (4) logits for selecting next top-k candidates; (5) selected tokens scores
+                # (model confidence minus degeneration penalty); (6) decoder hidden_states
+                next_tokens = top_k_ids[range(len(top_k_ids)), selected_idx]
+                next_hidden = torch.stack(torch.split(next_hidden.squeeze(dim=1), top_k))
+                next_hidden = next_hidden[range(batch_size), selected_idx, :]
+                last_hidden_states = torch.cat([last_hidden_states, next_hidden.unsqueeze(1)], dim=1)
+    
+                next_decoder_hidden_states = ()
+                for layer in full_hidden_states:
+                    layer = torch.stack(torch.split(layer, top_k))[range(batch_size), selected_idx, :]
+                    next_decoder_hidden_states += (layer,)
+    
+                # generate past_key_values cache of only the selected token
+                if sequential:
+                    next_model_input = self.prepare_inputs_for_generation(
+                        top_k_ids[:, selected_idx].view(-1, 1), **model_kwargs
+                    )
+    
+                    selected_outputs = self(
+                        **next_model_input,
+                        return_dict=True,
+                        output_hidden_states=False,
+                        output_attentions=False,
+                    )
+                    next_past_key_values = selected_outputs["past_key_values"]
+    
+                else:
+                    next_past_key_values = None
+                    for possible_cache_name in ALL_CACHE_NAMES:
+                        next_past_key_values = next_past_key_values or getattr(outputs, possible_cache_name, None)
+                    # Do it in-place layer per layer to save memory
+                    if isinstance(next_past_key_values, DynamicCache) or (
+                        isinstance(next_past_key_values, EncoderDecoderCache)
+                        and isinstance(next_past_key_values.self_attention_cache, DynamicCache)
+                    ):
+                        next_past_key_values.batch_select_indices(augmented_idx)
+                    else:
+                        new_key_values = []
+                        for layer in next_past_key_values:
+                            items = []
+                            # item is either the key or the value matrix
+                            for item in layer:
+                                items.append(item[augmented_idx, ...])
+                            new_key_values.append(tuple(items))
+    
+                        next_past_key_values = tuple(new_key_values)
+    
+                logit_for_next_step = torch.stack(torch.split(logits, top_k))[range(batch_size), selected_idx, :]
+                logit_for_next_step = logit_for_next_step.to(input_ids.device)
+    
+                # Rebuilds the relevant parts of the model output for the selected token, for use in the next iteration
+                if self.config.is_encoder_decoder:
+                    next_step_cross_attentions = ()
+                    next_step_decoder_attentions = ()
+                    if output_attentions:
+                        for layer in outputs.cross_attentions:
+                            layer = torch.stack(torch.split(layer, top_k, dim=0))[range(batch_size), selected_idx, ...]
+                            next_step_cross_attentions += (layer,)
+                        for layer in outputs.decoder_attentions:
+                            layer = torch.stack(torch.split(layer, top_k, dim=0))[range(batch_size), selected_idx, ...]
+                            next_step_decoder_attentions += (layer,)
+                    outputs = Seq2SeqLMOutput(
+                        past_key_values=next_past_key_values,
+                        decoder_hidden_states=next_decoder_hidden_states,
+                        decoder_attentions=next_step_decoder_attentions or None,
+                        cross_attentions=next_step_cross_attentions or None,
+                    )
+                else:
+                    next_step_attentions = ()
+                    if output_attentions:
+                        for layer in outputs.attentions:
+                            layer = torch.stack(torch.split(layer, top_k, dim=0))[range(batch_size), selected_idx, ...]
+                            next_step_attentions += (layer,)
+                    outputs = CausalLMOutputWithPast(
+                        past_key_values=next_past_key_values,
+                        hidden_states=next_decoder_hidden_states,
+                        attentions=next_step_attentions or None,
+                    )
+                # contrastive_search main logic end
+    
+                # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+                model_kwargs = self._update_model_kwargs_for_generation(
+                    outputs,
+                    model_kwargs,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                )
+                if synced_gpus and this_peer_finished:
+                    continue
+    
+                # finished sentences should have their next token be a padding token
+                if has_eos_stopping_criteria:
+                    next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+    
+                # update generated ids, model inputs, and length for next step
+                input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
+                if streamer is not None:
+                    streamer.put(next_tokens.cpu())
+    
+                # stop when each sentence is finished
+                unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores)
+                this_peer_finished = unfinished_sequences.max() == 0
+    
+            if streamer is not None:
+                streamer.end()
+    
+            if return_dict_in_generate:
+                # Contrastive search works by forward looking at the next token, so we need to exclude it from
+                # `past_key_values` to be consistent with the other decoding methods
+                if model_kwargs.get("past_key_values") is not None:
+                    if isinstance(model_kwargs["past_key_values"], DynamicCache) or (
+                        isinstance(model_kwargs["past_key_values"], EncoderDecoderCache)
+                        and isinstance(model_kwargs["past_key_values"].self_attention_cache, DynamicCache)
+                    ):
+                        model_kwargs["past_key_values"].crop(-1)
+                    else:
+                        past_key_values = []
+                        for layer in model_kwargs["past_key_values"]:
+                            layer_past_key_values = []
+                            for item in layer:
+                                layer_past_key_values.append(item[..., :-1, :])
+                            past_key_values.append(tuple(layer_past_key_values))
+                        model_kwargs["past_key_values"] = tuple(past_key_values)
+    
+                if self.config.is_encoder_decoder:
+                    return GenerateEncoderDecoderOutput(
+                        sequences=input_ids,
+                        scores=scores,
+                        logits=raw_logits,
+                        encoder_attentions=encoder_attentions,
+                        encoder_hidden_states=encoder_hidden_states,
+                        decoder_attentions=decoder_attentions,
+                        cross_attentions=cross_attentions,
+                        decoder_hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+                else:
+                    return GenerateDecoderOnlyOutput(
+                        sequences=input_ids,
+                        scores=scores,
+                        logits=raw_logits,
+                        attentions=decoder_attentions,
+                        hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+            else:
+                return input_ids
+    
+        def _sample(
+            self,
+            input_ids: torch.LongTensor,
+            logits_processor: LogitsProcessorList,
+            stopping_criteria: StoppingCriteriaList,
+            generation_config: GenerationConfig,
+            synced_gpus: bool,
+            streamer: Optional["BaseStreamer"],
+            **model_kwargs,
+        ) -> Union[GenerateNonBeamOutput, torch.LongTensor]:
+            r"""
+            Generates sequences of token ids for models with a language modeling head using **multinomial sampling** and
+            can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.
+    
+            Parameters:
+                input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                    The sequence used as a prompt for the generation.
+                logits_processor (`LogitsProcessorList`):
+                    An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
+                    used to modify the prediction scores of the language modeling head applied at each generation step.
+                stopping_criteria (`StoppingCriteriaList`):
+                    An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
+                    used to tell if the generation loop should stop.
+                generation_config ([`~generation.GenerationConfig`]):
+                    The generation configuration to be used as parametrization of the decoding method.
+                synced_gpus (`bool`):
+                    Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+                    `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+                streamer (`BaseStreamer`, *optional*):
+                    Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+                    through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+                model_kwargs:
+                    Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is
+                    an encoder-decoder model the kwargs should include `encoder_outputs`.
+    
+            Return:
+                [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or `torch.LongTensor`:
+                A `torch.LongTensor` containing the generated tokens (default behaviour) or a
+                [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
+                `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if
+                `model.config.is_encoder_decoder=True`.
+            """
+            # init values
+            pad_token_id = generation_config._pad_token_tensor
+            output_attentions = generation_config.output_attentions
+            output_hidden_states = generation_config.output_hidden_states
+            output_scores = generation_config.output_scores
+            output_logits = generation_config.output_logits
+            return_dict_in_generate = generation_config.return_dict_in_generate
+            has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria)
+            do_sample = generation_config.do_sample
+    
+            # init attention / hidden states / scores tuples
+            scores = () if (return_dict_in_generate and output_scores) else None
+            raw_logits = () if (return_dict_in_generate and output_logits) else None
+            decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+            cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+            decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+    
+            # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+            if return_dict_in_generate and self.config.is_encoder_decoder:
+                encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
+                encoder_hidden_states = (
+                    model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+                )
+    
+            # keep track of which sequences are already finished
+            batch_size, cur_len = input_ids.shape[:2]
+            this_peer_finished = False
+            unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+            model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+    
+            model_forward = self.__call__
+            compile_forward = self._valid_auto_compile_criteria(model_kwargs, generation_config)
+            if compile_forward:
+                os.environ["TOKENIZERS_PARALLELISM"] = "0"
+                # If we use FA2 and a static cache, we cannot compile with fullgraph
+                if self.config._attn_implementation == "flash_attention_2" and getattr(
+                    model_kwargs.get("past_key_values"), "is_compileable", False
+                ):
+                    if generation_config.compile_config is None:
+                        generation_config.compile_config = CompileConfig(fullgraph=False)
+                    # only raise warning if the user passed an explicit compile-config (otherwise, simply change the default without confusing the user)
+                    elif generation_config.compile_config.fullgraph:
+                        logger.warning_once(
+                            "When using Flash Attention 2 and a static cache, you cannot use the option `CompileConfig(fullgraph=True)` as "
+                            "FA2 introduces graph breaks. We overrode the option with `fullgraph=False`."
+                        )
+                        generation_config.compile_config.fullgraph = False
+                model_forward = self.get_compiled_call(generation_config.compile_config)
+    
+            if generation_config.prefill_chunk_size is not None:
+                model_kwargs = self._prefill_chunking(input_ids, generation_config, **model_kwargs)
+                is_prefill = False
+            else:
+                is_prefill = True
+    
+            while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+                # prepare model inputs
+                model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+    
+                # prepare variable output controls (note: some models won't accept all output controls)
+                model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
+                model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
+    
+                if is_prefill:
+                    outputs = self(**model_inputs, return_dict=True)
+                    is_prefill = False
+                else:
+                    outputs = model_forward(**model_inputs, return_dict=True)
+    
+                # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+                model_kwargs = self._update_model_kwargs_for_generation(
+                    outputs,
+                    model_kwargs,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                )
+                if synced_gpus and this_peer_finished:
+                    continue
+    
+                # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
+                # (the clone itself is always small)
+                next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device)
+    
+                # pre-process distribution
+                next_token_scores = logits_processor(input_ids, next_token_logits)
+    
+                # Store scores, attentions and hidden_states when required
+                if return_dict_in_generate:
+                    if output_scores:
+                        scores += (next_token_scores,)
+                    if output_logits:
+                        raw_logits += (next_token_logits,)
+                    if output_attentions:
+                        decoder_attentions += (
+                            (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
+                        )
+                        if self.config.is_encoder_decoder:
+                            cross_attentions += (outputs.cross_attentions,)
+    
+                    if output_hidden_states:
+                        decoder_hidden_states += (
+                            (outputs.decoder_hidden_states,)
+                            if self.config.is_encoder_decoder
+                            else (outputs.hidden_states,)
+                        )
+    
+                # token selection
+                if do_sample:
+                    probs = nn.functional.softmax(next_token_scores, dim=-1)
+                    # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution
+                    next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+                else:
+                    next_tokens = torch.argmax(next_token_scores, dim=-1)
+    
+                # finished sentences should have their next token be a padding token
+                if has_eos_stopping_criteria:
+                    next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+    
+                # update generated ids, model inputs, and length for next step
+                input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
+                if streamer is not None:
+                    streamer.put(next_tokens.cpu())
+    
+                unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores)
+                this_peer_finished = unfinished_sequences.max() == 0
+                cur_len += 1
+    
+                # This is needed to properly delete outputs.logits which may be very large for first iteration
+                # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration
+                del outputs
+    
+            if streamer is not None:
+                streamer.end()
+    
+            if return_dict_in_generate:
+                if self.config.is_encoder_decoder:
+                    return GenerateEncoderDecoderOutput(
+                        sequences=input_ids,
+                        scores=scores,
+                        logits=raw_logits,
+                        encoder_attentions=encoder_attentions,
+                        encoder_hidden_states=encoder_hidden_states,
+                        decoder_attentions=decoder_attentions,
+                        cross_attentions=cross_attentions,
+                        decoder_hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+                else:
+                    return GenerateDecoderOnlyOutput(
+                        sequences=input_ids,
+                        scores=scores,
+                        logits=raw_logits,
+                        attentions=decoder_attentions,
+                        hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+            else:
+                return input_ids
+    
+        # Auxiliary functions for beam search
+        def _temporary_reorder_cache(self, past_key_values, beam_idx):
+            """
+            Temporary function to handle the different types of cache reordering processes while we roll out `Cache`.
+    
+            TODO: standardize cache formats and make all models compatible with `Cache`. It would remove the need
+            for this function, with `Cache.reorder_cache` being the sole remaining code path
+            """
+            model_class = self.__class__.__name__.lower()
+            # Exception 1: code path for models using the legacy cache format
+            if isinstance(past_key_values, (tuple, list)):
+                past_key_values = self._reorder_cache(past_key_values, beam_idx)
+            # Exception 2: models with different cache formats. These are limited to `DynamicCache` until their
+            # cache format is standardized, to avoid adding complexity to the codebase.
+            elif "gptbigcode" in model_class:
+                if not isinstance(past_key_values, (DynamicCache, EncoderDecoderCache)):
+                    raise ValueError(
+                        f"Using an unsupported cache format with {model_class}. Currently, it only supports the "
+                        "legacy tuple format or `DynamicCache`"
+                    )
+                past_key_values = self._reorder_cache(past_key_values, beam_idx)
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+            # Standard code path: use the `Cache.reorder_cache`
+            else:
+                past_key_values.reorder_cache(beam_idx)
+            return past_key_values
+    
+        @staticmethod
+        def _flatten_beam_dim(tensor: torch.Tensor) -> torch.Tensor:
+            """[batch_size, num_beams, ...] -> [batch_size * num_beams, ...]"""
+            shape = list(tensor.shape)
+            return torch.reshape(tensor, [shape[0] * shape[1]] + shape[2:])
+    
+        @staticmethod
+        def _unflatten_beam_dim(tensor: torch.Tensor, batch_size: int, num_beams: int) -> torch.Tensor:
+            """[batch_size * num_beams, ...] -> [batch_size, num_beams, ...]"""
+            shape = list(tensor.shape)
+            return torch.reshape(tensor, [batch_size, num_beams] + shape[1:])
+    
+        @staticmethod
+        def _gather_beams(tensor: torch.Tensor, beam_indices: torch.Tensor) -> torch.Tensor:
+            """
+            Gathers the beam slices indexed by beam_indices into new beam array.
+    
+            Args:
+                tensor (`torch.Tensor`): A tensor containing data to be gathered. The tensor is a 2D or a 3D tensor
+                    with the two first dimensions depicting the batch and the beam dimensions.
+                beam_indices (`torch.Tensor` of shape `(batch_size, num_beams_to_select)`): The indices of the beams to
+                    select .
+    
+            Returns:
+                A tensor with the selected beams
+            """
+            # `take_along_dim` requires its indices arg to have the same number of dims as `input`
+            while len(beam_indices.shape) < len(tensor.shape):
+                beam_indices = beam_indices.unsqueeze(-1)
+            gathered_tensor = torch.take_along_dim(input=tensor, indices=beam_indices, dim=1)
+            return gathered_tensor
+    
+        @staticmethod
+        def _beam_search_has_unfinished_sequences(
+            running_beam_scores: torch.Tensor,
+            beam_scores: torch.Tensor,
+            is_sent_finished: torch.Tensor,
+            next_token_hits_stopping_criteria: torch.Tensor,
+            cur_len: int,
+            max_length: int,
+            decoder_prompt_len: int,
+            early_stopping: Union[bool, str],
+            length_penalty: float,
+        ):
+            """
+            Beam Search stopping condition -- halts the generation loop if any of these conditions becomes False
+            """
+            # a. Can the open beams improve the top completed scores?
+            # early_stopping == False -> apply heuristic = always get the best score from
+            #   `cur_len - decoder_prompt_len`. See the discussion below for more details.
+            #   https://github.com/huggingface/transformers/pull/20901#issuecomment-1369845565
+            # early_stopping == "never" -> compute the best score from `max_length` or `cur_len`, depending on the
+            #   sign of `length_penalty`. Positive `length_penalty` favors longer sequences, thus we use
+            #   `max_length` there.
+            if early_stopping == "never" and length_penalty > 0.0:
+                best_hypothetical_length = max_length - decoder_prompt_len
+            else:
+                best_hypothetical_length = cur_len - decoder_prompt_len
+            best_possible_running_score = running_beam_scores[:, :1] / (best_hypothetical_length**length_penalty)
+            worst_finished_score = torch.where(is_sent_finished, torch.min(beam_scores, dim=1, keepdim=True)[0], -1.0e9)
+            improvement_possible = torch.any(best_possible_running_score > worst_finished_score)
+    
+            # b. Is there still a beam without fully completed sequences? This is only relevant if early_stopping is
+            # enabled, where we want to finish as soon as all beams have a completed sequence.
+            exists_open_beam = ~(torch.all(is_sent_finished) & (early_stopping is True))
+    
+            # c. Have we hit a stopping criteria with all running sequences and have no way to continue? e.g. we have
+            # reached `max_length``
+            valid_continuations = ~torch.all(next_token_hits_stopping_criteria)
+    
+            return improvement_possible & exists_open_beam & valid_continuations
+    
+        def _get_top_k_continuations(
+            self,
+            accumulated_log_probs: torch.Tensor,
+            running_sequences: torch.Tensor,
+            running_beam_indices: torch.Tensor,
+            cur_len: int,
+            decoder_prompt_len: int,
+            do_sample: bool,
+            beams_to_keep: int,
+            num_beams: int,
+            vocab_size: int,
+            batch_size: int,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            """
+            Get top-K continuations given the accumulated log probs on the next token.
+    
+            A few notes to understand what's going on:
+            1. Each item in batch has `num_beams` * `vocab_size` candidate continuations. For each item, get the
+            top K [K = (number of EOS tokens + 1) * `num_beams`] candidates with the highest accumulated
+            log-probabilities, or sample them without replacement using the accumulated scores
+            2. We gather the top K (as opposed to `num_beams`, or any number lower than K) here so that we have at
+            least `num_beams` sequences remaining to continue the live beam search.
+            3. Note that other stopping criteria might result in impossible to continue beams, i.e. all continuations
+            selected in this step hit the stopping criteria.
+            """
+            # TODO (joao): This function should take an optional beam scorer function, to manipulate the scores after
+            # token selection. The function should be an argument exposed, so that custom scoring functions can be
+            # defined.
+    
+            # Gather the top K scores from _all_ beams.
+            if do_sample:
+                topk_indices = torch.multinomial(
+                    nn.functional.softmax(accumulated_log_probs, dim=-1), num_samples=beams_to_keep
+                )
+                topk_log_probs = torch.gather(input=accumulated_log_probs, dim=1, index=topk_indices)
+            else:
+                topk_log_probs, topk_indices = torch.topk(accumulated_log_probs, k=beams_to_keep)
+    
+            # Gather K top beams, recover the beam index by floor division and token id by modulo division
+            topk_current_beam_indices = topk_indices // vocab_size
+            topk_running_beam_indices = self._gather_beams(running_beam_indices, topk_current_beam_indices)
+            topk_running_sequences = self._gather_beams(running_sequences, topk_current_beam_indices)
+            topk_ids = topk_indices % vocab_size
+    
+            # Update sequences for the K top-k new sequences.
+            topk_running_sequences[:, :, cur_len] = topk_ids
+    
+            # we want to store the beam indices with batch information -> real beam index = beam index % num beams
+            batch_offset = torch.arange(batch_size, device=topk_ids.device).view(-1, 1) * num_beams
+            batch_modified_indices = topk_current_beam_indices + batch_offset
+            topk_running_beam_indices[:, :, cur_len - decoder_prompt_len] = batch_modified_indices
+    
+            return topk_log_probs, topk_running_sequences, topk_running_beam_indices
+    
+        def _get_running_beams_for_next_iteration(
+            self,
+            topk_log_probs: torch.Tensor,
+            topk_running_sequences: torch.Tensor,
+            topk_running_beam_indices: torch.Tensor,
+            next_token_hits_stopping_criteria: torch.Tensor,
+            num_beams: int,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            """
+            Given the top-K continuations, their scores, and whether they hit a stopping criteria, select the
+            best non-finished beams to continue beam search in the next iteration.
+            """
+            # To prevent these just finished sequences from being used in subsequent iterations, set their log probs
+            # to a very large negative value
+            topk_running_log_probs = topk_log_probs + next_token_hits_stopping_criteria.to(torch.float32) * -1.0e9
+    
+            next_topk_indices = torch.topk(topk_running_log_probs, k=num_beams)[1]
+            running_sequences = self._gather_beams(topk_running_sequences, next_topk_indices)
+            running_beam_scores = self._gather_beams(topk_running_log_probs, next_topk_indices)
+            running_beam_indices = self._gather_beams(topk_running_beam_indices, next_topk_indices)
+            return running_sequences, running_beam_scores, running_beam_indices
+    
+        def _update_finished_beams(
+            self,
+            sequences: torch.Tensor,
+            topk_running_sequences: torch.Tensor,
+            beam_scores: torch.Tensor,
+            topk_log_probs: torch.Tensor,
+            beam_indices: torch.Tensor,
+            topk_running_beam_indices: torch.Tensor,
+            is_sent_finished: torch.Tensor,
+            next_token_hits_stopping_criteria: torch.Tensor,
+            top_num_beam_mask: torch.Tensor,
+            num_beams: int,
+            cur_len: int,
+            decoder_prompt_len: int,
+            length_penalty: float,
+            early_stopping: Union[bool, str],
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+            """
+            Updates the finished beams if (and only if) there are new completed sequences that have a higher score than
+            the current finished sequences.
+            """
+            # Only the top `num_beam` sequences can be considered for the final returned sequences. Remember: the
+            # remaining sequences only exist as a backup to ensure that we have at least `num_beams` sequences to
+            # continue.
+            did_top_num_beams_just_finished = next_token_hits_stopping_criteria & top_num_beam_mask[None, :]
+    
+            # Further process topk logits for the finished beams
+            # - add length penalty
+            topk_log_probs = topk_log_probs / ((cur_len + 1 - decoder_prompt_len) ** length_penalty)
+            # - make sure no scores can be added anymore if beam is full and early stopping is on
+            beams_in_batch_are_full = torch.all(is_sent_finished, axis=-1, keepdims=True) & (early_stopping is True)
+            topk_log_probs += beams_in_batch_are_full.to(torch.float32) * -1.0e9
+            # - make sure still running sequences cannot be chosen as finalized beam
+            topk_log_probs += (~did_top_num_beams_just_finished) * -1.0e9
+    
+            # Get finalized  `num_beam` sequences for the next generation step -- combine the previous finalized
+            # data with the new finalized sequences (if any, non-finalized sequences have a very large negative score
+            # in this step), and keep the best `num_beams` sequences.
+            merged_sequences = torch.cat((sequences, topk_running_sequences), dim=1)
+            merged_scores = torch.cat((beam_scores, topk_log_probs), dim=1)
+            merged_beam_indices = torch.cat((beam_indices, topk_running_beam_indices), dim=1)
+            merged_is_sent_finished = torch.cat((is_sent_finished, did_top_num_beams_just_finished), dim=1)
+            topk_merged_indices = torch.topk(merged_scores, k=num_beams)[1]
+            sequences = self._gather_beams(merged_sequences, topk_merged_indices)
+            beam_scores = self._gather_beams(merged_scores, topk_merged_indices)
+            beam_indices = self._gather_beams(merged_beam_indices, topk_merged_indices)
+            is_sent_finished = self._gather_beams(merged_is_sent_finished, topk_merged_indices)
+            return sequences, beam_scores, beam_indices, is_sent_finished
+    
+        # end of auxiliary functions for beam search
+    
+        def _beam_search(
+            self,
+            input_ids: torch.LongTensor,
+            logits_processor: LogitsProcessorList,
+            stopping_criteria: StoppingCriteriaList,
+            generation_config: GenerationConfig,
+            synced_gpus: bool,
+            **model_kwargs,
+        ) -> Union[GenerateBeamOutput, torch.LongTensor]:
+            r"""
+            Generates sequences of token ids for models with a language modeling head using **beam search decoding** and
+            can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.
+    
+            If it's the first time you're diving into Beam Search, we recommend you read the following blog post:
+            https://huggingface.co/blog/how-to-generate (especially the beam search section).
+    
+            You can recompute the sequence scores from the individual scores using the `compute_transition_scores` function
+            (https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores)
+    
+            Parameters:
+                input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`):
+                    The sequence used as a prompt for the generation.
+                logits_processor (`LogitsProcessorList`):
+                    An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
+                    used to modify the prediction scores of the language modeling head applied at each generation step.
+                stopping_criteria (`StoppingCriteriaList`:
+                    An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
+                    used to tell if the generation loop should stop.
+                generation_config ([`~generation.GenerationConfig`]):
+                    The generation configuration to be used as parametrization of the decoding method.
+                synced_gpus (`bool`):
+                    Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+                    `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+                model_kwargs:
+                    Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is
+                    an encoder-decoder model the kwargs should include `encoder_outputs`.
+    
+            Return:
+                [`generation.GenerateBeamDecoderOnlyOutput`], [`~generation.GenerateBeamEncoderDecoderOutput`] or
+                `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a
+                [`~generation.GenerateBeamDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
+                `return_dict_in_generate=True` or a [`~generation.GenerateBeamEncoderDecoderOutput`] if
+                `model.config.is_encoder_decoder=True`.
+            """
+    
+            # 1. init beam_search values
+            pad_token_id = generation_config._pad_token_tensor
+            eos_token_id = generation_config._eos_token_tensor
+            output_attentions = generation_config.output_attentions
+            output_hidden_states = generation_config.output_hidden_states
+            output_scores = generation_config.output_scores
+            output_logits = generation_config.output_logits
+            return_dict_in_generate = generation_config.return_dict_in_generate
+            do_sample = generation_config.do_sample
+            early_stopping = generation_config.early_stopping
+            length_penalty = generation_config.length_penalty
+            max_length = generation_config.max_length
+            num_beams = generation_config.num_beams
+            num_return_sequences = generation_config.num_return_sequences
+    
+            batch_size_unflattened, cur_len = input_ids.shape[:2]
+            batch_size = batch_size_unflattened // num_beams
+            # TODO (joao): standardize special cases
+            if self.__class__.__name__ == "MoshiDepthDecoder":
+                vocab_size = self.config.audio_vocab_size
+            elif self.__class__.__name__ == "ImageGPTForCausalImageModeling":
+                vocab_size = self.get_output_embeddings().out_features
+            else:
+                vocab_size = self.config.get_text_config().vocab_size
+            decoder_prompt_len = cur_len
+            this_peer_finished = False
+    
+            # At each beam search step, we want to keep top K [K = (number of EOS tokens + 1) * `num_beams`] candidates
+            # with the highest log-probabilities, or sample K continuations without replacement. We gather the top K
+            # (as opposed to `num_beams`, or any number lower than K) so that we have at least `num_beams` sequences
+            # non-finished to continue the live beam search, in case the top `num_beams` all select an EOS token.
+            n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0
+            beams_to_keep = max(2, 1 + n_eos_tokens) * num_beams
+            top_num_beam_mask = torch.cat(
+                (torch.ones((num_beams), dtype=torch.bool), torch.zeros((beams_to_keep - num_beams), dtype=torch.bool)),
+                dim=0,
+            ).to(input_ids.device)
+    
+            model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+    
+            # (joao) feature lost in the refactor. Probably won't implement, hurts readability with minimal gains (there
+            # are newer low-memory alternatives like the offloaded cache)
+            sequential = generation_config.low_memory
+            if sequential:
+                raise ValueError(
+                    "`low_memory=True` is not supported after the beam search refactor. Please check the discussion in "
+                    "#35802 *after the PR got merged*, and add a comment there if your questions are not yet answered."
+                )
+    
+            # 2. init output tuples
+            all_scores = () if (return_dict_in_generate and output_scores) else None
+            raw_logits = () if (return_dict_in_generate and output_logits) else None
+            beam_indices = () if (return_dict_in_generate and output_logits) else None
+            decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+            cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+            decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+    
+            # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+            if return_dict_in_generate and self.config.is_encoder_decoder:
+                encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
+                encoder_hidden_states = (
+                    model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+                )
+    
+            # 3. init running tensors and static-shaped placeholders
+    
+            # per batch, beam-item holding current token in loop and completed sequences
+            output_fill_value = pad_token_id or eos_token_id[0] if eos_token_id is not None else -1
+            running_sequences = torch.full(
+                (batch_size, num_beams, max_length),
+                fill_value=output_fill_value,
+                dtype=torch.int64,
+                device=input_ids.device,
+            )
+            running_sequences[:, :, :cur_len] = self._unflatten_beam_dim(input_ids, batch_size, num_beams)
+            sequences = running_sequences.detach().clone()
+    
+            # per batch, beam-item score, logprobs
+            # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens
+            # of the first beam are considered to avoid sampling the exact same tokens across all beams.
+            running_beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
+            running_beam_scores[:, 1:] = -1e9
+            beam_scores = torch.full((batch_size, num_beams), fill_value=-1e9, dtype=torch.float, device=input_ids.device)
+    
+            # per batch, beam-item state bit indicating if sentence has finished.
+            is_sent_finished = torch.zeros((batch_size, num_beams), dtype=torch.bool, device=input_ids.device)
+    
+            # per batch, beam-item state bit indicating if there are valid continuations.
+            next_token_hits_stopping_criteria = torch.zeros(
+                (batch_size, num_beams), dtype=torch.bool, device=input_ids.device
+            )
+    
+            # per batch selected beam indices
+            running_beam_indices = torch.full(
+                (batch_size, num_beams, max_length - cur_len), fill_value=-1, dtype=torch.int32, device=input_ids.device
+            )
+            beam_indices = running_beam_indices.detach().clone()
+    
+            # 4. run the generation loop
+            while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+                # a. Forward current tokens, obtain the logits
+                flat_running_sequences = self._flatten_beam_dim(running_sequences[:, :, :cur_len])
+                model_inputs = self.prepare_inputs_for_generation(flat_running_sequences, **model_kwargs)
+    
+                # prepare variable output controls (note: some models won't accept all output controls)
+                model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
+                model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
+    
+                model_outputs = self(**model_inputs, return_dict=True)
+    
+                # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+                model_kwargs = self._update_model_kwargs_for_generation(
+                    model_outputs,
+                    model_kwargs,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                )
+                if synced_gpus and this_peer_finished:
+                    continue
+    
+                # Copy is needed to avoid keeping a hanging ref
+                logits = model_outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device)
+    
+                # b. Compute log probs -- get log probabilities from logits, process logits with processors (*e.g.*
+                # `temperature`, ...), and add new logprobs to existing running logprobs scores.
+                log_probs = nn.functional.log_softmax(logits, dim=-1)
+                log_probs = logits_processor(flat_running_sequences, log_probs)
+    
+                # Store logits, attentions and hidden_states when required
+                if return_dict_in_generate:
+                    if output_logits:
+                        raw_logits += (logits.clone(),)
+                    if return_dict_in_generate and output_scores:
+                        all_scores += (log_probs.clone(),)
+    
+                    if output_attentions:
+                        decoder_attentions += (
+                            (model_outputs.decoder_attentions,)
+                            if self.config.is_encoder_decoder
+                            else (model_outputs.attentions,)
+                        )
+                        if self.config.is_encoder_decoder:
+                            cross_attentions += (model_outputs.cross_attentions,)
+    
+                    if output_hidden_states:
+                        decoder_hidden_states += (
+                            (model_outputs.decoder_hidden_states,)
+                            if self.config.is_encoder_decoder
+                            else (model_outputs.hidden_states,)
+                        )
+    
+                # This is needed to properly delete logits which may be very large for first iteration
+                # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration
+                del model_outputs
+    
+                log_probs = self._unflatten_beam_dim(log_probs, batch_size, num_beams)
+                log_probs = log_probs + running_beam_scores[:, :, None]
+                log_probs = torch.reshape(log_probs, (batch_size, num_beams * vocab_size))
+    
+                # c. Retrieve top-K continuations, i.e. select the next token (greedy or sampling) and then keep the best
+                # continuations among all beams based on the accumulated scores.
+                topk_log_probs, topk_running_sequences, topk_running_beam_indices = self._get_top_k_continuations(
+                    accumulated_log_probs=log_probs,
+                    running_sequences=running_sequences,
+                    running_beam_indices=running_beam_indices,
+                    cur_len=cur_len,
+                    decoder_prompt_len=decoder_prompt_len,
+                    do_sample=do_sample,
+                    beams_to_keep=beams_to_keep,
+                    num_beams=num_beams,
+                    vocab_size=vocab_size,
+                    batch_size=batch_size,
+                )
+    
+                # d. Check which running sequences have finished
+                next_token_hits_stopping_criteria = stopping_criteria(
+                    self._flatten_beam_dim(topk_running_sequences[:, :, : cur_len + 1]),  # remove unfilled token indexes
+                    all_scores,
+                )
+                next_token_hits_stopping_criteria = self._unflatten_beam_dim(
+                    next_token_hits_stopping_criteria, batch_size, beams_to_keep
+                )
+    
+                # e. Get the non-finished running `num_beams` sequences for the next generation step
+                running_sequences, running_beam_scores, running_beam_indices = self._get_running_beams_for_next_iteration(
+                    topk_log_probs=topk_log_probs,
+                    topk_running_sequences=topk_running_sequences,
+                    topk_running_beam_indices=topk_running_beam_indices,
+                    next_token_hits_stopping_criteria=next_token_hits_stopping_criteria,
+                    num_beams=num_beams,
+                )
+    
+                # f. Update the completed beams if a new high score in a finished sequence is found
+                sequences, beam_scores, beam_indices, is_sent_finished = self._update_finished_beams(
+                    sequences=sequences,
+                    topk_running_sequences=topk_running_sequences,
+                    beam_scores=beam_scores,
+                    topk_log_probs=topk_log_probs,
+                    beam_indices=beam_indices,
+                    topk_running_beam_indices=topk_running_beam_indices,
+                    is_sent_finished=is_sent_finished,
+                    next_token_hits_stopping_criteria=next_token_hits_stopping_criteria,
+                    top_num_beam_mask=top_num_beam_mask,
+                    num_beams=num_beams,
+                    cur_len=cur_len,
+                    decoder_prompt_len=decoder_prompt_len,
+                    length_penalty=length_penalty,
+                    early_stopping=early_stopping,
+                )
+    
+                # g. Prepare remaining data for the next iteration, including computing the stopping condition for
+                # beam search as a whole (as opposed to individual beams, i.e. `stopping_criteria`)
+    
+                # pluck the cache from the beam indices that will be used in the next iteration
+                if model_kwargs.get("past_key_values", None) is not None:
+                    model_kwargs["past_key_values"] = self._temporary_reorder_cache(
+                        past_key_values=model_kwargs["past_key_values"],
+                        beam_idx=self._flatten_beam_dim(running_beam_indices[..., cur_len - decoder_prompt_len]),
+                    )
+    
+                cur_len = cur_len + 1
+                this_peer_finished = not self._beam_search_has_unfinished_sequences(
+                    running_beam_scores,
+                    beam_scores,
+                    is_sent_finished,
+                    next_token_hits_stopping_criteria,
+                    cur_len,
+                    max_length,
+                    decoder_prompt_len,
+                    early_stopping,
+                    length_penalty,
+                )
+    
+            # 5. prepare outputs
+            # Take best beams for each batch (the score is sorted in descending order)
+            sequences = self._flatten_beam_dim(sequences[:, :num_return_sequences, :])
+            beam_scores = self._flatten_beam_dim(beam_scores[:, :num_return_sequences])
+            beam_indices = self._flatten_beam_dim(beam_indices[:, :num_return_sequences, :])
+    
+            # Crop the static-shaped tensors to the actual size.
+            # `beam_indices` is initialized with -1s, and is updated with the beam index of the generated token at each
+            # step. We can use it to detect the generated length, which may be != `cur_len`  (e.g. selected beam is from a
+            # previous decoding iteration)
+            max_generated_length = ((beam_indices + 1).bool()).sum(dim=1).max()
+            output_length = decoder_prompt_len + max_generated_length
+            sequences = sequences[:, :output_length]
+            beam_indices = beam_indices[:, :max_generated_length]
+    
+            if return_dict_in_generate:
+                if not output_scores:
+                    beam_scores = None
+    
+                if self.config.is_encoder_decoder:
+                    return GenerateBeamEncoderDecoderOutput(
+                        sequences=sequences,
+                        sequences_scores=beam_scores,
+                        scores=all_scores,
+                        logits=raw_logits,
+                        beam_indices=beam_indices,
+                        encoder_attentions=encoder_attentions,
+                        encoder_hidden_states=encoder_hidden_states,
+                        decoder_attentions=decoder_attentions,
+                        cross_attentions=cross_attentions,
+                        decoder_hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+                else:
+                    return GenerateBeamDecoderOnlyOutput(
+                        sequences=sequences,
+                        sequences_scores=beam_scores,
+                        scores=all_scores,
+                        logits=raw_logits,
+                        beam_indices=beam_indices,
+                        attentions=decoder_attentions,
+                        hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+            else:
+                return sequences
+    
+        def _group_beam_search(
+            self,
+            input_ids: torch.LongTensor,
+            beam_scorer: BeamScorer,
+            logits_processor: LogitsProcessorList,
+            stopping_criteria: StoppingCriteriaList,
+            generation_config: GenerationConfig,
+            synced_gpus: bool,
+            **model_kwargs,
+        ):
+            r"""
+            Generates sequences of token ids for models with a language modeling head using **diverse beam search
+            decoding** and can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.
+    
+            Parameters:
+                input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`):
+                    The sequence used as a prompt for the generation.
+                beam_scorer (`BeamScorer`):
+                    An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and
+                    sorted during generation. For more information, the documentation of [`BeamScorer`] should be read.
+                logits_processor (`LogitsProcessorList`):
+                    An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
+                    used to modify the prediction scores of the language modeling head applied at each generation step.
+                stopping_criteria (`StoppingCriteriaList`):
+                    An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
+                    used to tell if the generation loop should stop.
+                generation_config ([`~generation.GenerationConfig`]):
+                    The generation configuration to be used as parametrization of the decoding method.
+                synced_gpus (`bool`):
+                    Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+                    `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+                model_kwargs:
+                    Additional model specific kwargs that will be forwarded to the `forward` function of the model. If
+                    model is an encoder-decoder model the kwargs should include `encoder_outputs`.
+    
+            Return:
+                [`~generation.GenerateBeamDecoderOnlyOutput`], [`~generation.GenerateBeamEncoderDecoderOutput`] or
+                `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a
+                [`~generation.GenerateBeamDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
+                `return_dict_in_generate=True` or a [`~generation.GenerateBeamEncoderDecoderOutput`] if
+                `model.config.is_encoder_decoder=True`.
+            """
+            # init values
+            pad_token_id = generation_config._pad_token_tensor
+            eos_token_id = generation_config._eos_token_tensor
+            output_attentions = generation_config.output_attentions
+            output_hidden_states = generation_config.output_hidden_states
+            output_scores = generation_config.output_scores
+            output_logits = generation_config.output_logits
+            return_dict_in_generate = generation_config.return_dict_in_generate
+    
+            num_beams = beam_scorer.num_beams
+            num_beam_groups = beam_scorer.num_beam_groups
+            num_sub_beams = num_beams // num_beam_groups
+            batch_size = len(beam_scorer._beam_hyps) // num_beam_groups
+            device = input_ids.device
+    
+            batch_beam_size, cur_len = input_ids.shape
+            model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+    
+            if return_dict_in_generate and output_scores:
+                beam_indices = [tuple(() for _ in range(num_sub_beams * batch_size)) for _ in range(num_beam_groups)]
+            else:
+                beam_indices = None
+    
+            if num_beams * batch_size != batch_beam_size:
+                raise ValueError(
+                    f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
+                )
+    
+            # init attention / hidden states / scores tuples
+            scores = () if (return_dict_in_generate and output_scores) else None
+            raw_logits = () if (return_dict_in_generate and output_logits) else None
+            decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+            cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+            decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+    
+            # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+            if return_dict_in_generate and self.config.is_encoder_decoder:
+                encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
+                encoder_hidden_states = (
+                    model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+                )
+    
+            # initialise score of first beam of each group with 0 and the rest with -1e9. This ensures that the beams in
+            # the same group don't produce same tokens every time.
+            beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device)
+            beam_scores[:, ::num_sub_beams] = 0
+            beam_scores = beam_scores.view((batch_size * num_beams,))
+    
+            this_peer_finished = False
+    
+            decoder_prompt_len = input_ids.shape[1]  # record the prompt length of decoder
+            while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+                # predicted tokens in cur_len step
+                current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device)
+    
+                # indices which will form the beams in the next time step
+                reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device)
+    
+                # do one decoder step on all beams of all sentences in batch
+                model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+    
+                # prepare variable output controls (note: some models won't accept all output controls)
+                model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
+                model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
+    
+                outputs = self(**model_inputs, return_dict=True)
+    
+                # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+                model_kwargs = self._update_model_kwargs_for_generation(
+                    outputs,
+                    model_kwargs,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                )
+                if synced_gpus and this_peer_finished:
+                    cur_len = cur_len + 1
+                    continue
+    
+                if output_scores:
+                    processed_score = torch.zeros_like(outputs.logits[:, -1, :])
+                if output_logits:
+                    # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
+                    # (the clone itself is always small)
+                    raw_logit_score = outputs.logits[:, -1, :].to(copy=True, device=input_ids.device)
+    
+                for beam_group_idx in range(num_beam_groups):
+                    group_start_idx = beam_group_idx * num_sub_beams
+                    group_end_idx = min(group_start_idx + num_sub_beams, num_beams)
+                    group_size = group_end_idx - group_start_idx
+    
+                    # indices of beams of current group among all sentences in batch
+                    batch_group_indices = []
+    
+                    for batch_idx in range(batch_size):
+                        batch_group_indices.extend(
+                            [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)]
+                        )
+                    group_input_ids = input_ids[batch_group_indices]
+    
+                    # select outputs of beams of current group only
+                    # No need to clone() the logits here as they will not retain outputs.logits at the end of the loop
+                    # .float() is needed to retain precision for later logits manipulations
+                    next_token_logits = outputs.logits[batch_group_indices, -1, :].to(
+                        dtype=torch.float32, device=input_ids.device
+                    )
+    
+                    next_token_scores = nn.functional.log_softmax(
+                        next_token_logits, dim=-1
+                    )  # (batch_size * group_size, vocab_size)
+                    vocab_size = next_token_scores.shape[-1]
+    
+                    next_token_scores_processed = logits_processor(
+                        group_input_ids, next_token_scores, current_tokens=current_tokens, beam_group_idx=beam_group_idx
+                    )
+                    next_token_scores = next_token_scores_processed + beam_scores[batch_group_indices].unsqueeze(-1)
+                    next_token_scores = next_token_scores.expand_as(next_token_scores_processed)
+    
+                    if output_scores:
+                        processed_score[batch_group_indices] = next_token_scores_processed
+    
+                    # reshape for beam search
+                    next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size)
+    
+                    # Sample 1 + len(eos_token_id) next tokens for each beam so we have at least 1 non eos token per beam.
+                    n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0
+                    next_token_scores, next_tokens = torch.topk(
+                        next_token_scores, max(2, 1 + n_eos_tokens) * group_size, dim=1, largest=True, sorted=True
+                    )
+    
+                    next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
+                    next_tokens = next_tokens % vocab_size
+    
+                    # stateless
+                    process_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+                    beam_outputs = beam_scorer.process(
+                        group_input_ids,
+                        next_token_scores,
+                        next_tokens,
+                        next_indices,
+                        pad_token_id=pad_token_id,
+                        eos_token_id=eos_token_id,
+                        beam_indices=process_beam_indices,
+                        group_index=beam_group_idx,
+                        decoder_prompt_len=decoder_prompt_len,
+                    )
+                    beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"]
+                    beam_next_tokens = beam_outputs["next_beam_tokens"]
+                    beam_idx = beam_outputs["next_beam_indices"]
+    
+                    if return_dict_in_generate and output_scores:
+                        beam_indices[beam_group_idx] = tuple(
+                            beam_indices[beam_group_idx][beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices[0]))
+                        )
+    
+                    input_ids[batch_group_indices] = group_input_ids[beam_idx]
+                    group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
+                    current_tokens[batch_group_indices] = group_input_ids[:, -1]
+    
+                    # (beam_idx // group_size) -> batch_idx
+                    # (beam_idx % group_size) -> offset of idx inside the group
+                    reordering_indices[batch_group_indices] = (
+                        num_beams * torch.div(beam_idx, group_size, rounding_mode="floor")
+                        + group_start_idx
+                        + (beam_idx % group_size)
+                    )
+    
+                # Store scores, attentions and hidden_states when required
+                if return_dict_in_generate:
+                    if output_scores:
+                        scores += (processed_score,)
+                    if output_logits:
+                        raw_logits += (raw_logit_score,)
+                    if output_attentions:
+                        decoder_attentions += (
+                            (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
+                        )
+                        if self.config.is_encoder_decoder:
+                            cross_attentions += (outputs.cross_attentions,)
+    
+                    if output_hidden_states:
+                        decoder_hidden_states += (
+                            (outputs.decoder_hidden_states,)
+                            if self.config.is_encoder_decoder
+                            else (outputs.hidden_states,)
+                        )
+    
+                input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1)
+    
+                # This is needed to properly delete outputs.logits which may be very large for first iteration
+                # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration
+                # IMPORTANT: Note that this should appear BEFORE the call to _reorder_cache() to save the maximum memory
+                # (that way the memory peak does not include outputs.logits)
+                del outputs
+    
+                if model_kwargs.get("past_key_values", None) is not None:
+                    model_kwargs["past_key_values"] = self._temporary_reorder_cache(
+                        model_kwargs["past_key_values"], reordering_indices
+                    )
+    
+                # increase cur_len
+                cur_len = cur_len + 1
+    
+                if beam_scorer.is_done or all(stopping_criteria(input_ids, scores)):
+                    this_peer_finished = True
+    
+            final_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+            sequence_outputs = beam_scorer.finalize(
+                input_ids,
+                beam_scores,
+                next_tokens,
+                next_indices,
+                pad_token_id=pad_token_id,
+                eos_token_id=eos_token_id,
+                max_length=stopping_criteria.max_length,
+                beam_indices=final_beam_indices,
+                decoder_prompt_len=decoder_prompt_len,
+            )
+    
+            if return_dict_in_generate:
+                if not output_scores:
+                    sequence_outputs["sequence_scores"] = None
+    
+                if self.config.is_encoder_decoder:
+                    return GenerateBeamEncoderDecoderOutput(
+                        sequences=sequence_outputs["sequences"],
+                        sequences_scores=sequence_outputs["sequence_scores"],
+                        scores=scores,
+                        logits=raw_logits,
+                        beam_indices=sequence_outputs["beam_indices"],
+                        encoder_attentions=encoder_attentions,
+                        encoder_hidden_states=encoder_hidden_states,
+                        decoder_attentions=decoder_attentions,
+                        cross_attentions=cross_attentions,
+                        decoder_hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+                else:
+                    return GenerateBeamDecoderOnlyOutput(
+                        sequences=sequence_outputs["sequences"],
+                        sequences_scores=sequence_outputs["sequence_scores"],
+                        scores=scores,
+                        logits=raw_logits,
+                        beam_indices=sequence_outputs["beam_indices"],
+                        attentions=decoder_attentions,
+                        hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+            else:
+                return sequence_outputs["sequences"]
+    
+        def _constrained_beam_search(
+            self,
+            input_ids: torch.LongTensor,
+            constrained_beam_scorer: ConstrainedBeamSearchScorer,
+            logits_processor: LogitsProcessorList,
+            stopping_criteria: StoppingCriteriaList,
+            generation_config: GenerationConfig,
+            synced_gpus: bool,
+            **model_kwargs,
+        ) -> Union[GenerateBeamOutput, torch.LongTensor]:
+            r"""
+            Generates sequences of token ids for models with a language modeling head using **constrained beam search
+            decoding** and can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.
+    
+            Parameters:
+                input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`):
+                    The sequence used as a prompt for the generation.
+                constrained_beam_scorer (`ConstrainedBeamSearchScorer`):
+                    A derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and
+                    sorted during generation, while satisfying a list of positive constraints. For more information, the
+                    documentation of [`ConstrainedBeamSearchScorer`] should be read.
+                logits_processor (`LogitsProcessorList`):
+                    An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
+                    used to modify the prediction scores of the language modeling head applied at each generation step.
+                stopping_criteria (`StoppingCriteriaList`):
+                    An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
+                    used to tell if the generation loop should stop.
+                generation_config ([`~generation.GenerationConfig`]):
+                    The generation configuration to be used as parametrization of the decoding method.
+                synced_gpus (`bool`):
+                    Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+                    `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+                model_kwargs:
+                    Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is
+                    an encoder-decoder model the kwargs should include `encoder_outputs`.
+    
+            Return:
+                [`~generation.GenerateBeamDecoderOnlyOutput`], [`~generation.GenerateBeamEncoderDecoderOutput`] or
+                `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a
+                [`~generation.GenerateBeamDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
+                `return_dict_in_generate=True` or a [`~generation.GenerateBeamEncoderDecoderOutput`] if
+                `model.config.is_encoder_decoder=True`.
+            """
+            # init values
+            pad_token_id = generation_config._pad_token_tensor
+            eos_token_id = generation_config._eos_token_tensor
+            output_attentions = generation_config.output_attentions
+            output_hidden_states = generation_config.output_hidden_states
+            output_scores = generation_config.output_scores
+            output_logits = generation_config.output_logits
+            return_dict_in_generate = generation_config.return_dict_in_generate
+    
+            batch_size = len(constrained_beam_scorer._beam_hyps)
+            num_beams = constrained_beam_scorer.num_beams
+    
+            batch_beam_size, cur_len = input_ids.shape[:2]
+            model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+    
+            if num_beams * batch_size != batch_beam_size:
+                raise ValueError(
+                    f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
+                )
+    
+            # init attention / hidden states / scores tuples
+            scores = () if (return_dict_in_generate and output_scores) else None
+            raw_logits = () if (return_dict_in_generate and output_logits) else None
+            beam_indices = (
+                tuple(() for _ in range(batch_beam_size)) if (return_dict_in_generate and output_scores) else None
+            )
+            decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+            cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+            decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+    
+            # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+            if return_dict_in_generate and self.config.is_encoder_decoder:
+                encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
+                encoder_hidden_states = (
+                    model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+                )
+    
+            # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens
+            # of the first beam are considered to avoid sampling the exact same tokens across all beams.
+            beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
+            beam_scores[:, 1:] = -1e9
+            beam_scores = beam_scores.view((batch_size * num_beams,))
+    
+            this_peer_finished = False
+    
+            decoder_prompt_len = input_ids.shape[1]  # record the prompt length of decoder
+            while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+                model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+    
+                # prepare variable output controls (note: some models won't accept all output controls)
+                model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
+                model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
+    
+                outputs = self(**model_inputs, return_dict=True)
+    
+                # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+                model_kwargs = self._update_model_kwargs_for_generation(
+                    outputs,
+                    model_kwargs,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                )
+                if synced_gpus and this_peer_finished:
+                    cur_len = cur_len + 1
+                    continue
+    
+                # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
+                # (the clone itself is always small)
+                # .float() is needed to retain precision for later logits manipulations
+                next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device)
+                next_token_scores = nn.functional.log_softmax(
+                    next_token_logits, dim=-1
+                )  # (batch_size * num_beams, vocab_size)
+    
+                next_token_scores_processed = logits_processor(input_ids, next_token_scores)
+    
+                next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(
+                    next_token_scores_processed
+                )
+    
+                scores_for_all_vocab = next_token_scores.clone()
+    
+                # Store scores, attentions and hidden_states when required
+                if return_dict_in_generate:
+                    if output_scores:
+                        scores += (next_token_scores,)
+                    if output_logits:
+                        raw_logits += (next_token_logits,)
+                    if output_attentions:
+                        decoder_attentions += (
+                            (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
+                        )
+                        if self.config.is_encoder_decoder:
+                            cross_attentions += (outputs.cross_attentions,)
+    
+                    if output_hidden_states:
+                        decoder_hidden_states += (
+                            (outputs.decoder_hidden_states,)
+                            if self.config.is_encoder_decoder
+                            else (outputs.hidden_states,)
+                        )
+    
+                # reshape for beam search
+                vocab_size = next_token_scores.shape[-1]
+                next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)
+    
+                # Sample 1 + len(eos_token_id) next tokens for each beam so we have at least 1 non eos token per beam.
+                n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0
+                next_token_scores, next_tokens = torch.topk(
+                    next_token_scores, max(2, 1 + n_eos_tokens) * num_beams, dim=1, largest=True, sorted=True
+                )
+    
+                next_indices = (next_tokens / vocab_size).long()
+                next_tokens = next_tokens % vocab_size
+    
+                # stateless
+                beam_outputs = constrained_beam_scorer.process(
+                    input_ids,
+                    next_token_scores,
+                    next_tokens,
+                    next_indices,
+                    scores_for_all_vocab,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    beam_indices=beam_indices,
+                    decoder_prompt_len=decoder_prompt_len,
+                )
+                beam_scores = beam_outputs["next_beam_scores"]
+                beam_next_tokens = beam_outputs["next_beam_tokens"]
+                beam_idx = beam_outputs["next_beam_indices"]
+    
+                input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
+    
+                # This is needed to properly delete outputs.logits which may be very large for first iteration
+                # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration
+                # IMPORTANT: Note that this should appear BEFORE the call to _reorder_cache() to save the maximum memory
+                # (that way the memory peak does not include outputs.logits)
+                del outputs
+    
+                if model_kwargs.get("past_key_values", None) is not None:
+                    model_kwargs["past_key_values"] = self._temporary_reorder_cache(
+                        model_kwargs["past_key_values"], beam_idx
+                    )
+    
+                if return_dict_in_generate and output_scores:
+                    beam_indices = tuple(beam_indices[beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices)))
+    
+                # increase cur_len
+                cur_len = cur_len + 1
+    
+                if constrained_beam_scorer.is_done or all(stopping_criteria(input_ids, scores)):
+                    this_peer_finished = True
+    
+            sequence_outputs = constrained_beam_scorer.finalize(
+                input_ids,
+                beam_scores,
+                next_tokens,
+                next_indices,
+                pad_token_id=pad_token_id,
+                eos_token_id=eos_token_id,
+                max_length=stopping_criteria.max_length,
+                beam_indices=beam_indices,
+                decoder_prompt_len=decoder_prompt_len,
+            )
+    
+            if return_dict_in_generate:
+                if not output_scores:
+                    sequence_outputs["sequence_scores"] = None
+                if self.config.is_encoder_decoder:
+                    return GenerateBeamEncoderDecoderOutput(
+                        sequences=sequence_outputs["sequences"],
+                        sequences_scores=sequence_outputs["sequence_scores"],
+                        scores=scores,
+                        logits=raw_logits,
+                        beam_indices=sequence_outputs["beam_indices"],
+                        encoder_attentions=encoder_attentions,
+                        encoder_hidden_states=encoder_hidden_states,
+                        decoder_attentions=decoder_attentions,
+                        cross_attentions=cross_attentions,
+                        decoder_hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+                else:
+                    return GenerateBeamDecoderOnlyOutput(
+                        sequences=sequence_outputs["sequences"],
+                        sequences_scores=sequence_outputs["sequence_scores"],
+                        scores=scores,
+                        logits=raw_logits,
+                        beam_indices=sequence_outputs["beam_indices"],
+                        attentions=decoder_attentions,
+                        hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+            else:
+                return sequence_outputs["sequences"]
+    
+        def _assisted_decoding(
+            self,
+            input_ids: torch.LongTensor,
+            candidate_generator: CandidateGenerator,
+            logits_processor: LogitsProcessorList,
+            stopping_criteria: StoppingCriteriaList,
+            generation_config: GenerationConfig,
+            synced_gpus: bool,
+            streamer: Optional["BaseStreamer"],
+            **model_kwargs,
+        ) -> Union[GenerateNonBeamOutput, torch.LongTensor]:
+            r"""
+            Generates sequences of token ids for models with a language modeling head using **greedy decoding** or
+            **sample** (depending on `do_sample`), assisted by candidate sequences. Assisted generation is an example of a
+            candidate decoding strategy. Can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text
+            models.
+    
+            Parameters:
+                input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                    The sequence used as a prompt for the generation.
+                candidate_generator (`CandidateGenerator`):
+                    A derived instance of [`CandidateGenerator`] that defines how candidate sequences are generated. For
+                    more information, the documentation of [`CandidateGenerator`] should be read.
+                logits_processor (`LogitsProcessorList`):
+                    An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
+                    used to modify the prediction scores of the language modeling head applied at each generation step.
+                stopping_criteria (`StoppingCriteriaList`):
+                    An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
+                    used to tell if the generation loop should stop.
+                generation_config ([`~generation.GenerationConfig`]):
+                    The generation configuration to be used as parametrization of the decoding method.
+                synced_gpus (`bool`):
+                    Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+                    `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+                streamer (`BaseStreamer`, *optional*):
+                    Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+                    through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+                model_kwargs:
+                    Additional model specific keyword arguments will be forwarded to the `forward` function of the model.
+                    If model is an encoder-decoder model the kwargs should include `encoder_outputs`.
+    
+            Return:
+                [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or
+                `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a
+                [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
+                `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if
+                `model.config.is_encoder_decoder=True`.
+            """
+            # init values
+            do_sample = generation_config.do_sample
+            output_attentions = generation_config.output_attentions
+            output_hidden_states = generation_config.output_hidden_states
+            output_scores = generation_config.output_scores
+            output_logits = generation_config.output_logits
+            return_dict_in_generate = generation_config.return_dict_in_generate
+    
+            # init attention / hidden states / scores tuples
+            scores = () if (return_dict_in_generate and output_scores) else None
+            raw_logits = () if (return_dict_in_generate and output_logits) else None
+            decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+            cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+            decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+    
+            # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+            if return_dict_in_generate and self.config.is_encoder_decoder:
+                encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
+                encoder_hidden_states = (
+                    model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+                )
+    
+            # keep track of which sequences are already finished
+            batch_size, cur_len = input_ids.shape[:2]
+            unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
+            model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs)
+    
+            this_peer_finished = False
+            is_first_iteration = True  # to preserve the same API in the output as other generation methods
+            while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
+                cur_len = input_ids.shape[1]
+    
+                #  1. Fetch candidate sequences from a `CandidateGenerator` and move to the correct device
+                candidate_input_ids, candidate_logits = candidate_generator.get_candidates(input_ids)
+                candidate_input_ids = candidate_input_ids.to(self.device)
+                if candidate_logits is not None:
+                    candidate_logits = candidate_logits.to(self.device)
+    
+                candidate_length = candidate_input_ids.shape[1] - input_ids.shape[1]
+                is_done_candidate = stopping_criteria(candidate_input_ids, None)
+    
+                # 2. Use the original model to obtain the next token logits given the candidate sequence. We obtain
+                # `candidate_length + 1` relevant logits from this process: in the event that all candidates are correct,
+                # we use this forward pass to also pick the subsequent logits in the original model.
+    
+                # 2.1. Prepare the model inputs
+                candidate_kwargs = copy.copy(model_kwargs)
+                candidate_kwargs = _prepare_attention_mask(
+                    candidate_kwargs, candidate_input_ids.shape[1], self.config.is_encoder_decoder
+                )
+                candidate_kwargs = _prepare_token_type_ids(candidate_kwargs, candidate_input_ids.shape[1])
+                if "cache_position" in candidate_kwargs:
+                    candidate_kwargs["cache_position"] = torch.cat(
+                        (
+                            candidate_kwargs["cache_position"],
+                            torch.arange(cur_len, cur_len + candidate_length, device=input_ids.device, dtype=torch.long),
+                        ),
+                        dim=0,
+                    )
+    
+                model_inputs = self.prepare_inputs_for_generation(candidate_input_ids, **candidate_kwargs)
+                if "logits_to_keep" in model_inputs:
+                    model_inputs["logits_to_keep"] = candidate_length + 1
+    
+                # 2.2. Run a forward pass on the candidate sequence
+                # prepare variable output controls (note: some models won't accept all output controls)
+                model_inputs.update({"output_attentions": output_attentions} if output_attentions else {})
+                model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {})
+    
+                outputs = self(**model_inputs)
+    
+                # 2.3. Process the new logits
+                # .float() is needed to retain precision for later logits manipulations
+                new_logits = outputs.logits[:, -candidate_length - 1 :].to(
+                    dtype=torch.float32, device=input_ids.device
+                )  # excludes the input prompt if present
+                next_token_logits = new_logits.clone()
+                if len(logits_processor) > 0:
+                    for i in range(candidate_length + 1):
+                        new_logits[:, i, :] = logits_processor(candidate_input_ids[:, : cur_len + i], new_logits[:, i, :])
+    
+                # 3. Select the accepted tokens. There are two possible cases:
+                # Case 1: `do_sample=True` and we have logits for the candidates (originally from speculative decoding)
+                # 👉 Apply algorithm 1 from the speculative decoding paper (https://huggingface.co/papers/2211.17192).
+                if do_sample and candidate_logits is not None:
+                    valid_tokens, n_matches = _speculative_sampling(
+                        candidate_input_ids,
+                        candidate_logits,
+                        candidate_length,
+                        new_logits,
+                        is_done_candidate,
+                    )
+    
+                # Case 2: all other cases (originally from assisted generation) 👉 Compare the tokens selected from the
+                # original model logits with the candidate tokens. We can keep the candidate tokens until the first
+                # mismatch, or until the max length is reached.
+                else:
+                    if do_sample:
+                        probs = new_logits.softmax(dim=-1)
+                        selected_tokens = torch.multinomial(probs[0, :, :], num_samples=1).squeeze(1)[None, :]
+                    else:
+                        selected_tokens = new_logits.argmax(dim=-1)
+    
+                    candidate_new_tokens = candidate_input_ids[:, cur_len:]
+                    n_matches = ((~(candidate_new_tokens == selected_tokens[:, :-1])).cumsum(dim=-1) < 1).sum()
+    
+                    # Ensure we don't generate beyond max_len or an EOS token
+                    if is_done_candidate and n_matches == candidate_length:
+                        n_matches -= 1
+                    valid_tokens = selected_tokens[:, : n_matches + 1]
+    
+                # 4. Update variables according to the number of matching assistant tokens. Remember: the token generated
+                # by the model after the last candidate match is also valid, as it is generated from a correct sequence.
+                # Because of this last token, assisted generation search reduces to a normal greedy search/sample if there
+                # is no match.
+    
+                # 4.1. Get the valid continuation, after the matching tokens
+                input_ids = torch.cat((input_ids, valid_tokens), dim=-1)
+                if streamer is not None:
+                    streamer.put(valid_tokens.cpu())
+                new_cur_len = input_ids.shape[1]
+    
+                # 4.2. Discard past key values relative to unused assistant tokens
+                new_cache_size = new_cur_len - 1
+                outputs.past_key_values = _crop_past_key_values(self, outputs.past_key_values, new_cache_size)
+    
+                # 5. Update the candidate generation strategy if needed
+                candidate_generator.update_candidate_strategy(input_ids, new_logits, n_matches)
+    
+                # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+                model_kwargs = self._update_model_kwargs_for_generation(
+                    outputs,
+                    model_kwargs,
+                    is_encoder_decoder=self.config.is_encoder_decoder,
+                    num_new_tokens=n_matches + 1,
+                )
+                if synced_gpus and this_peer_finished:
+                    continue
+    
+                # Store scores, attentions and hidden_states when required
+                # Assistant: modified to append one tuple element per token, as in the other generation methods.
+                if return_dict_in_generate:
+                    newly_added_length = n_matches + 1
+                    if output_scores:
+                        scores += tuple(new_logits[:, i, :] for i in range(newly_added_length))
+                    if output_logits:
+                        raw_logits += tuple(next_token_logits[:, i, :] for i in range(newly_added_length))
+    
+                    newly_added_length = new_cur_len if is_first_iteration else newly_added_length
+                    if output_attentions:
+                        if self.config.is_encoder_decoder:
+                            cross_attentions = _split_model_outputs(
+                                cross_attentions, outputs.cross_attentions, cur_len, newly_added_length
+                            )
+                            decoder_attentions = _split_model_outputs(
+                                decoder_attentions,
+                                outputs.decoder_attentions,
+                                cur_len,
+                                newly_added_length,
+                                is_decoder_attention=True,
+                            )
+                        # some (V)LLMs have hard requirement on SDPA and thus never return attn
+                        elif outputs.attentions[0] is not None:
+                            decoder_attentions = _split_model_outputs(
+                                decoder_attentions,
+                                outputs.attentions,
+                                cur_len,
+                                newly_added_length,
+                                is_decoder_attention=True,
+                            )
+                    if output_hidden_states:
+                        if self.config.is_encoder_decoder:
+                            decoder_hidden_states = _split_model_outputs(
+                                decoder_hidden_states, outputs.decoder_hidden_states, cur_len, newly_added_length
+                            )
+                        else:
+                            decoder_hidden_states = _split_model_outputs(
+                                decoder_hidden_states, outputs.hidden_states, cur_len, newly_added_length
+                            )
+    
+                unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores)
+                this_peer_finished = unfinished_sequences.max() == 0
+                is_first_iteration = False
+    
+            if streamer is not None:
+                streamer.end()
+    
+            if (
+                hasattr(candidate_generator, "assistant_model")
+                and candidate_generator.assistant_model.generation_config.num_assistant_tokens_schedule == "heuristic"
+            ):
+                candidate_generator.assistant_model.generation_config.num_assistant_tokens = (
+                    candidate_generator.num_assistant_tokens
+                )
+            if return_dict_in_generate:
+                if self.config.is_encoder_decoder:
+                    return GenerateEncoderDecoderOutput(
+                        sequences=input_ids,
+                        scores=scores,
+                        logits=raw_logits,
+                        encoder_attentions=encoder_attentions,
+                        encoder_hidden_states=encoder_hidden_states,
+                        decoder_attentions=decoder_attentions,
+                        cross_attentions=cross_attentions,
+                        decoder_hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+                else:
+                    return GenerateDecoderOnlyOutput(
+                        sequences=input_ids,
+                        scores=scores,
+                        logits=raw_logits,
+                        attentions=decoder_attentions,
+                        hidden_states=decoder_hidden_states,
+                        past_key_values=model_kwargs.get("past_key_values"),
+                    )
+            else:
+                return input_ids
+    
+        def _prefill_chunking(self, input_ids: torch.LongTensor, generation_config: GenerationConfig, **model_kwargs):
+            # Even if we are not compiling the forward, flex is always compiled when used. With chunk prefill, we may
+            # end up needing just a bit more graphs than the default (which is 8). Doing this avoids very cryptic warnings
+            torch._dynamo.config.cache_size_limit = 64
+    
+            chunk_size = generation_config.prefill_chunk_size
+            # Only chunk up the token just before last, so that decoding is completely performed outside this function
+            # (here we simply prefill the cache)
+            input_chunks = torch.split(input_ids[:, :-1], chunk_size, dim=-1)
+    
+            if "past_key_values" not in model_kwargs:
+                raise ValueError("Cannot use prefill chunking without a cache")
+    
+            model_forward = self.forward
+    
+            compile_forward = self._valid_auto_compile_criteria(model_kwargs, generation_config)
+            if compile_forward:
+                model_forward = self.get_compiled_call(generation_config.compile_config)
+    
+            attention_mask = model_kwargs.pop("attention_mask", None)
+    
+            past_length = 0
+            for input_chunk in input_chunks:
+                current_length = past_length + input_chunk.shape[-1]
+                # Prepare inputs
+                if attention_mask is not None:
+                    model_kwargs["attention_mask"] = attention_mask[:, :current_length]
+                model_kwargs["cache_position"] = torch.arange(
+                    past_length, current_length, dtype=torch.long, device=input_chunk.device
+                )
+                model_kwargs["position_ids"] = model_kwargs["cache_position"].unsqueeze(0)
+                model_inputs = self.prepare_inputs_for_generation(input_chunk, **model_kwargs)
+    
+                outputs = model_forward(**model_inputs, return_dict=True)
+    
+                model_kwargs["past_key_values"] = outputs.past_key_values
+                past_length = current_length
+    
+            model_kwargs["attention_mask"] = attention_mask
+            model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + 1
+            _ = model_kwargs.pop("position_ids", None)
+    
+            return model_kwargs
+    
+    
+    def _speculative_sampling(
+        candidate_input_ids,
+        candidate_logits,
+        candidate_length,
+        new_logits,
+        is_done_candidate,
+    ):
+        """
+        Applies sampling as in the speculative decoding paper (https://huggingface.co/papers/2211.17192, algorithm 1). Returns
+        the selected tokens, as well as the number of candidate matches.
+    
+        NOTE: Unless otherwise stated, the variable names match those in the paper.
+        """
+        new_candidate_input_ids = candidate_input_ids[:, -candidate_length:]
+        # Gets the probabilities from the logits. q_i and p_i denote the assistant and model probabilities of the tokens
+        # selected by the assistant, respectively.
+        q = candidate_logits.softmax(dim=-1)
+        q_i = q[:, torch.arange(candidate_length), new_candidate_input_ids].squeeze(0, 1)
+        p = new_logits.softmax(dim=-1)
+        p_i = p[:, torch.arange(candidate_length), new_candidate_input_ids].squeeze(0, 1)
+        probability_ratio = p_i / q_i
+    
+        # When probability_ratio > 1 (i.e. q_i(x) < p_i(x), or "assistant probability of the candidate token is smaller
+        # than the model probability for the same token"), keep the token. Otherwise reject with p = 1 - probability_ratio
+        # (= keep with p = probability_ratio). Keep all the tokens until the first rejection
+        r_i = torch.rand_like(probability_ratio)
+        is_accepted = r_i <= probability_ratio
+        n_matches = ((~is_accepted).cumsum(dim=-1) < 1).sum()  # this is `n` in algorithm 1
+    
+        # Ensure we don't generate beyond max_len or an EOS token (not in algorithm 1, but needed for correct behavior)
+        if is_done_candidate and n_matches == candidate_length:
+            # Output length is assumed to be `n_matches + 1`. Since we won't generate another token with the target model
+            # due to acceptance on EOS we fix `n_matches`
+            n_matches -= 1
+            valid_tokens = new_candidate_input_ids[:, : n_matches + 1]
+        else:
+            # Next token selection: if there is a rejection, adjust the distribution from the main model before sampling.
+            gamma = candidate_logits.shape[1]
+            p_n_plus_1 = p[:, n_matches, :]
+            if n_matches < gamma:
+                q_n_plus_1 = q[:, n_matches, :]
+                p_prime = torch.clamp((p_n_plus_1 - q_n_plus_1), min=0)
+                p_prime.div_(p_prime.sum())
+            else:
+                p_prime = p_n_plus_1
+            t = torch.multinomial(p_prime, num_samples=1).squeeze(1)[None, :]
+    
+            # The selected tokens include the matches (if any) plus the next sampled tokens
+            if n_matches > 0:
+                valid_tokens = torch.cat((new_candidate_input_ids[:, :n_matches], t), dim=-1)
+            else:
+                valid_tokens = t
+    
+        return valid_tokens, n_matches
+    
+    
+    def _split_model_outputs(outputs, new_outputs, cur_len, added_len, is_decoder_attention=False):
+        """
+        Given the (decoder/cross attentions)/(decoder hidden states) for multiple generated tokens, splits it into a tuple
+        where each member corresponds to a single generated token.
+        """
+        # Retrocompatibility: in our generation functions, the first iteration includes the attention/hidden states for the
+        # prompt.
+        if len(outputs) == 0:
+            new_tuple = ()
+            for layer in new_outputs:
+                last_dim_size = cur_len if is_decoder_attention else layer.shape[-1]
+                new_tuple += (layer[..., :cur_len, :last_dim_size],)
+            outputs += (new_tuple,)
+            # The first iteration contains the prompt + 1 generated token, let's update the length variables accordingly
+            cur_len += 1
+            added_len -= cur_len
+    
+        for i in range(added_len):
+            new_tuple = ()
+            for layer in new_outputs:
+                last_dim_size = cur_len + i if is_decoder_attention else layer.shape[-1]
+                new_tuple += (layer[..., i : i + 1, :last_dim_size],)
+            outputs += (new_tuple,)
+        return outputs
+    
+    
+    def _ranking_fast(
+        context_hidden: torch.FloatTensor,
+        next_hidden: torch.FloatTensor,
+        next_top_k_probs: torch.FloatTensor,
+        cosine_matrix_mask: torch.LongTensor,
+        alpha: float,
+        beam_width: int,
+    ) -> torch.FloatTensor:
+        """
+        Reranks the top_k candidates based on a degeneration penalty (cosine similarity with previous tokens), as described
+        in the paper "A Contrastive Framework for Neural Text Generation". Returns the index of the best candidate for each
+        row in the batch.
+        """
+        norm_context_hidden = context_hidden / context_hidden.norm(dim=2, keepdim=True)
+        norm_next_hidden = next_hidden / next_hidden.norm(dim=2, keepdim=True)
+        cosine_matrix = torch.matmul(norm_context_hidden, norm_next_hidden.transpose(1, 2)).squeeze(-1)  # [B*K, S]
+    
+        # Penalize cosine_matrix based on the cosine_matrix_mask (ignore padding positions)
+        # Using a large negative value for masked positions
+        cosine_matrix_mask = cosine_matrix_mask.to(dtype=cosine_matrix.dtype)
+        cosine_matrix_mask = (1 - cosine_matrix_mask) * torch.finfo(cosine_matrix.dtype).min
+        cosine_matrix = cosine_matrix + cosine_matrix_mask
+    
+        degeneration_penalty, _ = torch.max(cosine_matrix, dim=-1)  # [B*K]
+        next_top_k_probs = next_top_k_probs.view(-1)  # [B*K]
+        contrastive_score = (1.0 - alpha) * next_top_k_probs - alpha * degeneration_penalty
+        contrastive_score = torch.stack(torch.split(contrastive_score, beam_width))  # [B, K]
+        _, selected_idx = contrastive_score.max(dim=-1)  # [B]
+        return selected_idx
+    
+    
+    def _split(data, full_batch_size: int, split_size: int):
+        """
+        Takes care of three cases:
+        1. data is a tensor: e.g. last_hidden_state, pooler_output etc. split them on the batch_size dim
+        2. data is a tuple: e.g. hidden_states, attentions etc. Keep the tuple as it is and split each tensor in it and
+            return a list of tuples
+        3. data is a tuple of tuples, e.g. past_key_values. Keep the tuple as it is and split each tuple in it and
+            return a list of tuples of tuples
+        (see documentation of ModelOutput)
+        """
+        if data is None:
+            return [None] * (full_batch_size // split_size)
+        if isinstance(data, torch.Tensor):
+            return [data[i : i + split_size] for i in range(0, full_batch_size, split_size)]
+        # New cache format
+        elif isinstance(data, DynamicCache) or (
+            isinstance(data, EncoderDecoderCache) and isinstance(data.self_attention_cache, DynamicCache)
+        ):
+            return data.batch_split(full_batch_size, split_size)
+        elif isinstance(data, tuple):
+            # If the elements of the tuple are also tuples (e.g., past_key_values in our earlier example)
+            if isinstance(data[0], tuple):
+                return [
+                    tuple(tuple(tensor[i : i + split_size] for tensor in inner_tuple) for inner_tuple in data)
+                    for i in range(0, full_batch_size, split_size)
+                ]
+    
+            else:
+                return [
+                    tuple(sub_tensor[i : i + split_size] for sub_tensor in data)
+                    for i in range(0, full_batch_size, split_size)
+                ]
+        else:
+            raise TypeError(f"Unexpected attribute type: {type(data)}")
+    
+    
+    def _split_model_inputs(
+        model_input: Union[ModelOutput, dict], split_size: int, full_batch_size: int, config: PretrainedConfig
+    ) -> list[Union[ModelOutput, dict]]:
+        """
+        Split a ModelOutput object (or its subclasses) or Dict into a list of same-class objects based on a specified split
+        size. The input object is dict when it was prepared for forward pass and ModelOutput when it was returned from
+        previous forward pass.
+        """
+        # Edge case: if model_input is None, return a list of Nones
+        # this happens with Whisper where encoder_outputs is None
+        if model_input is None:
+            return [model_input] * (full_batch_size // split_size)
+        # Infer the class from the object
+        model_output_cls = type(model_input)
+        if (full_batch_size % split_size) != 0:
+            raise ValueError("`full_batch_size` must be divisible by `split_size`")
+    
+        if split_size > full_batch_size:
+            raise ValueError("`split_size` must be smaller or equal to `full_batch_size`")
+    
+        # Helper function to split tensors or tuples of tensors
+    
+        # Find all the dataclass fields (e.g., last_hidden_state, pooler_output etc.) and split them
+        keys = (
+            model_input.__dataclass_fields__.keys() if hasattr(model_input, "__dataclass_fields__") else model_input.keys()
+        )
+        # We only keep keys that are in the model_input
+        keys = [k for k in keys if k in model_input]
+        # Here we can have four types of values: tensors, tuples of tensors and booleans, and encoder_outputs which is a
+        # ModelOutput object.
+        # bool should not be split but replicated for each split
+        bool_keys = [k for k in keys if isinstance(model_input[k], bool) or k == "cache_position"]
+        keys_to_ignore = ["cache_position", "encoder_outputs", "logits_to_keep"]
+        non_bool_keys = [k for k in keys if not isinstance(model_input[k], bool) and k not in keys_to_ignore]
+    
+        # we split the tensors and tuples of tensors
+        data_split_list = [
+            {k: _split(model_input[k], full_batch_size, split_size)[i] for k in non_bool_keys}
+            for i in range(full_batch_size // split_size)
+        ]
+        # bool values are the same and replicated for each split
+        bool_data = {k: model_input[k] for k in bool_keys}
+        # encoder_outputs is a ModelOutput object and should be split by its own
+        if "encoder_outputs" in model_input:
+            encoder_outputs_split = _split_model_inputs(
+                model_input["encoder_outputs"], split_size, full_batch_size, config.get_text_config()
+            )
+            data_split_list = [
+                {**data_split, "encoder_outputs": encoder_outputs_split[i]} for i, data_split in enumerate(data_split_list)
+            ]
+        # logits_to_keep should be replicated for each split, similar to bool values
+        if "logits_to_keep" in model_input:
+            data_split_list = [
+                {**data_split, "logits_to_keep": model_input["logits_to_keep"]} for data_split in data_split_list
+            ]
+    
+        # Convert each dictionary in the list to an object of the inferred class
+        split_model_inputs: list[Union[ModelOutput, dict]] = [
+            model_output_cls(**data_split, **bool_data) for data_split in data_split_list
+        ]
+    
+        return split_model_inputs
+    
+    
+    def stack_model_outputs(model_outputs: list[ModelOutput], config: PretrainedConfig) -> ModelOutput:
+        """
+        Stack a list of ModelOutput objects (or its subclasses) along the batch_size dimension. The function infers the
+        specific ModelOutput subclass from the list provided.
+        """
+        if not model_outputs:
+            raise ValueError("Input list is empty.")
+    
+        # Infer the class from the first object in the list
+        model_output_cls = type(model_outputs[0])
+    
+        # Ensure all objects are of the same type
+        if not all(isinstance(obj, model_output_cls) for obj in model_outputs):
+            raise ValueError("All elements in the list should be of the same type.")
+    
+        # Helper function to concat tensors or tuples of tensors
+        def _concat(data):
+            """
+            Reverse of `_split` function above.
+            """
+            if any(data is None for data in data):
+                return None
+            if isinstance(data[0], torch.Tensor):
+                return torch.cat(data, dim=0)
+            # New cache format
+            elif isinstance(data[0], DynamicCache):
+                return DynamicCache.from_batch_splits(data)
+            elif isinstance(data[0], EncoderDecoderCache):
+                return EncoderDecoderCache.from_batch_splits(data)
+            elif isinstance(data[0], tuple):
+                # If the elements of the tuple are also tuples (e.g., past_key_values in our earlier example)
+                if isinstance(data[0][0], tuple):
+                    return tuple(
+                        tuple(torch.cat([attr[i][j] for attr in data], dim=0) for j in range(len(data[0][0])))
+                        for i in range(len(data[0]))
+                    )
+                else:
+                    return tuple(torch.cat([attr[i] for attr in data], dim=0) for i in range(len(data[0])))
+            elif isinstance(data[0], (int, float)):
+                # If the elements are integers or floats, return a tensor
+                return torch.tensor(data)
+            else:
+                raise TypeError(f"Unexpected attribute type: {type(data[0])}")
+    
+        # Use a dictionary comprehension to gather attributes from all objects and concatenate them
+        concatenated_data = {
+            k: _concat([getattr(model_output, k) for model_output in model_outputs])
+            for k in model_output_cls.__dataclass_fields__.keys()
+        }
+    
+        # Return a new object of the inferred class with the concatenated attributes
+        return model_output_cls(**concatenated_data)
+    
+    
+    def _relative_top_filter(
+        scores: torch.FloatTensor,
+        baseline_scores: torch.FloatTensor,
+        relative_top: float = 0.1,
+        filter_value: float = -float("Inf"),
+        base_filter_value=-1e-3,
+        min_tokens_to_keep: int = 1,
+    ) -> torch.FloatTensor:
+        """
+        Reference: https://github.com/XiangLi1999/ContrastiveDecoding/blob/170e9142e92159c1237d731e240f5eb14aabf428/transformers/src/transformers/generation_logits_process.py#L235
+        Apply filtering to only keep tokens with a probability above a certain threshold. The threshold is defined as `relative_top` * max probability in the distribution.
+        """
+        scores_normalized = scores.log_softmax(dim=-1)
+        baseline_scores_normalized = baseline_scores.log_softmax(dim=-1)
+        sorted_logits, sorted_indices = torch.sort(scores_normalized, descending=True)
+        min_thresh = sorted_logits[..., min_tokens_to_keep - 1]
+        probs_max = torch.max(scores_normalized, dim=-1).values
+        probs_thresh = probs_max + np.log(relative_top)
+        probs_thresh = torch.min(min_thresh, probs_thresh)
+        probs_thresh = probs_thresh.unsqueeze(-1)
+        baseline_scores_normalized[scores_normalized < probs_thresh] = base_filter_value
+        scores_normalized[scores_normalized < probs_thresh] = filter_value
+        return scores_normalized, baseline_scores_normalized
+    
+    
+    def _dola_select_contrast(
+        candidate_premature_layers: list[int],
+        candidate_premature_logits: dict[int, torch.FloatTensor],
+        final_logits: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        if len(candidate_premature_layers) == 1:
+            base_logits = candidate_premature_logits[candidate_premature_layers[0]]
+            final_logits, base_logits = _relative_top_filter(final_logits, base_logits)
+            logits = final_logits - base_logits
+            return logits
+    
+        # 1. Stacking all premature_layers into a new dimension
+        stacked_premature_layers = torch.stack([candidate_premature_logits[i] for i in candidate_premature_layers], dim=0)
+    
+        # 2. Calculate the softmax values for mature_layer and all premature_layers
+        # shape: (batch_size, vocab_size)
+        softmax_mature_layer = F.softmax(final_logits, dim=-1)
+        # shape: (num_premature_layers, batch_size, vocab_size)
+        softmax_premature_layers = F.softmax(stacked_premature_layers, dim=-1)
+    
+        # 3. Calculate the average distribution
+        # shape: (num_premature_layers, batch_size, vocab_size)
+        avg_dist = 0.5 * (softmax_mature_layer[None, :, :] + softmax_premature_layers)
+    
+        # 4. Calculate log-softmax for the KL divergence
+        # shape: (batch_size, vocab_size)
+        log_softmax_mature_layer = F.log_softmax(final_logits, dim=-1)
+        # shape: (num_premature_layers, batch_size, vocab_size)
+        log_softmax_premature_layers = F.log_softmax(stacked_premature_layers, dim=-1)
+    
+        # 5. Calculate the KL divergences and then the JS divergences
+        # shape: (num_premature_layers, batch_size)
+        kl1 = F.kl_div(log_softmax_mature_layer[None, :, :], avg_dist, reduction="none").mean(-1)
+        # shape: (num_premature_layers, batch_size)
+        kl2 = F.kl_div(log_softmax_premature_layers, avg_dist, reduction="none").mean(-1)
+        js_divs = 0.5 * (kl1 + kl2)  # shape: (num_premature_layers, batch_size)
+    
+        # 6. Reduce the batchmean
+        js_divs = js_divs.mean(-1)  # shape: (num_premature_layers,)
+        premature_layer = candidate_premature_layers[int(js_divs.argmax().item())]
+    
+        base_logits = candidate_premature_logits[premature_layer]
+        final_logits, base_logits = _relative_top_filter(final_logits, base_logits)
+        logits = final_logits - base_logits
+        return logits
+    
+    
+    
+    </script>
+        
 <body class="bg-gray-900 text-gray-200 flex flex-col h-screen">
     <header class="bg-gray-800/50 backdrop-blur-sm border-b border-gray-700 p-4 shadow-lg">
         <h1 class="text-2xl font-bold text-center text-white">Bloatedness Visualizer</h1>
@@ -120,43 +7989,10 @@
         const addTabBtn = document.getElementById('add-tab-btn');
 
         // --- Example Code ---
-        const exampleCodeMain = `
-# A specific model implementation that uses an external class
-class MyAwesomeModel(BaseModel, AttentionMixin):
-    def __init__(self, config):
-        super().__init__(config)
-        self._init_attention(config)
-
-    def forward(self, input_ids) -> BaseModelOutputWithPast:
-        """The main forward pass for the model."""
-        x = self._split_heads(input_ids)
-        # The return type is a class not defined in this file
-        return BaseModelOutputWithPast(last_hidden_state=x)
-`;
-
-        const exampleCodeDeps = `
-# This class is an external dependency, not defined here
-# from modeling_outputs import BaseModelOutputWithPast
-
-# A common mixin for attention mechanisms
-class AttentionMixin:
-    def _init_attention(self, config):
-        self.num_heads = config.num_heads
-        self.head_dim = config.hidden_size // config.num_heads
-
-    def _split_heads(self, tensor):
-        # Logic to split tensor for multi-head attention
-        return tensor
-
-# A base class for all models, similar to PreTrainedModel
-class BaseModel:
-    def __init__(self, config):
-        self.config = config
-
-    def from_pretrained(cls, model_name):
-        # Load weights
-        pass
-`;
+        const exampleCodeMain = document.getElementById("main-code").textContent;
+
+        const exampleCodeDeps = document.getElementById("dependencies").textContent;
+
 
         // --- Tab Management ---
         let tabCounter = 0;