#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/interns2_preview/modular_interns2_preview.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_interns2_preview.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 The Intern team, The Qwen Team and The 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 math
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from typing import Any, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import initialization as init
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache
from transformers.conversion_mapping import register_checkpoint_conversion_mapping
from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter
from transformers.generation import GenerationMixin
from transformers.integrations import use_experts_implementation, use_kernelized_func
from transformers.masking_utils import create_bidirectional_mask, create_causal_mask
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPooling,
    BaseModelOutputWithPoolingAndCrossAttentions,
    ModelOutput,
)
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from transformers.processing_utils import Unpack
from transformers.pytorch_utils import apply_chunking_to_forward
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check
from transformers.utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults
from transformers.utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
from transformers.utils.output_capturing import OutputRecorder, capture_outputs

from .configuration_interns2_preview import (
    InternS2PreviewBertConfig,
    InternS2PreviewConfig,
    InternS2PreviewTextConfig,
    InternS2PreviewTimeSeriesConfig,
    InternS2PreviewTimeSeriesForecasterConfig,
    InternS2PreviewVisionConfig,
)


if is_causal_conv1d_available():
    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
else:
    causal_conv1d_update, causal_conv1d_fn = None, None

if is_flash_linear_attention_available():
    from fla.modules import FusedRMSNormGated
    from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
else:
    chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None
    FusedRMSNormGated = None

logger = logging.get_logger(__name__)


class InternS2PreviewVisionRotaryEmbedding(nn.Module):
    inv_freq: torch.Tensor  # fix linting for `register_buffer`

    def __init__(self, dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        self.dim = dim
        self.theta = theta
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, seqlen: int) -> torch.Tensor:
        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(seq, self.inv_freq)
        return freqs


class InternS2PreviewDynamicCache:
    """
    A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the linear attention
    cache (which has a constant shape regardless of seq_len).

    This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states`
    and `ssm_states` for gated deltanet cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor
    For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`,
    while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors).
    For linear attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors),
    while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`,
    and `recurrent_states` represents the recurrent state and has a shape of `(batch_size, d_inner, d_state)`.
    """

    is_compileable = False

    def __init__(self, config: InternS2PreviewConfig):
        super().__init__()
        self.layer_types = config.layer_types
        self.transformer_layers = [
            i for i in range(config.num_hidden_layers) if self.layer_types[i] == "full_attention"
        ]
        self.last_linear_layer = len(self.layer_types) - 1 - self.layer_types[::-1].index("linear_attention")

        # Initialize everything to None -> will be lazy initialized to allow multi-gpu (device_map) inference
        self.conv_states = [None for _ in range(config.num_hidden_layers)]
        self.recurrent_states = [None for _ in range(config.num_hidden_layers)]
        self.key_cache = [None for _ in range(config.num_hidden_layers)]
        self.value_cache = [None for _ in range(config.num_hidden_layers)]

    def __len__(self):
        return len(self.layer_types)

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: dict[str, Any] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if self.key_cache[layer_idx] is None:
            self.key_cache[layer_idx] = key_states
            self.value_cache[layer_idx] = value_states
        else:
            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2)
            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2)

        return self.key_cache[layer_idx], self.value_cache[layer_idx]

    def reorder_cache(self, beam_idx: torch.LongTensor):
        """Reorders the cache for beam search, given the selected beam indices."""
        for layer_idx in range(len(self.key_cache)):
            if self.key_cache[layer_idx] is not None:
                device = self.key_cache[layer_idx].device
                beam_idx = beam_idx.to(device)
                self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx)
                self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx)

            if self.conv_states[layer_idx] is not None:
                device = self.conv_states[layer_idx].device
                beam_idx = beam_idx.to(device)
                self.conv_states[layer_idx] = self.conv_states[layer_idx].index_select(0, beam_idx)
                self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select(0, beam_idx)

    def get_seq_length(self, layer_idx: int | None = 0) -> int:
        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
        # take any layer that contains cache and not empty tensor
        layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx
        if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None:
            return 0
        return self.key_cache[layer_idx].shape[-2]

    def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]:
        """
        Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for
        the given layer at `layer_idx`.
        The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.
        """
        kv_offset = 0
        query_length = cache_position.shape[0]
        past_seen_tokens = self.get_seq_length(layer_idx)
        kv_length = query_length + past_seen_tokens
        return kv_length, kv_offset

    @property
    def has_previous_state(self):
        """We have a previous state if the last linear (conv) layer was already updated."""
        return self.conv_states[self.last_linear_layer] is not None


class InternS2PreviewRMSNormGated(nn.Module):
    def __init__(self, hidden_size, eps=1e-6, **kwargs):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states, gate=None):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        # Norm before gate
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        hidden_states = self.weight * hidden_states.to(input_dtype)
        hidden_states = hidden_states * F.silu(gate.to(torch.float32))

        return hidden_states.to(input_dtype)


def apply_mask_to_padding_states(hidden_states, attention_mask):
    """
    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
    """
    # NOTE: attention mask is a 2D boolean tensor
    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
        dtype = hidden_states.dtype
        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)

    return hidden_states


is_fast_path_available = all(
    (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)
)


def torch_causal_conv1d_update(
    hidden_states,
    conv_state,
    weight,
    bias=None,
    activation=None,
):
    _, hidden_size, seq_len = hidden_states.shape
    state_len = conv_state.shape[-1]

    hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
    conv_state.copy_(hidden_states_new[:, :, -state_len:])
    out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
    out = F.silu(out[:, :, -seq_len:])
    out = out.to(hidden_states.dtype)
    return out


def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
    """This function is intended to align with the l2norm implementation in the FLA library."""
    inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
    return x * inv_norm


def torch_chunk_gated_delta_rule(
    query,
    key,
    value,
    g,
    beta,
    chunk_size=64,
    initial_state=None,
    output_final_state=False,
    use_qk_l2norm_in_kernel=False,
):
    initial_dtype = query.dtype
    if use_qk_l2norm_in_kernel:
        query = l2norm(query, dim=-1, eps=1e-6)
        key = l2norm(key, dim=-1, eps=1e-6)
    query, key, value, beta, g = [
        x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
    ]

    batch_size, num_heads, sequence_length, k_head_dim = key.shape
    v_head_dim = value.shape[-1]
    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
    query = F.pad(query, (0, 0, 0, pad_size))
    key = F.pad(key, (0, 0, 0, pad_size))
    value = F.pad(value, (0, 0, 0, pad_size))
    beta = F.pad(beta, (0, pad_size))
    g = F.pad(g, (0, pad_size))
    total_sequence_length = sequence_length + pad_size
    scale = 1 / (query.shape[-1] ** 0.5)
    query = query * scale

    v_beta = value * beta.unsqueeze(-1)
    k_beta = key * beta.unsqueeze(-1)
    # reshape to chunks
    query, key, value, k_beta, v_beta = [
        x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) for x in (query, key, value, k_beta, v_beta)
    ]
    g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)

    # chunk decay
    g = g.cumsum(dim=-1)
    decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
    attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
    for i in range(1, chunk_size):
        row = attn[..., i, :i].clone()
        sub = attn[..., :i, :i].clone()
        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
    value = attn @ v_beta
    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
    last_recurrent_state = (
        torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
        if initial_state is None
        else initial_state.to(value)
    )
    core_attn_out = torch.zeros_like(value)
    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)

    # for each chunk
    for i in range(0, total_sequence_length // chunk_size):
        q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
        attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
        v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
        v_new = v_i - v_prime
        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
        core_attn_out[:, :, i] = attn_inter + attn @ v_new
        last_recurrent_state = (
            last_recurrent_state * g[:, :, i, -1, None, None].exp()
            + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
        )

    if not output_final_state:
        last_recurrent_state = None
    core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1])
    core_attn_out = core_attn_out[:, :, :sequence_length]
    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
    return core_attn_out, last_recurrent_state


def torch_recurrent_gated_delta_rule(
    query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False
):
    initial_dtype = query.dtype
    if use_qk_l2norm_in_kernel:
        query = l2norm(query, dim=-1, eps=1e-6)
        key = l2norm(key, dim=-1, eps=1e-6)
    query, key, value, beta, g = [
        x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
    ]

    batch_size, num_heads, sequence_length, k_head_dim = key.shape
    v_head_dim = value.shape[-1]
    scale = 1 / (query.shape[-1] ** 0.5)
    query = query * scale

    core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(value)
    last_recurrent_state = (
        torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
        if initial_state is None
        else initial_state.to(value)
    )

    for i in range(sequence_length):
        q_t = query[:, :, i]
        k_t = key[:, :, i]
        v_t = value[:, :, i]
        g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1)
        beta_t = beta[:, :, i].unsqueeze(-1)

        last_recurrent_state = last_recurrent_state * g_t
        kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2)
        delta = (v_t - kv_mem) * beta_t
        last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2)
        core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2)

    if not output_final_state:
        last_recurrent_state = None
    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
    return core_attn_out, last_recurrent_state


class InternS2PreviewGatedDeltaNet(nn.Module):
    def __init__(self, config: InternS2PreviewConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_v_heads = config.linear_num_value_heads
        self.num_k_heads = config.linear_num_key_heads
        self.head_k_dim = config.linear_key_head_dim
        self.head_v_dim = config.linear_value_head_dim
        self.key_dim = self.head_k_dim * self.num_k_heads
        self.value_dim = self.head_v_dim * self.num_v_heads

        self.conv_kernel_size = config.linear_conv_kernel_dim
        self.layer_idx = layer_idx
        self.activation = config.hidden_act
        self.act = ACT2FN[config.hidden_act]
        self.layer_norm_epsilon = config.rms_norm_eps

        # QKV
        self.conv_dim = self.key_dim * 2 + self.value_dim
        self.conv1d = nn.Conv1d(
            in_channels=self.conv_dim,
            out_channels=self.conv_dim,
            bias=False,
            kernel_size=self.conv_kernel_size,
            groups=self.conv_dim,
            padding=self.conv_kernel_size - 1,
        )

        # time step projection (discretization)
        # instantiate once and copy inv_dt in init_weights of PretrainedModel
        self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads))

        A = torch.empty(self.num_v_heads).uniform_(0, 16)
        self.A_log = nn.Parameter(torch.log(A))

        self.norm = (
            InternS2PreviewRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
            if FusedRMSNormGated is None
            else FusedRMSNormGated(
                self.head_v_dim,
                eps=self.layer_norm_epsilon,
                activation=self.activation,
                device=torch.cuda.current_device(),
                dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(),
            )
        )

        self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)

        self.causal_conv1d_fn = causal_conv1d_fn
        self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update
        self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
        self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule

        if not is_fast_path_available:
            logger.warning_once(
                "The fast path is not available because one of the required library is not installed. Falling back to "
                "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and"
                " https://github.com/Dao-AILab/causal-conv1d"
            )

        self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False)
        self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False)
        self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)
        self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cache_params: InternS2PreviewDynamicCache | None = None,
        cache_position: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
    ):
        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)

        # Set up dimensions for reshapes later
        batch_size, seq_len, _ = hidden_states.shape

        use_precomputed_states = (
            cache_params is not None
            and cache_params.has_previous_state
            and seq_len == 1
            and cache_position is not None
        )

        # getting projected states from cache if it exists
        if cache_params is not None:
            conv_state = cache_params.conv_states[self.layer_idx]
            recurrent_state = cache_params.recurrent_states[self.layer_idx]

        mixed_qkv = self.in_proj_qkv(hidden_states)
        mixed_qkv = mixed_qkv.transpose(1, 2)

        z = self.in_proj_z(hidden_states)
        z = z.reshape(batch_size, seq_len, -1, self.head_v_dim)

        b = self.in_proj_b(hidden_states)
        a = self.in_proj_a(hidden_states)

        if use_precomputed_states:
            # 2. Convolution sequence transformation
            # NOTE: the conv state is updated in `causal_conv1d_update`
            mixed_qkv = self.causal_conv1d_update(
                mixed_qkv,
                conv_state,
                self.conv1d.weight.squeeze(1),
                self.conv1d.bias,
                self.activation,
            )
        else:
            if cache_params is not None:
                conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
                cache_params.conv_states[self.layer_idx] = conv_state
            if self.causal_conv1d_fn is not None:
                mixed_qkv = self.causal_conv1d_fn(
                    x=mixed_qkv,
                    weight=self.conv1d.weight.squeeze(1),
                    bias=self.conv1d.bias,
                    activation=self.activation,
                    seq_idx=None,
                )
            else:
                mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len])

        mixed_qkv = mixed_qkv.transpose(1, 2)
        query, key, value = torch.split(
            mixed_qkv,
            [
                self.key_dim,
                self.key_dim,
                self.value_dim,
            ],
            dim=-1,
        )

        query = query.reshape(batch_size, seq_len, -1, self.head_k_dim)
        key = key.reshape(batch_size, seq_len, -1, self.head_k_dim)
        value = value.reshape(batch_size, seq_len, -1, self.head_v_dim)

        beta = b.sigmoid()
        # If the model is loaded in fp16, without the .float() here, A might be -inf
        g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
        if self.num_v_heads // self.num_k_heads > 1:
            query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
            key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)

        if not use_precomputed_states:
            core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
                query,
                key,
                value,
                g=g,
                beta=beta,
                initial_state=None,
                output_final_state=cache_params is not None,
                use_qk_l2norm_in_kernel=True,
            )

        else:
            core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
                query,
                key,
                value,
                g=g,
                beta=beta,
                initial_state=recurrent_state,
                output_final_state=cache_params is not None,
                use_qk_l2norm_in_kernel=True,
            )

        # Update cache
        if cache_params is not None:
            cache_params.recurrent_states[self.layer_idx] = last_recurrent_state

        # reshape input data into 2D tensor
        core_attn_out = core_attn_out.reshape(-1, self.head_v_dim)
        z = z.reshape(-1, self.head_v_dim)
        core_attn_out = self.norm(core_attn_out, z)
        core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1)

        output = self.out_proj(core_attn_out)
        return output


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)


# Adapted from transformers.models.glm.modular_glm.apply_rotary_pos_emb
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    """Applies Rotary Position Embedding to the query and key tensors.

    Removes the interleaving of cos and sin from GLM

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
        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)

    # Keep half or full tensor for later concatenation
    rotary_dim = cos.shape[-1]
    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]

    # Apply rotary embeddings on the first half or full tensor
    q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
    k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)

    # Concatenate back to full shape
    q_embed = torch.cat([q_embed, q_pass], dim=-1)
    k_embed = torch.cat([k_embed, k_pass], dim=-1)
    return q_embed, k_embed


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: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    **kwargs: Unpack[TransformersKwargs],
):
    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 attention_mask is not None:
        attn_weights = attn_weights + attention_mask

    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


class InternS2PreviewRMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.zeros(dim))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float())
        # Llama does x.to(float16) * w whilst InternS2Preview is (x * w).to(float16)
        # See https://github.com/huggingface/transformers/pull/29402
        output = output * (1.0 + self.weight.float())
        return output.type_as(x)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.eps}"


@use_kernelized_func(apply_rotary_pos_emb)
class InternS2PreviewAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: InternS2PreviewConfig, layer_idx: int):
        super().__init__()
        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.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True
        self.q_proj = nn.Linear(
            config.hidden_size, config.num_attention_heads * self.head_dim * 2, 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.q_norm = InternS2PreviewRMSNorm(
            self.head_dim, eps=config.rms_norm_eps
        )  # unlike olmo, only on the head dim!
        self.k_norm = InternS2PreviewRMSNorm(
            self.head_dim, eps=config.rms_norm_eps
        )  # thus post q_norm does not need reshape

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None,
        past_key_values: Cache | None = None,
        cache_position: torch.LongTensor | None = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states, gate = torch.chunk(
            self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1
        )
        gate = gate.reshape(*input_shape, -1)

        query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_values 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}
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = attn_output * torch.sigmoid(gate)

        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


class InternS2PreviewMLP(nn.Module):
    def __init__(self, config: InternS2PreviewConfig, intermediate_size: int):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = intermediate_size
        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_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


@use_experts_implementation
class InternS2PreviewExperts(nn.Module):
    """Collection of expert weights stored as 3D tensors."""

    def __init__(self, config):
        super().__init__()
        self.num_experts = config.num_experts
        self.hidden_dim = config.hidden_size
        self.intermediate_dim = config.moe_intermediate_size
        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(
        self,
        hidden_states: torch.Tensor,
        top_k_index: torch.Tensor,
        top_k_weights: torch.Tensor,
    ) -> torch.Tensor:
        final_hidden_states = torch.zeros_like(hidden_states)
        with torch.no_grad():
            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
            expert_mask = expert_mask.permute(2, 1, 0)
            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

        for expert_idx in expert_hit:
            expert_idx = expert_idx[0]
            if expert_idx == self.num_experts:
                continue
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
            current_hidden_states = self.act_fn(gate) * up
            current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))

        return final_hidden_states


class InternS2PreviewTopKRouter(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.top_k = config.num_experts_per_tok
        self.num_experts = config.num_experts
        self.hidden_dim = config.hidden_size
        self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))

    def forward(self, hidden_states):
        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
        router_logits = F.linear(hidden_states, self.weight)  # (seq_len, num_experts)
        router_logits = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1)
        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
        router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
        router_top_value = router_top_value.to(router_logits.dtype)
        router_scores = router_top_value
        return router_logits, router_scores, router_indices


class InternS2PreviewSparseMoeBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate = InternS2PreviewTopKRouter(config)
        self.experts = InternS2PreviewExperts(config)
        self.shared_expert = InternS2PreviewMLP(config, intermediate_size=config.shared_expert_intermediate_size)
        self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        batch_size, sequence_length, hidden_dim = hidden_states.shape
        hidden_states_reshaped = hidden_states.view(-1, hidden_dim)
        shared_expert_output = self.shared_expert(hidden_states_reshaped)
        _, routing_weights, selected_experts = self.gate(hidden_states_reshaped)
        expert_output = self.experts(hidden_states_reshaped, selected_experts, routing_weights)

        shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states_reshaped)) * shared_expert_output

        expert_output += shared_expert_output
        expert_output = expert_output.reshape(batch_size, sequence_length, hidden_dim)
        return expert_output


# The MoE experts are kept fused in memory (`gate_up_proj` / `down_proj` 3D
# tensors) for an efficient grouped forward, but the checkpoint stores each
# expert separately as `...mlp.experts.{i}.{gate,up,down}_proj.weight`. Mirroring
# the upstream Qwen3-MoE implementation, we register a bidirectional
# `WeightConverter`: it merges the sharded expert weights into the fused
# parameters on load, and `revert_weight_conversion` (called from
# `save_pretrained`) splits them back into the sharded layout on save.
_INTERN_S2_EXPERTS_CONVERSIONS = [
    WeightConverter(
        source_patterns=[
            "mlp.experts.*.gate_proj.weight",
            "mlp.experts.*.up_proj.weight",
        ],
        target_patterns="mlp.experts.gate_up_proj",
        operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
    ),
    WeightConverter(
        source_patterns="mlp.experts.*.down_proj.weight",
        target_patterns="mlp.experts.down_proj",
        operations=[MergeModulelist(dim=0)],
    ),
]
# Lookup is by class name first, then `config.model_type`; register under both
# the multimodal model type and the text model type so the conversion applies
# whether the full model or the text backbone is loaded standalone.
for _model_type in ("intern_s2_preview", "qwen3_5_moe_text"):
    register_checkpoint_conversion_mapping(_model_type, _INTERN_S2_EXPERTS_CONVERSIONS, overwrite=True)


class InternS2PreviewDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: InternS2PreviewTextConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.layer_type = config.layer_types[layer_idx]
        if self.layer_type == "linear_attention":
            self.linear_attn = InternS2PreviewGatedDeltaNet(config, layer_idx)
        elif self.layer_type == "full_attention":
            self.self_attn = InternS2PreviewAttention(config, layer_idx)
        self.mlp = InternS2PreviewSparseMoeBlock(config)
        self.input_layernorm = InternS2PreviewRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = InternS2PreviewRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        cache_position: torch.LongTensor | None = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> torch.FloatTensor:
        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)

        # Token Mixer
        if self.layer_type == "linear_attention":
            hidden_states = self.linear_attn(
                hidden_states=hidden_states,
                cache_params=past_key_values,
                cache_position=cache_position,
                attention_mask=attention_mask,
            )
        elif self.layer_type == "full_attention":
            # Self Attention
            hidden_states, _ = self.self_attn(
                hidden_states=hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
                **kwargs,
            )

        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        # For the MoE layers, we need to unpack
        if isinstance(hidden_states, tuple):
            hidden_states, _ = hidden_states
        hidden_states = residual + hidden_states

        return hidden_states


class InternS2PreviewPreTrainedModel(PreTrainedModel):
    config: InternS2PreviewConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["InternS2PreviewDecoderLayer", "InternS2PreviewVisionBlock"]
    _skip_keys_device_placement = "past_key_values"
    _supports_flash_attn = True
    _supports_sdpa = True
    _keys_to_ignore_on_load_unexpected = [r"^mtp.*"]
    _can_record_outputs = {
        "router_logits": OutputRecorder(InternS2PreviewTopKRouter, index=0),
        "hidden_states": InternS2PreviewDecoderLayer,
        "attentions": InternS2PreviewAttention,
    }
    _is_stateful = True

    @torch.no_grad()
    def _init_weights(self, module):
        super()._init_weights(module)
        if isinstance(module, InternS2PreviewGatedDeltaNet):
            init.ones_(module.dt_bias)
            init.copy_(module.A_log, torch.empty_like(module.A_log).uniform_(0, 16).log_())
        # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
        elif isinstance(module, InternS2PreviewRMSNorm):
            init.zeros_(module.weight)
        elif isinstance(module, InternS2PreviewExperts):
            init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range)
            init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range)
        elif isinstance(module, InternS2PreviewSparseMoeBlock):
            init.normal_(module.gate.weight, mean=0.0, std=self.config.initializer_range)
        elif isinstance(module, InternS2PreviewVisionRotaryEmbedding):
            inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim))
            init.copy_(module.inv_freq, inv_freq)


class InternS2PreviewVisionMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
        self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, hidden_state):
        return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))


class InternS2PreviewVisionPatchEmbed(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.patch_size = config.patch_size
        self.temporal_patch_size = config.temporal_patch_size
        self.in_channels = config.in_channels
        self.embed_dim = config.hidden_size

        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
        self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        target_dtype = self.proj.weight.dtype
        hidden_states = hidden_states.view(
            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
        )
        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
        return hidden_states


class InternS2PreviewVisionPatchMerger(nn.Module):
    def __init__(self, config: InternS2PreviewVisionConfig, use_postshuffle_norm=False) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
        self.use_postshuffle_norm = use_postshuffle_norm
        self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)
        self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
        self.act_fn = nn.GELU()
        self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size)
        x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
        return x


def apply_rotary_pos_emb_vision(
    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    orig_q_dtype = q.dtype
    orig_k_dtype = k.dtype
    q, k = q.float(), k.float()
    cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    q_embed = q_embed.to(orig_q_dtype)
    k_embed = k_embed.to(orig_k_dtype)
    return q_embed, k_embed


class InternS2PreviewVisionAttention(nn.Module):
    def __init__(self, config: InternS2PreviewVisionConfig) -> None:
        super().__init__()
        self.dim = config.hidden_size
        self.num_heads = config.num_heads
        self.head_dim = self.dim // self.num_heads
        self.num_key_value_groups = 1  # needed for eager attention
        self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
        self.proj = nn.Linear(self.dim, self.dim)
        self.scaling = self.head_dim**-0.5
        self.config = config
        self.attention_dropout = 0.0
        self.is_causal = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        rotary_pos_emb: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs,
    ) -> torch.Tensor:
        seq_length = hidden_states.shape[0]
        query_states, key_states, value_states = (
            self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
        )
        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)

        query_states = query_states.transpose(0, 1).unsqueeze(0)
        key_states = key_states.transpose(0, 1).unsqueeze(0)
        value_states = value_states.transpose(0, 1).unsqueeze(0)

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )

        if is_flash_attention_requested(self.config):
            # Flash Attention: Use cu_seqlens for variable length attention
            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
            attn_output, _ = attention_interface(
                self,
                query_states,
                key_states,
                value_states,
                attention_mask=None,
                scaling=self.scaling,
                dropout=0.0 if not self.training else self.attention_dropout,
                cu_seq_lens_q=cu_seqlens,
                cu_seq_lens_k=cu_seqlens,
                max_length_q=max_seqlen,
                max_length_k=max_seqlen,
                is_causal=False,
                **kwargs,
            )
        else:
            # Other implementations: Process each chunk separately
            lengths = cu_seqlens[1:] - cu_seqlens[:-1]
            splits = [
                torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
            ]

            attn_outputs = [
                attention_interface(
                    self,
                    q,
                    k,
                    v,
                    attention_mask=None,
                    scaling=self.scaling,
                    dropout=0.0 if not self.training else self.attention_dropout,
                    is_causal=False,
                    **kwargs,
                )[0]
                for q, k, v in zip(*splits)
            ]
            attn_output = torch.cat(attn_outputs, dim=1)

        attn_output = attn_output.reshape(seq_length, -1).contiguous()
        attn_output = self.proj(attn_output)
        return attn_output


class InternS2PreviewVisionBlock(GradientCheckpointingLayer):
    def __init__(self, config, attn_implementation: str = "sdpa") -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
        self.attn = InternS2PreviewVisionAttention(config=config)
        self.mlp = InternS2PreviewVisionMLP(config=config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        rotary_pos_emb: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs,
    ) -> torch.Tensor:
        hidden_states = hidden_states + self.attn(
            self.norm1(hidden_states),
            cu_seqlens=cu_seqlens,
            rotary_pos_emb=rotary_pos_emb,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
        return hidden_states


class InternS2PreviewVisionModel(InternS2PreviewPreTrainedModel):
    config: InternS2PreviewVisionConfig
    _no_split_modules = ["InternS2PreviewVisionBlock"]
    _can_record_outputs = {
        "hidden_states": InternS2PreviewVisionBlock,
        "attentions": InternS2PreviewVisionAttention,
    }

    def __init__(self, config, *inputs, **kwargs) -> None:
        super().__init__(config, *inputs, **kwargs)
        self.spatial_merge_size = config.spatial_merge_size
        self.patch_size = config.patch_size
        self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size

        self.patch_embed = InternS2PreviewVisionPatchEmbed(
            config=config,
        )

        self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
        self.num_grid_per_side = int(config.num_position_embeddings**0.5)

        head_dim = config.hidden_size // config.num_heads
        self.rotary_pos_emb = InternS2PreviewVisionRotaryEmbedding(head_dim // 2)

        self.blocks = nn.ModuleList([InternS2PreviewVisionBlock(config) for _ in range(config.depth)])
        self.merger = InternS2PreviewVisionPatchMerger(
            config=config,
            use_postshuffle_norm=False,
        )

        self.gradient_checkpointing = False

        self.post_init()

    def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
        merge_size = self.spatial_merge_size
        grid_thw_list = grid_thw.tolist()

        max_hw = max(max(h, w) for _, h, w in grid_thw_list)
        freq_table = self.rotary_pos_emb(max_hw)  # (max_hw, dim // 2)
        device = freq_table.device

        total_tokens = sum(t * h * w for t, h, w in grid_thw_list)
        pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device)

        offset = 0
        for num_frames, height, width in grid_thw_list:
            merged_h, merged_w = height // merge_size, width // merge_size

            block_rows = torch.arange(merged_h, device=device)  # block row indices
            block_cols = torch.arange(merged_w, device=device)  # block col indices
            intra_row = torch.arange(merge_size, device=device)  # intra-block row offsets
            intra_col = torch.arange(merge_size, device=device)  # intra-block col offsets

            # Compute full-resolution positions
            row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]
            col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]

            row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
            col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)

            coords = torch.stack((row_idx, col_idx), dim=-1)

            if num_frames > 1:
                coords = coords.repeat(num_frames, 1)

            num_tokens = coords.shape[0]
            pos_ids[offset : offset + num_tokens] = coords
            offset += num_tokens

        embeddings = freq_table[pos_ids]  # lookup rotary embeddings
        embeddings = embeddings.flatten(1)
        return embeddings

    def fast_pos_embed_interpolate(self, grid_thw):
        grid_thw_list = grid_thw.tolist()
        grid_ts = [row[0] for row in grid_thw_list]
        grid_hs = [row[1] for row in grid_thw_list]
        grid_ws = [row[2] for row in grid_thw_list]
        device = self.pos_embed.weight.device

        idx_list = [[] for _ in range(4)]
        weight_list = [[] for _ in range(4)]

        for t, h, w in grid_thw_list:
            h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h)
            w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w)

            h_idxs_floor = h_idxs.int()
            w_idxs_floor = w_idxs.int()
            h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
            w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)

            dh = h_idxs - h_idxs_floor
            dw = w_idxs - w_idxs_floor

            base_h = h_idxs_floor * self.num_grid_per_side
            base_h_ceil = h_idxs_ceil * self.num_grid_per_side

            indices = [
                (base_h[None].T + w_idxs_floor[None]).flatten(),
                (base_h[None].T + w_idxs_ceil[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
            ]

            weights = [
                ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
                ((1 - dh)[None].T * dw[None]).flatten(),
                (dh[None].T * (1 - dw)[None]).flatten(),
                (dh[None].T * dw[None]).flatten(),
            ]

            for i in range(4):
                idx_list[i].extend(indices[i].tolist())
                weight_list[i].extend(weights[i].tolist())

        idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=device)
        weight_tensor = torch.tensor(weight_list, dtype=self.pos_embed.weight.dtype, device=device)
        pos_embeds = self.pos_embed(idx_tensor).to(device) * weight_tensor[:, :, None]
        patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]

        patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)])

        patch_pos_embeds_permute = []
        merge_size = self.config.spatial_merge_size
        for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
            pos_embed = pos_embed.repeat(t, 1)
            pos_embed = (
                pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1)
                .permute(0, 1, 3, 2, 4, 5)
                .flatten(0, 4)
            )
            patch_pos_embeds_permute.append(pos_embed)
        patch_pos_embeds = torch.cat(patch_pos_embeds_permute)
        return patch_pos_embeds

    @merge_with_config_defaults
    @capture_outputs
    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
        """
        Args:
            hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
                The final hidden states of the model.
            grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
                The temporal, height and width of feature shape of each image in LLM.

        Returns:
            `torch.Tensor`: hidden_states.
        """
        hidden_states = self.patch_embed(hidden_states)

        pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
        hidden_states = hidden_states + pos_embeds

        rotary_pos_emb = self.rot_pos_emb(grid_thw)

        seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.reshape(seq_len, -1)
        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
        position_embeddings = (emb.cos(), emb.sin())

        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
            dim=0,
            # Select dtype based on the following factors:
            #  - FA2 requires that cu_seqlens_q must have dtype int32
            #  - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
            # See https://github.com/huggingface/transformers/pull/34852 for more information
            dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
        )
        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

        for blk in self.blocks:
            hidden_states = blk(
                hidden_states,
                cu_seqlens=cu_seqlens,
                position_embeddings=position_embeddings,
                **kwargs,
            )

        merged_hidden_states = self.merger(hidden_states)

        return BaseModelOutputWithPooling(
            last_hidden_state=hidden_states,
            pooler_output=merged_hidden_states,
        )


class InternS2PreviewTextRotaryEmbedding(nn.Module):
    inv_freq: torch.Tensor  # fix linting for `register_buffer`

    def __init__(self, config: InternS2PreviewTextConfig, device=None):
        super().__init__()
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config

        self.rope_type = self.config.rope_parameters["rope_type"]
        rope_init_fn: Callable = self.compute_default_rope_parameters
        if self.rope_type != "default":
            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)

        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
        self.mrope_section = config.rope_parameters.get("mrope_section", [11, 11, 10])

    @staticmethod
    def compute_default_rope_parameters(
        config: InternS2PreviewTextConfig | None = None,
        device: Optional["torch.device"] = None,
        seq_len: int | None = None,
    ) -> tuple["torch.Tensor", float]:
        """
        Computes the inverse frequencies according to the original RoPE implementation
        Args:
            config ([`~transformers.PreTrainedConfig`]):
                The model configuration.
            device (`torch.device`):
                The device to use for initialization of the inverse frequencies.
            seq_len (`int`, *optional*):
                The current sequence length. Unused for this type of RoPE.
        Returns:
            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
        """
        base = config.rope_parameters["rope_theta"]
        partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
        dim = int(head_dim * partial_rotary_factor)

        attention_factor = 1.0  # Unused in this type of RoPE

        # Compute the inverse frequencies
        inv_freq = 1.0 / (
            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
        )
        return inv_freq, attention_factor

    @torch.no_grad()
    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
    def forward(self, x, position_ids):
        # In contrast to other models, InternS2Preview has different position ids for the grids
        # So we expand the inv_freq to shape (3, ...)
        if position_ids.ndim == 2:
            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (3, bs, 1, positions)

        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
            freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
            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 apply_interleaved_mrope(self, freqs, mrope_section):
        """Apply interleaved MRoPE to 3D rotary embeddings.
        Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
        interleaved [THWTHWTHW...TT], preserving frequency continuity.
        args:
            x: (3, bs, seq_len, head_dim // 2)
            mrope_section: (3,)
        returns:
            x_t: (bs, seq_len, head_dim // 2)
        """
        freqs_t = freqs[0]  # just overwrite the first dimension T
        for dim, offset in enumerate((1, 2), start=1):  # H, W
            length = mrope_section[dim] * 3
            idx = slice(offset, length, 3)
            freqs_t[..., idx] = freqs[dim, ..., idx]
        return freqs_t


class InternS2PreviewTextModel(InternS2PreviewPreTrainedModel):
    def __init__(self, config: InternS2PreviewTextConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
        self.layers = nn.ModuleList(
            [InternS2PreviewDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = InternS2PreviewRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = InternS2PreviewTextRotaryEmbedding(config=config)
        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()

    @merge_with_config_defaults
    @capture_outputs
    @auto_docstring
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        cache_position: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BaseModelOutputWithPast:
        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 None:
            inputs_embeds = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = InternS2PreviewDynamicCache(config=self.config)

        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
            )

        # mrope: the hard coded `3` is for temporal, height and width.
        if position_ids is None:
            position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
        elif position_ids.ndim == 2:
            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)

        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
            text_position_ids = position_ids[0]
            position_ids = position_ids[1:]
        else:
            text_position_ids = position_ids[0]

        causal_mask = create_causal_mask(
            config=self.config,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            cache_position=cache_position,
            past_key_values=past_key_values,
            position_ids=text_position_ids,
        )
        linear_attn_mask = self._update_linear_attn_mask(attention_mask, cache_position)

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
            layer_mask = linear_attn_mask if decoder_layer.layer_type == "linear_attention" else causal_mask

            hidden_states = decoder_layer(
                hidden_states,
                position_embeddings=position_embeddings,
                attention_mask=layer_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
                cache_position=cache_position,
                **kwargs,
            )

        hidden_states = self.norm(hidden_states)

        return InternS2PreviewModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values,
        )

    def _update_linear_attn_mask(self, attention_mask, cache_position):
        """
        NOTE: Left-padding is used for linear attention mask.
        No need for zeroing states when
            1. Cached forward
            2. Attending to all inputs
        """
        linear_attn_mask = attention_mask
        if cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1)):
            linear_attn_mask = None
        return linear_attn_mask


class InternS2PreviewTimeSeriesAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        is_decoder: bool = False,
        bias: bool = True,
        is_causal: bool = False,
        layer_idx: int | None = None,
        config: InternS2PreviewTimeSeriesConfig | None = None,
        num_key_value_groups: int = 1,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        self.config = config

        if (self.head_dim * num_heads) != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
                f" and `num_heads`: {num_heads})."
            )
        self.scaling = self.head_dim**-0.5
        self.is_decoder = is_decoder
        self.is_causal = is_causal

        if layer_idx is None and is_decoder:
            logger.warning_once(
                f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "
                "will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
                "when creating this class."
            )
        self.layer_idx = layer_idx

        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.num_key_value_groups = num_key_value_groups

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_value_states: torch.Tensor | None = None,
        past_key_values: Cache | None = None,
        attention_mask: torch.Tensor | None = None,
        output_attentions: bool = False,
        cache_position: torch.Tensor | None = None,
        # TODO: we need a refactor so that the different attention modules can get their specific kwargs
        # ATM, we have mixed things encoder, decoder, and encoder-decoder attn
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
        """Input shape: Batch x Time x Channel"""

        # if key_value_states are provided this layer is used as a cross-attention layer
        # for the decoder
        is_cross_attention = key_value_states is not None

        # determine input shapes
        bsz, tgt_len = hidden_states.shape[:-1]
        q_input_shape = (bsz, tgt_len, -1, self.head_dim)

        # Scaling is susceptible to floating point arithmetics' inprecisions
        # which can lead to different results (this is dependent from model
        # to model, e.g. intern_s2_preview_time_series is one such case). We therefore keep the
        # original order of scaling to follow the original implementation
        # and enforce no scaling (1.0) in the attention call below.
        query_states = self.q_proj(hidden_states) * self.scaling
        query_states = query_states.view(*q_input_shape)
        query_states = query_states.transpose(1, 2).contiguous()

        # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache`
        if past_key_values is not None and isinstance(past_key_values, EncoderDecoderCache):
            is_updated = past_key_values.is_updated.get(self.layer_idx)
            if is_cross_attention:
                # after the first generated id, we can subsequently re-use all key/value_states from cache
                past_key_values.is_updated[self.layer_idx] = True
                past_key_values = past_key_values.cross_attention_cache
            else:
                past_key_values = past_key_values.self_attention_cache

        # use key_value_states if cross attention
        current_states = key_value_states if key_value_states is not None else hidden_states
        if is_cross_attention and past_key_values and is_updated:
            # reuse k,v, cross_attentions
            key_states = past_key_values.layers[self.layer_idx].keys
            value_states = past_key_values.layers[self.layer_idx].values
        else:
            key_states = self.k_proj(current_states).view(bsz, -1, self.num_heads, self.head_dim)
            value_states = self.v_proj(current_states).view(bsz, -1, self.num_heads, self.head_dim)
            key_states = key_states.transpose(1, 2).contiguous()
            value_states = value_states.transpose(1, 2).contiguous()
            if past_key_values is not None:
                # save all key/value_states to cache to be re-used for fast auto-regressive generation
                cache_position = cache_position if not is_cross_attention else None
                key_states, value_states = past_key_values.update(
                    key_states, value_states, self.layer_idx, {"cache_position": cache_position}
                )

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.dropout,
            scaling=1.0,
            output_attentions=output_attentions,
            **kwargs,
        )

        attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights


class InternS2PreviewTimeSeriesEncoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: InternS2PreviewTimeSeriesConfig):
        super().__init__()
        self.embed_dim = config.d_model

        self.self_attn = InternS2PreviewTimeSeriesAttention(
            embed_dim=self.embed_dim,
            num_heads=config.encoder_attention_heads,
            dropout=config.attention_dropout,
            config=config,
        )
        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout
        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = nn.LayerNorm(self.embed_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        output_attentions: bool = False,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
            attention_mask (`torch.FloatTensor`): attention mask of size
                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
        """
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)
        hidden_states, attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
        )
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states

        if hidden_states.dtype == torch.float16:
            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)

        return hidden_states, attn_weights


class InternS2PreviewTimeSeriesEncoder(PreTrainedModel):
    config: InternS2PreviewTimeSeriesConfig
    base_model_prefix = "model"
    main_input_name = "input_features"
    input_modalities = ("audio", "text")
    supports_gradient_checkpointing = True
    _no_split_modules = ["InternS2PreviewTimeSeriesEncoderLayer", "WhisperDecoderLayer"]
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = True

    _can_compile_fullgraph = True

    def __init__(self, config: InternS2PreviewTimeSeriesConfig):
        super().__init__(config)
        self.dropout = config.dropout
        self.layerdrop = config.encoder_layerdrop

        self.embed_dim = config.d_model
        self.num_mel_bins = config.num_mel_bins
        # self.padding_idx = config.pad_token_id
        self.max_source_positions = config.max_source_positions
        self.embed_scale = math.sqrt(self.embed_dim) if config.scale_embedding else 1.0

        self.conv1 = nn.Conv1d(self.num_mel_bins, self.embed_dim, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(self.embed_dim, self.embed_dim, kernel_size=3, stride=2, padding=1)
        self.embed_positions = nn.Embedding(self.max_source_positions, self.embed_dim)

        self.layers = nn.ModuleList(
            [InternS2PreviewTimeSeriesEncoderLayer(config) for _ in range(config.encoder_layers)]
        )
        self.layer_norm = nn.LayerNorm(config.d_model)

        self.gradient_checkpointing = False
        self.post_init()

        self.mask_type = None
        self.chunk_length = None

        self.adapt_in = nn.Linear(config.ts_adapt_in_dim, 80)
        self.adapt_out = nn.Linear(self.embed_dim, config.ts_adapt_out_dim)

    def _freeze_parameters(self):
        for param in self.parameters():
            param.requires_grad = False
        self._requires_grad = False

    def get_input_embeddings(self) -> nn.Module:
        return self.conv1

    def set_input_embeddings(self, value: nn.Module):
        self.conv1 = value

    def define_masktype(self, masktype, chunk_length=None):
        self.mask_type = masktype
        self.chunk_length = chunk_length

    def _make_causal_mask(
        self, input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
    ):
        """
        Make causal mask used for bi-directional self-attention.
        """
        bsz, tgt_len = input_ids_shape
        mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
        mask_cond = torch.arange(mask.size(-1), device=device)
        mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
        mask = mask.to(dtype)

        if past_key_values_length > 0:
            mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
        return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)

    # Copied from transformers.models.bart.modeling_bart._expand_mask
    def _expand_mask(self, mask: torch.Tensor, dtype: torch.dtype, tgt_len: int = None):
        """
        Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
        """
        # print(mask.size())
        bsz, src_len = mask.size()
        tgt_len = tgt_len if tgt_len is not None else src_len

        expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

        inverted_mask = 1.0 - expanded_mask

        return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)

    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
        # create causal mask
        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        combined_attention_mask = None

        if input_shape[-1] > 1:
            combined_attention_mask = self._make_causal_mask(
                input_shape,
                inputs_embeds.dtype,
                device=inputs_embeds.device,
                past_key_values_length=past_key_values_length,
            )

        if attention_mask is not None:
            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
            expanded_attn_mask = self._expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
            combined_attention_mask = (
                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
            )
        return combined_attention_mask

    def prepare_chunk_attention_mask(self, attention_mask, input_shape, inputs_embeds):

        block_size = round(self.chunk_length / 4 * 2)
        matrix_size = input_shape[1]

        matrix = torch.ones(matrix_size, matrix_size)

        num_full_blocks = round(matrix_size // block_size)
        remainder = matrix_size % block_size
        for i in range(num_full_blocks):
            row_start = i * block_size
            col_start = i * block_size
            matrix[row_start : row_start + block_size, col_start : col_start + block_size] = torch.zeros(
                block_size, block_size
            )

        if remainder > 0:
            last_row_start = num_full_blocks * block_size
            last_col_start = num_full_blocks * block_size
            matrix[last_row_start : last_row_start + remainder, last_col_start : last_col_start + remainder] = (
                torch.zeros(remainder, remainder)
            )

        matrix = matrix * -65504
        matrix = matrix.unsqueeze(0).unsqueeze(0).repeat(input_shape[0], 1, 1, 1)
        attention_mask = matrix.to(inputs_embeds.device)
        return attention_mask

    def prepare_padding_mask(self, attention_mask, input_shape, inputs_embeds, input_lens):

        matrix_size = input_shape[1]
        matrix_list = []

        for i in range(input_shape[0]):
            padding_matrix = torch.ones(matrix_size, matrix_size) * -65504
            padding_matrix[: input_lens[i], : input_lens[i]] = 0
            matrix_list.append(padding_matrix)

        attention_mask = torch.stack(matrix_list).unsqueeze(1).to(inputs_embeds.dtype).to(inputs_embeds.device)

        return attention_mask

    def forward(
        self,
        input_features,
        input_length=None,
        attention_mask=None,
        head_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        causal=True,
    ):
        # (N, T, C) -> (T, N, C) -> (N, C, T)
        input_features = input_features.permute(1, 0, 2)
        input_features = self.adapt_in(input_features)
        input_features = input_features.permute(1, 2, 0)

        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
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # (N, C, T) -> (N, C, T//2)
        inputs_embeds = nn.functional.gelu(self.conv1(input_features))
        inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))

        # (N, C, T) -> (N, T, C)
        inputs_embeds = inputs_embeds.permute(0, 2, 1)  # torch.Size([1, 100, 768])
        embed_pos = self.embed_positions.weight  # torch.Size([1500, 768])

        if inputs_embeds.shape[1] > embed_pos.shape[0]:
            target_len = inputs_embeds.shape[1]
            padding = [0, 0, 0, target_len - embed_pos.shape[0]]

            embed_pos = nn.functional.pad(embed_pos, pad=padding, mode="constant", value=0)
            hidden_states = inputs_embeds[:, : embed_pos.shape[0], :] + embed_pos
        else:
            hidden_states = inputs_embeds + embed_pos[: inputs_embeds.shape[1], :]
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        input_shape = inputs_embeds.size()[:-1]
        past_key_values_length = 0
        attention_mask = None
        if causal:
            if self.mask_type == "chunk":
                attention_mask = self.prepare_chunk_attention_mask(attention_mask, input_shape, inputs_embeds)
            else:
                attention_mask = self._prepare_decoder_attention_mask(
                    attention_mask, input_shape, inputs_embeds, past_key_values_length
                )
        else:
            attention_mask = self.prepare_padding_mask(
                attention_mask, input_shape, inputs_embeds, (input_length + 1) // 2
            )

        if head_mask is not None:
            assert head_mask.size()[0] == (len(self.layers)), (
                f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
            )

        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (self.layer_norm(hidden_states),)
            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
            to_drop = False
            if self.training:
                dropout_probability = torch.rand([])
                if dropout_probability < self.layerdrop:  # skip the layer
                    to_drop = True

            if to_drop:
                layer_outputs = (None, None)
            else:
                if self.gradient_checkpointing and self.training:

                    def create_custom_forward(module):
                        def custom_forward(*inputs):
                            return module(*inputs, output_attentions)

                        return custom_forward

                    layer_outputs = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(encoder_layer),
                        hidden_states,
                        attention_mask,
                        (head_mask[idx] if head_mask is not None else None),
                    )
                else:
                    layer_outputs = encoder_layer(
                        hidden_states,
                        attention_mask,
                        # TODO test whether OK
                        # layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                        output_attentions=output_attentions,
                    )

                hidden_states = layer_outputs[0]

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        # (N, T, C) -> (T, N, C)
        hidden_states = hidden_states.permute(1, 0, 2)
        hidden_states = self.layer_norm(hidden_states)
        hidden_states = self.adapt_out(hidden_states)

        # (T, N, C) -> (N, T, C)
        hidden_states = hidden_states.permute(1, 0, 2)
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
        return ModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)


class InternS2PreviewTimeSeriesFixPositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 200000):
        super().__init__()
        pe = torch.zeros(max_len, d_model, dtype=torch.float)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1).to(torch.float32)  # (max_len, 1, d_model)
        self.register_buffer("pe", pe, persistent=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (seq_len, batch_size, d_model)
        x = x + self.pe[: x.size(0), :].to(device=x.device, dtype=x.dtype)
        return x.clone()


class InternS2PreviewBertSelfAttention(nn.Module):
    def __init__(self, config, is_causal=False, layer_idx=None):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
                f"heads ({config.num_attention_heads})"
            )
        self.config = config

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.scaling = self.attention_head_size**-0.5

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

        self.is_decoder = config.is_decoder
        self.is_causal = is_causal
        self.layer_idx = layer_idx

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.FloatTensor | None = None,
        past_key_values: Cache | None = None,
        cache_position: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.attention_head_size)

        # get all proj
        query_layer = self.query(hidden_states).view(*hidden_shape).transpose(1, 2)
        key_layer = self.key(hidden_states).view(*hidden_shape).transpose(1, 2)
        value_layer = self.value(hidden_states).view(*hidden_shape).transpose(1, 2)

        if past_key_values is not None:
            # decoder-only intern_s2_preview_bert can have a simple dynamic cache for example
            current_past_key_values = past_key_values
            if isinstance(past_key_values, EncoderDecoderCache):
                current_past_key_values = past_key_values.self_attention_cache

            # save all key/value_layer to cache to be re-used for fast auto-regressive generation
            key_layer, value_layer = current_past_key_values.update(
                key_layer,
                value_layer,
                self.layer_idx,
                {"cache_position": cache_position},
            )

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )

        attn_output, attn_weights = attention_interface(
            self,
            query_layer,
            key_layer,
            value_layer,
            attention_mask,
            dropout=0.0 if not self.training else self.dropout.p,
            scaling=self.scaling,
            **kwargs,
        )
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        return attn_output, attn_weights


class InternS2PreviewBertCrossAttention(nn.Module):
    def __init__(self, config, is_causal=False, layer_idx=None):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
                f"heads ({config.num_attention_heads})"
            )
        self.config = config

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.scaling = self.attention_head_size**-0.5

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

        self.is_causal = is_causal
        self.layer_idx = layer_idx

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.FloatTensor | None = None,
        attention_mask: torch.FloatTensor | None = None,
        past_key_values: EncoderDecoderCache | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor]:
        # determine input shapes
        bsz, tgt_len = hidden_states.shape[:-1]
        src_len = encoder_hidden_states.shape[1]

        q_input_shape = (bsz, tgt_len, -1, self.attention_head_size)
        kv_input_shape = (bsz, src_len, -1, self.attention_head_size)

        # get query proj
        query_layer = self.query(hidden_states).view(*q_input_shape).transpose(1, 2)

        is_updated = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False
        if past_key_values is not None and is_updated:
            # reuse k,v, cross_attentions
            key_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].keys
            value_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].values
        else:
            key_layer = self.key(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2)
            value_layer = self.value(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2)

            if past_key_values is not None:
                # save all states to the cache
                key_layer, value_layer = past_key_values.cross_attention_cache.update(
                    key_layer, value_layer, self.layer_idx
                )
                # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
                past_key_values.is_updated[self.layer_idx] = True

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )

        attn_output, attn_weights = attention_interface(
            self,
            query_layer,
            key_layer,
            value_layer,
            attention_mask,
            dropout=0.0 if not self.training else self.dropout.p,
            scaling=self.scaling,
            **kwargs,
        )
        attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
        return attn_output, attn_weights


class InternS2PreviewBertSelfOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class InternS2PreviewBertAttention(nn.Module):
    def __init__(self, config, is_causal=False, layer_idx=None, is_cross_attention=False):
        super().__init__()
        self.is_cross_attention = is_cross_attention
        attention_class = InternS2PreviewBertCrossAttention if is_cross_attention else InternS2PreviewBertSelfAttention
        self.self = attention_class(config, is_causal=is_causal, layer_idx=layer_idx)
        self.output = InternS2PreviewBertSelfOutput(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.FloatTensor | None = None,
        encoder_hidden_states: torch.FloatTensor | None = None,
        encoder_attention_mask: torch.FloatTensor | None = None,
        past_key_values: Cache | None = None,
        cache_position: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor]:
        attention_mask = attention_mask if not self.is_cross_attention else encoder_attention_mask
        attention_output, attn_weights = self.self(
            hidden_states,
            encoder_hidden_states=encoder_hidden_states,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            cache_position=cache_position,
            **kwargs,
        )
        attention_output = self.output(attention_output, hidden_states)
        return attention_output, attn_weights


class InternS2PreviewBertIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states


class InternS2PreviewBertOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class InternS2PreviewBertLayer(GradientCheckpointingLayer):
    def __init__(self, config, layer_idx=None):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = InternS2PreviewBertAttention(config, is_causal=config.is_decoder, layer_idx=layer_idx)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        self.intermediate = InternS2PreviewBertIntermediate(config)
        self.output = InternS2PreviewBertOutput(config)
        if self.add_cross_attention:
            if not self.is_decoder:
                raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
            self.crossattention = InternS2PreviewBertAttention(
                config,
                is_causal=False,
                layer_idx=layer_idx,
                is_cross_attention=True,
            )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.FloatTensor | None = None,
        encoder_hidden_states: torch.FloatTensor | None = None,
        encoder_attention_mask: torch.FloatTensor | None = None,
        past_key_values: Cache | None = None,
        cache_position: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor]:
        self_attention_output, _ = self.attention(
            hidden_states,
            attention_mask,
            past_key_values=past_key_values,
            cache_position=cache_position,
            **kwargs,
        )
        attention_output = self_attention_output

        if self.is_decoder and encoder_hidden_states is not None:
            if not hasattr(self, "crossattention"):
                raise ValueError(
                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
                    " by setting `config.add_cross_attention=True`"
                )

            cross_attention_output, _ = self.crossattention(
                self_attention_output,
                None,  # attention_mask
                encoder_hidden_states,
                encoder_attention_mask,
                past_key_values=past_key_values,
                **kwargs,
            )
            attention_output = cross_attention_output

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
        )
        return layer_output

    def feed_forward_chunk(self, attention_output):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output


class InternS2PreviewBertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings."""

    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.register_buffer(
            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
        )
        self.register_buffer(
            "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
        )

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        token_type_ids: torch.LongTensor | None = None,
        position_ids: torch.LongTensor | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        past_key_values_length: int = 0,
    ) -> torch.Tensor:
        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        batch_size, seq_length = input_shape

        if position_ids is None:
            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]

        # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
        # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
        # issue #5664
        if token_type_ids is None:
            if hasattr(self, "token_type_ids"):
                # NOTE: We assume either pos ids to have bsz == 1 (broadcastable) or bsz == effective bsz (input_shape[0])
                buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1)
                buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids)
                token_type_ids = buffered_token_type_ids.expand(batch_size, seq_length)
            else:
                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)
        embeddings = inputs_embeds + token_type_embeddings

        position_embeddings = self.position_embeddings(position_ids)
        embeddings = embeddings + position_embeddings

        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings


class InternS2PreviewBertEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList(
            [InternS2PreviewBertLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)]
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.FloatTensor | None = None,
        encoder_hidden_states: torch.FloatTensor | None = None,
        encoder_attention_mask: torch.FloatTensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = None,
        cache_position: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor] | BaseModelOutputWithPastAndCrossAttentions:
        for i, layer_module in enumerate(self.layer):
            hidden_states = layer_module(
                hidden_states,
                attention_mask,
                encoder_hidden_states,  # as a positional argument for gradient checkpointing
                encoder_attention_mask=encoder_attention_mask,
                past_key_values=past_key_values,
                cache_position=cache_position,
                **kwargs,
            )

        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
        )


class InternS2PreviewBertPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output


@auto_docstring
class InternS2PreviewBertPreTrainedModel(PreTrainedModel):
    config_class = InternS2PreviewBertConfig
    base_model_prefix = "interns2_preview_bert"
    supports_gradient_checkpointing = True
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = True
    _supports_attention_backend = True
    _can_record_outputs = {
        "hidden_states": InternS2PreviewBertLayer,
        "attentions": InternS2PreviewBertSelfAttention,
        "cross_attentions": InternS2PreviewBertCrossAttention,
    }

    @torch.no_grad()
    def _init_weights(self, module):
        """Initialize the weights"""
        super()._init_weights(module)
        if isinstance(module, InternS2PreviewBertEmbeddings):
            init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
            init.zeros_(module.token_type_ids)


class InternS2PreviewBertModel(InternS2PreviewBertPreTrainedModel):
    _no_split_modules = ["InternS2PreviewBertEmbeddings", "InternS2PreviewBertLayer"]

    def __init__(self, config, add_pooling_layer=True):
        r"""
        add_pooling_layer (bool, *optional*, defaults to `True`):
            Whether to add a pooling layer
        """
        super().__init__(config)
        self.config = config
        self.gradient_checkpointing = False

        self.embeddings = InternS2PreviewBertEmbeddings(config)
        self.encoder = InternS2PreviewBertEncoder(config)

        self.pooler = InternS2PreviewBertPooler(config) if add_pooling_layer else None

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    @merge_with_config_defaults
    @capture_outputs
    @auto_docstring
    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        token_type_ids: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        encoder_hidden_states: torch.Tensor | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = None,
        cache_position: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor] | BaseModelOutputWithPoolingAndCrossAttentions:
        if self.config.is_decoder:
            use_cache = use_cache if use_cache is not None else self.config.use_cache
        else:
            use_cache = False

        if use_cache and past_key_values is None:
            past_key_values = (
                EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
                if encoder_hidden_states is not None or self.config.is_encoder_decoder
                else DynamicCache(config=self.config)
            )

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if input_ids is not None:
            device = input_ids.device
            seq_length = input_ids.shape[1]
        else:
            device = inputs_embeds.device
            seq_length = inputs_embeds.shape[1]

        past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
        if cache_position is None:
            cache_position = torch.arange(past_key_values_length, past_key_values_length + seq_length, device=device)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )

        attention_mask, encoder_attention_mask = self._create_attention_masks(
            attention_mask=attention_mask,
            encoder_attention_mask=encoder_attention_mask,
            embedding_output=embedding_output,
            encoder_hidden_states=encoder_hidden_states,
            cache_position=cache_position,
            past_key_values=past_key_values,
        )

        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            position_ids=position_ids,
            **kwargs,
        )
        sequence_output = encoder_outputs.last_hidden_state
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
        )

    def _create_attention_masks(
        self,
        attention_mask,
        encoder_attention_mask,
        embedding_output,
        encoder_hidden_states,
        cache_position,
        past_key_values,
    ):
        if self.config.is_decoder:
            attention_mask = create_causal_mask(
                config=self.config,
                inputs_embeds=embedding_output,
                attention_mask=attention_mask,
                cache_position=cache_position,
                past_key_values=past_key_values,
            )
        else:
            attention_mask = create_bidirectional_mask(
                config=self.config,
                inputs_embeds=embedding_output,
                attention_mask=attention_mask,
            )

        if encoder_attention_mask is not None:
            encoder_attention_mask = create_bidirectional_mask(
                config=self.config,
                inputs_embeds=embedding_output,
                attention_mask=encoder_attention_mask,
                encoder_hidden_states=encoder_hidden_states,
            )

        return attention_mask, encoder_attention_mask


class MRQFormer(nn.Module):
    def __init__(
        self,
        num_query_list,
        hidden_size: int = 1024,
        num_layers: int = 6,
        cross_attention_freq: int = 2,
        encoder_hidden_size: int = None,
        output_size: int = None,
    ):
        super().__init__()

        self.num_query_list = num_query_list
        self.hidden_size = hidden_size

        self.query_token_list = nn.ParameterList()
        for num_queries in self.num_query_list:
            self.query_token_list.append(nn.Parameter(torch.randn(1, num_queries, hidden_size)))

        bert_config = InternS2PreviewBertConfig(
            vocab_size=1,
            hidden_size=hidden_size,
            num_hidden_layers=num_layers,
            num_attention_heads=hidden_size // 16,
            intermediate_size=hidden_size * 4,
            max_position_embeddings=max(self.num_query_list),
            add_cross_attention=True,
            is_decoder=True,
            cross_attention_freq=cross_attention_freq,
        )
        self.transformer = InternS2PreviewBertModel(bert_config)

        if encoder_hidden_size is not None and encoder_hidden_size != hidden_size:
            self.encoder_proj = nn.Linear(encoder_hidden_size, hidden_size)
        else:
            self.encoder_proj = nn.Identity()

        if output_size == None:
            output_size = hidden_size

        self.proj_out = nn.Sequential(
            nn.Linear(in_features=hidden_size, out_features=4 * hidden_size),
            nn.ReLU(),
            nn.Linear(in_features=4 * hidden_size, out_features=output_size),
        )

    def forward(self, encoder_features: torch.Tensor, attention_mask=None, res_idx=0):
        B = encoder_features.shape[0]

        encoder_features = self.encoder_proj(encoder_features)

        query_tokens = self.query_token_list[res_idx]
        query_tokens = query_tokens.expand(B, -1, -1)

        encoder_attention_mask = None
        if attention_mask is not None:
            encoder_attention_mask = attention_mask

        outputs = self.transformer(
            inputs_embeds=query_tokens,
            encoder_hidden_states=encoder_features,
            encoder_attention_mask=encoder_attention_mask,
            return_dict=True,
        )

        outputs = outputs.last_hidden_state

        outputs = self.proj_out(outputs)

        return outputs


class TotalFixlenSingleResChunkQformerSubsampling(nn.Module):
    def __init__(self, config: InternS2PreviewTimeSeriesConfig):
        super().__init__()

        self.selfatt = config.subsampling_selfatt
        self.patch_list = [config.subsampling_patch]
        self.num_query_list = [config.subsampling_num_query]

        self.num_conv_layers = config.subsampling_num_conv_layers
        hidden_dim = config.subsampling_hidden_dim
        if self.num_conv_layers == 0:
            self.conv = nn.Conv1d(in_channels=1, out_channels=hidden_dim - 2, kernel_size=5, stride=1, padding=2)
            self.mask_pool = None
        else:
            self.conv_layers = []
            self.pool_layers = []
            current_channels = 1
            output_channels = hidden_dim - 2

            for i in range(self.num_conv_layers):
                if i == 0:
                    in_c = current_channels
                    out_c = output_channels
                else:
                    in_c = output_channels
                    out_c = output_channels

                self.conv_layers.append(nn.Conv1d(in_c, out_c, kernel_size=5, stride=2, padding=2))
                self.conv_layers.append(nn.ReLU(inplace=True))
                self.pool_layers.append(nn.MaxPool1d(kernel_size=5, stride=2, padding=2))

            self.conv = nn.Sequential(*self.conv_layers)
            self.mask_pool = nn.Sequential(*self.pool_layers)

        if self.selfatt:
            num_encoder_layers = 1
            encoder_layers = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=config.subsampling_nhead)
            self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_encoder_layers)
            self.mrqformer = None
        else:
            self.mrqformer = MRQFormer(
                num_query_list=self.num_query_list,
                hidden_size=hidden_dim,
                num_layers=6,
                cross_attention_freq=2,
                output_size=hidden_dim,
            )
            self.transformer_encoder = None

        self.pos_encoder = InternS2PreviewTimeSeriesFixPositionalEncoding(d_model=hidden_dim)

        num_channel_encoder_layers = 1
        channel_encoder_layers = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=32)
        self.channel_encoder = nn.TransformerEncoder(channel_encoder_layers, num_channel_encoder_layers)
        self.fuze_proj = nn.Linear(in_features=hidden_dim, out_features=256)

    def forward(self, inputs, mask=None, subrate=None):
        if mask is None:
            mask = torch.ones(inputs.shape).to(inputs.device)
        features, feature_lens, learnable_lens, strides, raw_outputs = self.forward_patch(inputs, mask, subrate)

        outputs = features
        output_lens = feature_lens

        return outputs, output_lens, learnable_lens, strides, raw_outputs

    def forward_patch(self, inputs, mask=None, subrate=None):

        pred_strides = None

        patched_outputs = []
        raw_outputs = []

        with torch.set_grad_enabled(self.training):
            seq = inputs  # (B, L, C)

            masked_seq = seq.masked_fill(mask == 0, float("nan"))

            mean = seq.mean(dim=1, keepdim=True)
            std = seq.std(dim=1, keepdim=True)

            mean_expand = mean.expand(seq.size())
            abs_mean_plus_1 = torch.abs(mean) + 1.0  # (B, 1, C), 所有值 ≥ 1
            log_abs_plus_1 = torch.log(abs_mean_plus_1)  # (B, 1, C), ≥ 0
            sign_mean = torch.sign(mean)  # (B, 1, C), 取值 -1, 0, +1
            mean_feat = sign_mean * log_abs_plus_1  # (B, 1, C), 保留原符号的 log 值

            std_expand = std.expand(seq.size())
            std_feat = torch.log(std + 1e-7)

            seq = (seq - mean) / (std + 1e-7)

            P_max = torch.tensor([max(self.patch_list)]).to(seq.device)
            Q_max = torch.tensor([max(self.num_query_list)]).to(seq.device)
            if subrate is None:
                subrate = P_max / Q_max
            st = subrate * Q_max
            ps = torch.ceil(st)
            le = torch.tensor(inputs.shape[1]).to(seq.device)

            output_len = torch.ceil((le - ps) / st + 1)
            pad_len = (torch.ceil((output_len - 1) * st + ps - le)).long().item()

            if seq.ndim == 2:
                seq = seq.unsqueeze(-1)
            if not isinstance(pad_len, int):
                print(f"⚠️ [DEBUG] pad_len is NOT an int! value={pad_len}, type={type(pad_len)}")
            seq = torch.nn.functional.pad(seq, (0, 0, 0, pad_len, 0, 0), "constant", 0)

            padded_mask = torch.nn.functional.pad(mask, (0, 0, 0, pad_len, 0, 0), "constant", 0)  # (B, L, C)

            B, L, C = seq.shape
            seq = seq.permute(0, 2, 1)  # (B, C, L)
            seq = seq.reshape(B * C, 1, L)  # (B*C, 1, L)
            if self.num_conv_layers == 0:
                seq = nn.functional.relu(self.conv(seq))  # (B*C, D, L)
            else:
                seq = self.conv(seq)  # (B*C, D, L)
            L = seq.shape[-1]
            seq = seq.reshape(B, C, -1, L)
            seq = seq.permute(0, 3, 1, 2)  # (B, L, C, D)

            mean_feat = mean_feat.expand([B, L, C]).unsqueeze(-1)
            std_feat = std_feat.expand([B, L, C]).unsqueeze(-1)

            seq = torch.cat([seq, mean_feat, std_feat], dim=-1)  # (B, L, C, D)

            D = seq.shape[-1]

            if self.mask_pool is not None:
                padded_mask = padded_mask.permute(0, 2, 1)
                padded_mask = self.mask_pool(padded_mask)
                padded_mask = padded_mask.permute(0, 2, 1)

            idx = torch.arange(L).to(seq.device)

            output_list = []
            for res_idx, raw_patch in enumerate(self.patch_list):
                num_query = self.num_query_list[res_idx]
                stride = subrate * num_query
                patch = torch.ceil(stride)
                patch_output_len = torch.floor(output_len * st / (stride * (2**self.num_conv_layers)))

                indices = (
                    torch.floor(torch.arange(0, patch_output_len.item()).to(seq.device) * stride).unsqueeze(1)
                    + torch.arange(torch.ceil(patch)).to(seq.device)
                ).long()
                patched = seq[:, indices, :, :]  # (B, N, P, C, D)
                _, N, P, _, _ = patched.shape
                patched = patched.permute(1, 2, 0, 3, 4)
                patched = patched.reshape(N, P, -1)  # (N, P, B*C*D)
                patched = patched.reshape(N, P, B, C, D)

                patched = patched.permute(2, 0, 1, 3, 4)  # (B, N, P, C, D)
                patched = patched.reshape(B * N, P, C, D)

                patched_mask = padded_mask[:, indices, :]  # (B, N, P, C)
                patched_mask = patched_mask.reshape(B * N, P, C)

                raw_output = self.forward_encoder(patched, patched_mask, res_idx)
                if self.selfatt:
                    output = raw_output.mean(dim=1)  # B*N, D
                else:
                    output = raw_output.reshape(-1, D)  # B*N*Q, D
                output_list.append(output)

            outputs = torch.cat(output_list, dim=1)
            outputs = self.fuze_proj(outputs)
            _, D_out = outputs.shape

        outputs = outputs.reshape(B, -1, D_out)
        raw_outputs = None

        if self.selfatt:
            output_lens = torch.ceil((padded_mask[:, :, 0].sum(dim=-1) - ps) / st + 1)
        else:
            output_lens = torch.ceil((padded_mask[:, :, 0].sum(dim=-1) - ps) / st + 1) * Q_max

        learnable_lens = None

        return outputs, output_lens, learnable_lens, pred_strides, raw_outputs

    def forward_encoder(self, inputs, mask, res_idx=0):

        x = inputs  # (B*N, P, C, D)
        num_patch, patch_len, C, D = x.shape
        x = x.permute(1, 0, 2, 3)  # (P, B*N, C, D)
        x = x.reshape(patch_len, num_patch * C, -1)  # [P, B*N*C, D]

        x = self.pos_encoder(x)  # [P, B*C, D]

        mask = mask.permute(0, 2, 1)  # (B*N, C, P)
        mask = mask.reshape(num_patch * C, patch_len)  # (B*N*C, P)
        bool_mask = mask == 0  # (B*N*C, P)

        chunk_size = 4096 * 8
        all_outputs = []

        for i in range(0, x.size(1), chunk_size):
            x_chunk = x[:, i : i + chunk_size, :].contiguous()  # [P, chunk_len, D]
            chunk_mask = mask[i : i + chunk_size]  # (chunk_len, P)
            chunk_bool_mask = bool_mask[i : i + chunk_size]  # (chunk_len, P)
            if self.selfatt:
                x_chunk = self.transformer_encoder(x_chunk, src_key_padding_mask=chunk_bool_mask)  # [P, chunk_len, D]
            else:
                x_chunk = x_chunk.permute(1, 0, 2)
                x_chunk = self.mrqformer(x_chunk, attention_mask=chunk_mask, res_idx=res_idx)
                x_chunk = x_chunk.permute(1, 0, 2)  # [Q, chunk_len, D]
            all_outputs.append(x_chunk)

        x = torch.cat(all_outputs, dim=1)  # [P, B*N*C, D] / [Q, B*N*C, D]

        x = x.reshape(-1, num_patch, C, D)  # [P, B*N, C, D] /# [Q, B*N, C, D]
        x = x.permute(2, 1, 0, 3)  # [C, B, P, D]
        x = x.reshape(C, -1, D)  # [C, B*N*P, D]

        all_outputs = []

        for i in range(0, x.size(1), chunk_size):
            x_chunk = x[:, i : i + chunk_size, :].contiguous()  # [C, chunk_len, D]
            x_chunk = self.channel_encoder(x_chunk)
            x_chunk = x_chunk.mean(0)  # [chunk_len, D]
            all_outputs.append(x_chunk)

        x = torch.cat(all_outputs, dim=0)  # [B*N*P, D]
        x = x.reshape(num_patch, -1, D)  # [B*N, P, D]

        return x


class InternS2PreviewTimeSeriesProjector(nn.Module):
    def __init__(self, config: InternS2PreviewTimeSeriesConfig):
        super().__init__()
        self.layer_norm = nn.LayerNorm(config.ts_hidden_dim)
        self.linear_1 = nn.Linear(config.ts_hidden_dim, config.out_hidden_size)
        self.act = ACT2FN[config.activation_function]
        self.linear_2 = nn.Linear(config.out_hidden_size, config.out_hidden_size)

    def forward(self, ts_features):
        hidden_states = self.layer_norm(ts_features)
        hidden_states = self.linear_1(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.linear_2(hidden_states)
        return hidden_states


@auto_docstring
class InternS2PreviewTimeSeriesModel(PreTrainedModel):
    main_input_name = "time_series_signals"
    _supports_flash_attn_2 = False
    config_class = InternS2PreviewTimeSeriesConfig
    _no_split_modules = ["InternS2PreviewTimeSeriesEncoderLayer"]

    def __init__(self, config: InternS2PreviewTimeSeriesConfig):
        super().__init__(config)
        self.config = config
        self.encoder_embed = TotalFixlenSingleResChunkQformerSubsampling(config)
        self.encoder = InternS2PreviewTimeSeriesEncoder(config)
        self.projector = InternS2PreviewTimeSeriesProjector(config)

    def get_input_embeddings(self):
        return self.encoder_embed

    def make_pad_mask(self, lengths: torch.Tensor) -> torch.Tensor:
        r"""
        lengths (`torch.Tensor` of shape `(batch_size,)`):
            A 1-D tensor containing sentence lengths.

        Returns:
        A 2-D bool tensor, where masked positions are filled with `True` and non-masked positions are filled with `False`.
        Example:
        ```python
        >>> lengths = torch.tensor([1, 3, 2, 5])
        >>> self.make_pad_mask(lengths)
        tensor([[False,  True,  True,  True,  True],
                [False, False, False,  True,  True],
                [False, False,  True,  True,  True],
                [False, False, False, False, False]])
        ```
        """
        assert lengths.ndim == 1, lengths.ndim
        max_len = lengths.max()
        n = lengths.size(0)
        seq_range = torch.arange(0, max_len, device=lengths.device)
        expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)
        return expaned_lengths >= lengths.unsqueeze(-1)

    def concat_chunks_simple(
        self,
        chunks: torch.Tensor,  # (N, T_chunk, C)
        chunk_lengths: torch.Tensor = None,  # (N,) 每个 chunk 的有效长度
    ) -> torch.Tensor:
        """
        简单首尾相接（无 overlap）
        """
        N, T_chunk, C = chunks.shape

        if chunk_lengths is None:
            signal = chunks.reshape(-1, C)  # (N*T_chunk, C)
            signal_len = N * T_chunk
        else:
            signal_list = []
            signal_len = 0
            for i in range(N):
                valid_len = int(chunk_lengths[i].item())
                signal_list.append(chunks[i, :valid_len, :])  # (T_i, C)
                signal_len = signal_len + valid_len
            signal = torch.cat(signal_list, dim=0)  # (sum(T_i), C)

        return signal, signal_len

    def chunk_tensor(self, tensor: torch.Tensor, return_indices: bool = False):

        T = tensor.shape[0]
        chunk_size = self.config.chunk_size
        step = self.config.chunk_step

        chunks = []
        indices = [] if return_indices else None
        masks = []

        if chunk_size > 0:
            start = 0
            while start < T:
                end = min(start + chunk_size, T)
                chunk = tensor[start:end, :]
                chunks.append(chunk)
                mask = torch.zeros(chunk.shape)
                mask[: end - start, :] = 1
                masks.append(mask)
                if return_indices:
                    indices.append((start, end))

                if end >= T:
                    break

                start += step

            output = nn.utils.rnn.pad_sequence(chunks, batch_first=True)
            output_mask = nn.utils.rnn.pad_sequence(masks, batch_first=True).to(output.device)
        else:
            output = tensor.unsqueeze(0)
            output_mask = torch.ones(output.shape).to(tensor.device)

        if return_indices:
            return output, indices, output_mask
        return output, output_mask

    def forward(
        self,
        time_series_signals: torch.FloatTensor = None,
        ts_lens: torch.Tensor = None,
        sr: torch.Tensor = None,
        channels: torch.LongTensor | None = None,
        output_hidden_states: bool = None,
        return_dict: bool = None,
        time_series_embeds: torch.FloatTensor = None,
    ):
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if time_series_signals is None and time_series_embeds is None:
            raise ValueError("You have to specify time_series_signals or time_series_embeds")
        if channels is None:
            raise ValueError("channels must be specified")

        if isinstance(time_series_signals, list):
            batch_iter = [time_series_signals[b][: ts_lens[b], : channels[b]] for b in range(len(time_series_signals))]
        elif time_series_signals.ndim == 3:
            batch_iter = [
                time_series_signals[b, : ts_lens[b], : channels[b]] for b in range(time_series_signals.shape[0])
            ]
        else:
            raise ValueError(f"wrong time_series_signals size: {time_series_signals[0].shape}")

        outputs = []
        output_lens = []
        output_chunks = []
        for b, seq in enumerate(batch_iter):
            chunks, mask = self.chunk_tensor(seq)

            subrate = max(ts_lens[b] / 500, torch.tensor(1.0))
            chunk_output, chunk_lens, learnable_lens, strides, raw_feature = self.encoder_embed(chunks, mask, subrate)

            outputs.append(chunk_output)
            output_lens.append(chunk_lens)
            output_chunks.append(chunk_output.shape[0])
        x_lens = torch.cat(output_lens, dim=0).long()
        max_len = x_lens.max().item()

        padded_list = []
        for t in outputs:
            _, T, _ = t.shape  # (B, T, D)
            pad_T = int(max_len - T)

            padded = nn.functional.pad(t, (0, 0, 0, pad_T, 0, 0), mode="constant", value=0)
            padded_list.append(padded)
        x = torch.cat(padded_list, dim=0)
        if x_lens.ndim == 0:
            x_lens = x_lens.unsqueeze(-1)

        src_key_padding_mask = self.make_pad_mask(x_lens)
        encoder_out = self.encoder(x, x_lens, src_key_padding_mask, output_hidden_states=True, return_dict=return_dict)
        encoder_out = encoder_out.last_hidden_state
        encoder_out_lens = (x_lens + 1) // 2

        features = []
        feature_lens = []
        start = 0

        for i, count in enumerate(output_chunks):
            end = start + count
            signal_chunks = encoder_out[start:end]  # (N_i, T_chunk, C)

            feature, feature_len = self.concat_chunks_simple(signal_chunks, encoder_out_lens[start:end])

            features.append(feature)
            start = end
            feature_lens.append(feature_len)

        encoder_out = nn.utils.rnn.pad_sequence(features, batch_first=True)
        encoder_out_lens = torch.tensor(feature_lens).to(encoder_out.device)

        src_key_padding_mask = self.make_pad_mask(encoder_out_lens)

        ts_pad_mask = src_key_padding_mask
        ts_embeds = self.projector(encoder_out)

        return ts_embeds, ts_pad_mask, encoder_out


class RMSNorm(nn.Module):
    """RMS normalization."""

    def __init__(self, num_features: int, *, epsilon: float = 1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.zeros(num_features))
        self.num_features = num_features
        self.epsilon = epsilon

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        var = torch.mean(torch.square(inputs), dim=-1, keepdim=True)
        normed_inputs = inputs * torch.rsqrt(var + self.epsilon)
        normed_inputs = normed_inputs * self.scale
        return normed_inputs


class ResidualBlock(nn.Module):
    """Residual block with two linear layers and a linear residual connection."""

    def __init__(self, input_dims, hidden_dims, output_dims, use_bias, activation="swish"):
        super().__init__()
        self.hidden_layer = nn.Linear(
            in_features=input_dims,
            out_features=hidden_dims,
            bias=use_bias,
        )
        self.output_layer = nn.Linear(
            in_features=hidden_dims,
            out_features=output_dims,
            bias=use_bias,
        )
        self.residual_layer = nn.Linear(
            in_features=input_dims,
            out_features=output_dims,
            bias=use_bias,
        )
        if activation == "relu":
            self.activation = nn.ReLU()
        elif activation == "swish":
            self.activation = nn.SiLU()
        elif activation == "none":
            self.activation = nn.Identity()
        else:
            raise ValueError(f"Activation: {activation} not supported.")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.output_layer(self.activation(self.hidden_layer(x))) + self.residual_layer(x)


@dataclass(frozen=False)
class DecodeCache:
    """Cache for decoding."""

    next_index: torch.Tensor
    num_masked: torch.Tensor
    key: torch.Tensor
    value: torch.Tensor
    kv_mask: torch.Tensor | None = None  # (B, cache_size), True=padded/masked


class RotaryPositionalEmbedding(nn.Module):
    """Rotary positional embedding."""

    def __init__(
        self,
        embedding_dims: int,
        min_timescale: float = 1.0,
        max_timescale: float = 10000.0,
    ):
        super().__init__()
        self.embedding_dims = embedding_dims
        self.min_timescale = min_timescale
        self.max_timescale = max_timescale

    def forward(
        self,
        inputs: torch.Tensor,
        position: torch.Tensor | None = None,
    ):
        """Generates a JTensor of sinusoids with different frequencies."""
        if self.embedding_dims != inputs.shape[-1]:
            raise ValueError(
                "The embedding dims of the rotary position embeddingmust match the hidden dimension of the inputs."
            )
        half_embedding_dim = self.embedding_dims // 2
        fraction = 2 * torch.arange(0, half_embedding_dim, device=inputs.device) / self.embedding_dims
        timescale = (self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction).to(inputs.device)
        if position is None:
            seq_length = inputs.shape[1]
            position = torch.arange(seq_length, dtype=torch.float32, device=inputs.device)[None, :]

        if len(inputs.shape) == 4:
            position = position[..., None, None]
            timescale = timescale[None, None, None, :]
        elif len(inputs.shape) == 3:
            position = position[..., None]
            timescale = timescale[None, None, :]
        else:
            raise ValueError("Inputs must be of rank 3 or 4.")

        sinusoid_inp = position / timescale
        sin = torch.sin(sinusoid_inp)
        cos = torch.cos(sinusoid_inp)
        first_half, second_half = torch.chunk(inputs, 2, dim=-1)
        first_part = first_half * cos - second_half * sin
        second_part = second_half * cos + first_half * sin
        return torch.cat([first_part, second_part], dim=-1)


class PerDimScale(nn.Module):
    """Per-dimension scaling."""

    def __init__(self, num_dims: int):
        super().__init__()
        self.num_dims = num_dims
        self.per_dim_scale = nn.Parameter(torch.zeros(num_dims))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        scale_factor = 1.442695041 / math.sqrt(self.num_dims) * nn.functional.softplus(self.per_dim_scale)
        return x * scale_factor


def _torch_dot_product_attention(query, key, value, mask=None):
    """Same (unscaled) attention as _dot_product_attention, fused kernel."""
    safe_mask = mask
    fully_masked_rows = None
    if mask is not None:
        # F.scaled_dot_product_attention can emit NaNs/NaN-grads when a query row is
        # fully masked. Keep the fused kernel by opening a single dummy key for
        # those rows, then zero their outputs.
        fully_masked_rows = ~mask.any(dim=-1, keepdim=True)
        if fully_masked_rows.any():
            dummy_key_mask = torch.zeros_like(mask)
            dummy_key_mask[..., 0] = True
            safe_mask = mask | (fully_masked_rows & dummy_key_mask)

    attention_dtype = value.dtype
    if query.dtype != attention_dtype:
        query = query.to(attention_dtype)
    if key.dtype != attention_dtype:
        key = key.to(attention_dtype)

    # 1. Permute inputs from (B, L, H, D) to the expected (B, H, L, D)
    query = query.permute(0, 2, 1, 3)
    key = key.permute(0, 2, 1, 3)
    value = value.permute(0, 2, 1, 3)

    # 2. Fused attention kernel with scale=1.0 (disable 1/sqrt(d_k) scaling).
    output = nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=safe_mask, scale=1.0)

    if fully_masked_rows is not None and fully_masked_rows.any():
        output = output.masked_fill(fully_masked_rows, 0.0)

    # 3. Permute back to (B, L, H, D)
    output = output.permute(0, 2, 1, 3)
    return output


class MultiHeadAttention(nn.Module):
    """Multi-head attention."""

    def __init__(
        self,
        num_heads: int,
        in_features: int,
        *,
        use_per_dim_scale: bool = True,
        use_rotary_position_embeddings: bool = True,
        use_bias: bool = False,
        attention_fn: Callable[..., torch.Tensor] = _torch_dot_product_attention,
        qk_norm: str = "rms",
        fuse_qkv: bool = False,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.in_features = in_features
        self.head_dim = in_features // num_heads
        self.use_bias = use_bias
        self.attention_fn = attention_fn
        self.qk_norm = qk_norm
        self.fuse_qkv = fuse_qkv

        if self.in_features % self.num_heads != 0:
            raise ValueError(
                f"Memory dimension ({self.in_features}) must be divisible by 'num_heads' heads ({self.num_heads})."
            )

        if self.fuse_qkv:
            self.qkv_proj = nn.Linear(self.in_features, 3 * self.in_features, bias=use_bias)
        else:
            self.query = nn.Linear(self.in_features, self.in_features, bias=use_bias)
            self.key = nn.Linear(self.in_features, self.in_features, bias=use_bias)
            self.value = nn.Linear(self.in_features, self.in_features, bias=use_bias)
        self.out = nn.Linear(self.in_features, self.in_features, bias=use_bias)

        if self.qk_norm == "rms":
            self.query_ln = RMSNorm(self.head_dim)
            self.key_ln = RMSNorm(self.head_dim)
        else:
            self.query_ln = nn.Identity()
            self.key_ln = nn.Identity()

        self.use_rotary_position_embeddings = use_rotary_position_embeddings
        if self.use_rotary_position_embeddings:
            self.rotary_position_embedding = RotaryPositionalEmbedding(
                embedding_dims=self.head_dim,
            )

        self.use_per_dim_scale = use_per_dim_scale
        if use_per_dim_scale:
            self.per_dim_scale = PerDimScale(num_dims=self.head_dim)

    def make_attn_mask(
        self,
        query_length: int,
        num_all_masked_kv: torch.Tensor,
        query_index_offset: torch.Tensor | None = None,
        kv_length: int = 0,
    ) -> torch.Tensor:
        """Makes attention mask."""
        if kv_length == 0:
            kv_length = query_length

        q_index = torch.arange(query_length, device=num_all_masked_kv.device)[None, None, :, None]
        if query_index_offset is not None:
            q_index = q_index + query_index_offset[:, None, None, None]
        kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[None, None, None, :]
        return torch.logical_and(
            q_index >= kv_index,
            kv_index >= num_all_masked_kv[:, None, None, None],
        )

    def forward(
        self,
        inputs_q: torch.Tensor,
        *,
        decode_cache: DecodeCache | None = None,
        patch_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, DecodeCache | None]:
        b, n_patches, _ = inputs_q.shape
        if patch_mask is None:
            patch_mask = torch.zeros(b, n_patches, dtype=torch.bool, device=inputs_q.device)

        if self.fuse_qkv:
            qkv = self.qkv_proj(inputs_q)
            query, key, value = torch.chunk(qkv, 3, dim=-1)
            query = query.view(b, n_patches, self.num_heads, self.head_dim)
            key = key.view(b, n_patches, self.num_heads, self.head_dim)
            value = value.view(b, n_patches, self.num_heads, self.head_dim)
        else:
            query = self.query(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
            key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
            value = self.value(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)

        if decode_cache is None:
            # Count only LEADING masked patches (not total). make_attn_mask assumes
            # all masked positions are a contiguous left-padding block.
            is_valid = ~patch_mask
            has_valid = is_valid.any(dim=-1)
            first_valid = torch.argmax(is_valid.to(torch.int32), dim=-1)
            seq_len = torch.tensor(patch_mask.shape[-1], dtype=torch.int32, device=patch_mask.device)
            num_masked = torch.where(has_valid, first_valid, seq_len)
            next_index = torch.zeros_like(num_masked, dtype=torch.int32)
        else:
            is_valid = ~patch_mask
            has_valid = is_valid.any(dim=-1)
            first_valid = torch.argmax(is_valid.to(torch.int32), dim=-1)
            seq_len = torch.tensor(patch_mask.shape[-1], dtype=torch.int32, device=patch_mask.device)
            leading_masked = torch.where(has_valid, first_valid, seq_len)
            num_masked = leading_masked + decode_cache.num_masked
            next_index = decode_cache.next_index.clone()

        if self.use_rotary_position_embeddings:
            position = (
                torch.arange(n_patches, device=inputs_q.device)[None, :] + next_index[:, None] - num_masked[:, None]
            )
            query = self.rotary_position_embedding(query, position)
            key = self.rotary_position_embedding(key, position)

        query = self.query_ln(query)
        key = self.key_ln(key)

        if self.use_per_dim_scale:
            query = self.per_dim_scale(query)

        if decode_cache is not None:
            _, decode_cache_size, _, _ = decode_cache.value.shape

            start = decode_cache.next_index[0]
            end = start + n_patches

            decode_cache.key[:, start:end] = key
            decode_cache.value[:, start:end] = value

            if decode_cache.kv_mask is None:
                decode_cache.kv_mask = torch.ones(
                    b,
                    decode_cache_size,
                    dtype=torch.bool,
                    device=patch_mask.device,
                )
            decode_cache.kv_mask[:, start:end] = patch_mask

            key = decode_cache.key
            value = decode_cache.value
            decode_cache.next_index += n_patches
            decode_cache.num_masked = num_masked
            attn_mask = self.make_attn_mask(
                query_length=n_patches,
                num_all_masked_kv=num_masked,
                query_index_offset=next_index,
                kv_length=decode_cache_size,
            )
            kv_valid = ~decode_cache.kv_mask  # (B, decode_cache_size)
            attn_mask = attn_mask & kv_valid[:, None, None, :]
        else:
            attn_mask = self.make_attn_mask(query_length=n_patches, num_all_masked_kv=num_masked)
            kv_valid = ~patch_mask  # (B, n_patches), True=valid
            attn_mask = attn_mask & kv_valid[:, None, None, :]

        x = self.attention_fn(query, key, value, mask=attn_mask)

        x = x.reshape(b, n_patches, self.in_features)
        out = self.out(x)
        return out, decode_cache


class MultiHeadCrossAttention(nn.Module):
    """Multi-head cross-attention from a query stream into a static KV context.

    Q comes from the transformer's residual stream; K/V come from an external
    context tensor (e.g. Q-former output). No RoPE, no decode_cache: the KV
    context is treated as a static set of tokens at every layer / decode step.
    """

    def __init__(
        self,
        num_heads: int,
        in_features: int,
        *,
        kv_features: int | None = None,
        use_bias: bool = False,
        attention_fn: Callable[..., torch.Tensor] = _torch_dot_product_attention,
        qk_norm: str = "rms",
    ):
        super().__init__()
        self.num_heads = num_heads
        self.in_features = in_features
        self.kv_features = kv_features if kv_features is not None else in_features
        self.head_dim = in_features // num_heads
        self.attention_fn = attention_fn

        if self.in_features % self.num_heads != 0:
            raise ValueError(
                f"Memory dimension ({self.in_features}) must be divisible by 'num_heads' heads ({self.num_heads})."
            )

        self.q_proj = nn.Linear(self.in_features, self.in_features, bias=use_bias)
        self.k_proj = nn.Linear(self.kv_features, self.in_features, bias=use_bias)
        self.v_proj = nn.Linear(self.kv_features, self.in_features, bias=use_bias)
        self.out = nn.Linear(self.in_features, self.in_features, bias=use_bias)

        if qk_norm == "rms":
            self.query_ln = RMSNorm(self.head_dim)
            self.key_ln = RMSNorm(self.head_dim)
        else:
            self.query_ln = nn.Identity()
            self.key_ln = nn.Identity()

        self.per_dim_scale = PerDimScale(num_dims=self.head_dim)

    def forward(
        self,
        inputs_q: torch.Tensor,
        *,
        kv: torch.Tensor,
        kv_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        b, n_q, _ = inputs_q.shape
        n_kv = kv.shape[1]

        if kv.dtype != inputs_q.dtype:
            kv = kv.to(dtype=inputs_q.dtype)

        query = self.q_proj(inputs_q).view(b, n_q, self.num_heads, self.head_dim)
        key = self.k_proj(kv).view(b, n_kv, self.num_heads, self.head_dim)
        value = self.v_proj(kv).view(b, n_kv, self.num_heads, self.head_dim)

        query = self.query_ln(query)
        key = self.key_ln(key)
        query = self.per_dim_scale(query)

        if kv_mask is None:
            attn_mask = torch.ones(
                b,
                1,
                n_q,
                n_kv,
                dtype=torch.bool,
                device=inputs_q.device,
            )
        else:
            kv_valid = ~kv_mask  # True = valid
            attn_mask = kv_valid[:, None, None, :].expand(b, 1, n_q, n_kv)

        x = self.attention_fn(query, key, value, mask=attn_mask)
        x = x.reshape(b, n_q, self.in_features)
        return self.out(x)


class Transformer(nn.Module):
    """Classic Transformer used in Forecaster."""

    def __init__(
        self,
        model_dims: int,
        hidden_dims: int,
        num_heads: int,
        *,
        attention_norm: str,
        feedforward_norm: str,
        qk_norm: str,
        use_bias: bool,
        use_rotary_position_embeddings: bool,
        ff_activation: str,
        fuse_qkv: bool,
        use_cross_attn: bool = False,
        cross_attn_kv_dim: int = 0,
        use_cross_attn_gate: bool = True,
    ):
        super().__init__()

        if attention_norm == "rms":
            self.pre_attn_ln = RMSNorm(num_features=model_dims)
            self.post_attn_ln = RMSNorm(num_features=model_dims)
        else:
            raise ValueError(f"Layer norm: {attention_norm} not supported.")

        self.attn = MultiHeadAttention(
            num_heads=num_heads,
            in_features=model_dims,
            use_per_dim_scale=True,
            use_rotary_position_embeddings=use_rotary_position_embeddings,
            qk_norm=qk_norm,
            fuse_qkv=fuse_qkv,
        )

        if use_cross_attn:
            _cross_kv_dim = cross_attn_kv_dim or model_dims
            self.cross_attn_ln = RMSNorm(num_features=model_dims)
            self.cross_attn = MultiHeadCrossAttention(
                num_heads=num_heads,
                in_features=model_dims,
                kv_features=_cross_kv_dim,
                use_bias=use_bias,
                qk_norm=qk_norm,
            )
            if use_cross_attn_gate:
                # tanh(0)=0 → cross-attn contributes zero at init.
                self.cross_attn_gate = nn.Parameter(torch.zeros(1))
            else:
                self.cross_attn_gate = None
        else:
            self.cross_attn_ln = None
            self.cross_attn = None
            self.cross_attn_gate = None

        if feedforward_norm == "rms":
            self.pre_ff_ln = RMSNorm(num_features=model_dims)
            self.post_ff_ln = RMSNorm(num_features=model_dims)
        else:
            raise ValueError(f"Layer norm: {feedforward_norm} not supported.")

        self.ff0 = nn.Linear(
            in_features=model_dims,
            out_features=hidden_dims,
            bias=use_bias,
        )
        self.ff1 = nn.Linear(
            in_features=hidden_dims,
            out_features=model_dims,
            bias=use_bias,
        )
        if ff_activation == "relu":
            self.activation = nn.ReLU()
        elif ff_activation == "swish":
            self.activation = nn.SiLU()
        elif ff_activation == "none":
            self.activation = nn.Identity()
        else:
            raise ValueError(f"Activation: {ff_activation} not supported.")

    def forward(
        self,
        input_embeddings: torch.Tensor,
        patch_mask: torch.Tensor,
        decode_cache: DecodeCache | None = None,
        cross_kv: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, DecodeCache | None]:
        attn_output, decode_cache = self.attn(
            inputs_q=self.pre_attn_ln(input_embeddings),
            decode_cache=decode_cache,
            patch_mask=patch_mask,
        )
        attn_output = self.post_attn_ln(attn_output) + input_embeddings

        if self.cross_attn is not None and cross_kv is not None:
            cross_out = self.cross_attn(
                self.cross_attn_ln(attn_output),
                kv=cross_kv,
            )
            if self.cross_attn_gate is not None:
                cross_out = torch.tanh(self.cross_attn_gate) * cross_out
            attn_output = attn_output + cross_out

        output_embeddings = (
            self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output))))) + attn_output
        )
        return output_embeddings, decode_cache


class _QFormerBlock(nn.Module):
    """One Q-former block: self-attn over queries, cross-attn into source, FFN."""

    def __init__(
        self,
        hidden_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        ffn_mult: float = 4.0,
    ):
        super().__init__()
        self.self_attn_ln = nn.LayerNorm(hidden_dim)
        self.self_attn = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True,
        )

        self.cross_attn_ln_q = nn.LayerNorm(hidden_dim)
        self.cross_attn_ln_kv = nn.LayerNorm(hidden_dim)
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True,
        )

        ffn_hidden = int(round(hidden_dim * ffn_mult))
        self.ffn_ln = nn.LayerNorm(hidden_dim)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_dim, ffn_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_hidden, hidden_dim),
            nn.Dropout(dropout),
        )

    def forward(
        self,
        queries: torch.Tensor,
        kv: torch.Tensor,
        kv_key_padding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        q_norm = self.self_attn_ln(queries)
        sa_out, _ = self.self_attn(q_norm, q_norm, q_norm, need_weights=False)
        queries = queries + sa_out

        q_norm = self.cross_attn_ln_q(queries)
        kv_norm = self.cross_attn_ln_kv(kv)
        ca_out, _ = self.cross_attn(
            q_norm,
            kv_norm,
            kv_norm,
            key_padding_mask=kv_key_padding_mask,
            need_weights=False,
        )
        queries = queries + ca_out

        queries = queries + self.ffn(self.ffn_ln(queries))
        return queries


class QFormer(nn.Module):
    """BLIP-2-style Q-former: learned queries cross-attend to a source sequence.

    Args:
      in_dim: Source feature dim (e.g. ts encoder hidden, or LLM hidden).
      out_dim: Output (query) hidden dim. Should match the consumer module dim.
      num_query_tokens: Number of learned query tokens.
      num_heads: Attention heads inside the Q-former.
      num_layers: Number of stacked Q-former blocks.
      dropout: Attention/FFN dropout.
      ffn_mult: FFN expansion ratio.
    """

    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        num_query_tokens: int,
        num_heads: int = 8,
        num_layers: int = 2,
        dropout: float = 0.0,
        ffn_mult: float = 4.0,
    ):
        super().__init__()
        if num_query_tokens <= 0:
            raise ValueError(f"num_query_tokens must be positive, got {num_query_tokens}")
        if num_layers <= 0:
            raise ValueError(f"num_layers must be positive, got {num_layers}")
        if out_dim % num_heads != 0:
            raise ValueError(f"out_dim ({out_dim}) must be divisible by num_heads ({num_heads})")

        self.in_dim = in_dim
        self.out_dim = out_dim
        self.num_query_tokens = num_query_tokens

        self.input_proj = nn.Linear(in_dim, out_dim)
        self.input_ln = nn.LayerNorm(out_dim)

        self.query_tokens = nn.Parameter(torch.zeros(1, num_query_tokens, out_dim))
        nn.init.trunc_normal_(self.query_tokens, std=0.02)

        self.blocks = nn.ModuleList(
            [
                _QFormerBlock(
                    hidden_dim=out_dim,
                    num_heads=num_heads,
                    dropout=dropout,
                    ffn_mult=ffn_mult,
                )
                for _ in range(num_layers)
            ]
        )

        self.output_ln = nn.LayerNorm(out_dim)

    def forward(
        self,
        src: torch.Tensor,
        src_key_padding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Compress `src` into `num_query_tokens` learned query vectors.

        Args:
          src: (B, T, in_dim) source token embeddings.
          src_key_padding_mask: (B, T) bool, True = padded position (ignored).

        Returns:
          (B, num_query_tokens, out_dim) query outputs.
        """
        if src.dim() != 3:
            raise ValueError(f"src must be (B, T, D), got shape {tuple(src.shape)}")
        if src.size(-1) != self.in_dim:
            raise ValueError(f"src last dim ({src.size(-1)}) does not match in_dim ({self.in_dim})")

        kv = self.input_ln(self.input_proj(src))

        if src_key_padding_mask is not None:
            if src_key_padding_mask.shape != src.shape[:2]:
                raise ValueError(
                    "src_key_padding_mask shape must equal src.shape[:2]:"
                    f" {tuple(src_key_padding_mask.shape)} != {tuple(src.shape[:2])}"
                )
            kpm = src_key_padding_mask.to(dtype=torch.bool, device=kv.device)
            # If a row is fully padded, MHA would NaN. Flip one position to valid.
            all_padded = kpm.all(dim=1)
            if all_padded.any():
                kpm = kpm.clone()
                kpm[all_padded, 0] = False
        else:
            kpm = None

        queries = self.query_tokens.expand(src.size(0), -1, -1).to(dtype=kv.dtype)
        for block in self.blocks:
            queries = block(queries, kv, kv_key_padding_mask=kpm)

        return self.output_ln(queries)


class Aligner(nn.Module):
    """Aligns precomputed LLM / TS-encoder embeddings into Forecaster's cross-attn
    KV space, and predicts the forecast horizon.

    Two per-modality Q-formers compress the raw hidden-state sequences into a
    fixed number of query tokens; their concatenation is the static KV stream the
    Forecaster transformer cross-attends to. The forecast horizon is predicted by
    a separate one-layer Transformer encoder over the original LLM hidden states,
    so it does not share the LLM Q-former used for Forecaster cross-attention.
    """

    def __init__(self, config: "InternS2PreviewTimeSeriesForecasterConfig"):
        super().__init__()
        qformer_kwargs = dict(
            out_dim=config.qformer_hidden_dim,
            num_query_tokens=config.qformer_num_query_tokens,
            num_heads=config.qformer_num_heads,
            num_layers=config.qformer_num_layers,
            dropout=config.qformer_dropout,
        )
        self.ts_qformer = QFormer(in_dim=config.d_ts_encoder, **qformer_kwargs)
        self.llm_qformer = QFormer(in_dim=config.d_llm, **qformer_kwargs)

        if config.use_horizon_head:
            horizon_hidden_dim = config.qformer_hidden_dim
            if horizon_hidden_dim % config.qformer_num_heads != 0:
                raise ValueError(
                    f"horizon_hidden_dim ({horizon_hidden_dim}) must be divisible by qformer_num_heads "
                    f"({config.qformer_num_heads}) for the horizon Transformer encoder."
                )
            self.horizon_input_proj = nn.Linear(config.d_llm, horizon_hidden_dim)
            self.horizon_input_ln = nn.LayerNorm(horizon_hidden_dim)
            self.horizon_encoder = nn.TransformerEncoderLayer(
                d_model=horizon_hidden_dim,
                nhead=config.qformer_num_heads,
                dim_feedforward=4 * horizon_hidden_dim,
                dropout=config.qformer_dropout,
                activation="gelu",
                batch_first=True,
                norm_first=True,
            )
            # Prediction-length head: LayerNorm -> Linear -> SiLU -> Linear(.,1),
            # regressing horizon from pooled horizon-encoder LLM features.
            self.horizon_head = nn.Sequential(
                nn.LayerNorm(horizon_hidden_dim),
                nn.Linear(horizon_hidden_dim, horizon_hidden_dim),
                nn.SiLU(),
                nn.Linear(horizon_hidden_dim, 1),
            )
        else:
            self.horizon_input_proj = None
            self.horizon_input_ln = None
            self.horizon_encoder = None
            self.horizon_head = None

    def forward(
        self,
        llm_embedding_input: torch.Tensor,
        ts_encoder_embedding_input: torch.Tensor,
        llm_embedding_mask: torch.Tensor | None = None,
        ts_encoder_embedding_mask: torch.Tensor | None = None,
    ):
        """Compress both modalities into the Forecaster cross-attention KV stream.

        Args:
          llm_embedding_input: (B, T_llm, d_llm) precomputed LLM hidden states.
          ts_encoder_embedding_input: (B, T_ts, d_ts_encoder) precomputed TS-encoder
            hidden states.
          llm_embedding_mask / ts_encoder_embedding_mask: optional (B, T) bool masks,
            True = valid token. Default: all valid.

        Returns:
          ctx: (B, Q_ts + Q_llm, qformer_hidden_dim) cross-attention KV stream.
        """
        # --- TS-encoder branch ---
        ts_param = next(self.ts_qformer.parameters())
        ts_hidden = ts_encoder_embedding_input.to(device=ts_param.device, dtype=ts_param.dtype)
        if ts_encoder_embedding_mask is None:
            ts_pad = torch.zeros(ts_hidden.shape[:2], dtype=torch.bool, device=ts_hidden.device)
        else:
            ts_pad = (~ts_encoder_embedding_mask.to(dtype=torch.bool)).to(device=ts_hidden.device)
        ts_chunk = self.ts_qformer(ts_hidden, src_key_padding_mask=ts_pad)

        # --- LLM branch ---
        llm_param = next(self.llm_qformer.parameters())
        llm_hidden = llm_embedding_input.to(device=llm_param.device, dtype=llm_param.dtype)
        if llm_embedding_mask is None:
            llm_pad = torch.zeros(llm_hidden.shape[:2], dtype=torch.bool, device=llm_hidden.device)
        else:
            llm_pad = (~llm_embedding_mask.to(dtype=torch.bool)).to(device=llm_hidden.device)
        llm_chunk = self.llm_qformer(llm_hidden, src_key_padding_mask=llm_pad)

        # --- Concatenate: ts tokens first, then llm tokens ---
        ctx = torch.cat([ts_chunk, llm_chunk], dim=1)
        return ctx

    def predict_horizon(
        self,
        llm_embedding_input: torch.Tensor,
        llm_embedding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Predict per-sample forecast horizon in LINEAR step space from LLM hidden states.

        Args:
          llm_embedding_input: (B, T_llm, d_llm) precomputed LLM hidden states.
          llm_embedding_mask: optional (B, T_llm) bool mask, True = valid token.

        Returns:
          (B,) float32 tensor of the predicted horizon (in steps) as float32. Linear (not log)
          so a small regression error stays a small ABSOLUTE step error — log1p/expm1 amplified
          errors at long horizons (e.g. ~0.05 in log ≈ 36 steps at H=720), hurting exact matching.
        """
        if (
            self.horizon_input_proj is None
            or self.horizon_input_ln is None
            or self.horizon_encoder is None
            or self.horizon_head is None
        ):
            raise RuntimeError("horizon_head is not enabled but predict_horizon was called")
        if llm_embedding_input.dim() != 3:
            raise ValueError(
                f"Expected llm_embedding_input with shape (B, T, D), got {tuple(llm_embedding_input.shape)}"
            )
        head_param = next(self.horizon_head.parameters())
        llm_hidden = llm_embedding_input.to(device=head_param.device, dtype=head_param.dtype)

        if llm_embedding_mask is None:
            valid_mask = torch.ones(llm_hidden.shape[:2], dtype=torch.bool, device=llm_hidden.device)
        else:
            valid_mask = llm_embedding_mask.to(device=llm_hidden.device, dtype=torch.bool)
            if valid_mask.shape != llm_hidden.shape[:2]:
                raise ValueError(
                    "llm_embedding_mask shape must equal llm_embedding_input.shape[:2]: "
                    f"{tuple(valid_mask.shape)} != {tuple(llm_hidden.shape[:2])}"
                )
        if not valid_mask.any(dim=1).all():
            raise ValueError("Each llm_embedding_mask row must contain at least one valid token.")

        src_key_padding_mask = ~valid_mask

        hidden = self.horizon_input_ln(self.horizon_input_proj(llm_hidden))
        encoded = self.horizon_encoder(
            hidden,
            src_key_padding_mask=src_key_padding_mask,
        )
        pool_mask = valid_mask.to(dtype=encoded.dtype).unsqueeze(-1)
        denom = pool_mask.sum(dim=1).clamp_min(1.0)
        pooled = (encoded * pool_mask).sum(dim=1) / denom
        return self.horizon_head(pooled).squeeze(-1).to(dtype=torch.float32)

    @staticmethod
    def decode_horizon(
        pred_horizon: torch.Tensor,
        min_h: int = 1,
        max_h: int | None = None,
    ) -> torch.Tensor:
        """Round the head's linear-space horizon prediction to a positive integer tensor."""
        horizon = pred_horizon.round().clamp_min(float(min_h))
        if max_h is not None:
            horizon = horizon.clamp_max(float(max_h))
        return horizon.long()


_TOLERANCE = 1e-6


def revin(
    x: torch.Tensor,
    mu: torch.Tensor,
    sigma: torch.Tensor,
    reverse: bool = False,
):
    """Reversible instance normalization."""
    if len(mu.shape) == len(x.shape) - 1:
        mu = mu[..., None]
        sigma = sigma[..., None]
    elif len(mu.shape) == len(x.shape) - 2:
        mu = mu[..., None, None]
        sigma = sigma[..., None, None]

    if reverse:
        return x * sigma + mu
    else:
        return (x - mu) / torch.where(sigma < _TOLERANCE, 1.0, sigma)


class ForecasterBackbone(nn.Module):
    """Forecaster 2.5 with 200M parameters (cross-attention capable)."""

    # Architecture constants (Forecaster 2.5 200M)
    context_limit = 16384
    _quantiles = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    _num_layers = 20

    # Residual block configs
    _tokenizer_kwargs = dict(input_dims=64, hidden_dims=1280, output_dims=1280, use_bias=True, activation="swish")
    _output_point_kwargs = dict(
        input_dims=1280, hidden_dims=1280, output_dims=1280, use_bias=False, activation="swish"
    )
    _output_quantile_kwargs = dict(
        input_dims=1280, hidden_dims=1280, output_dims=10240, use_bias=False, activation="swish"
    )

    # Transformer layer config
    _xf_kwargs = dict(
        model_dims=1280,
        hidden_dims=1280,
        num_heads=16,
        attention_norm="rms",
        feedforward_norm="rms",
        qk_norm="rms",
        use_bias=False,
        use_rotary_position_embeddings=True,
        ff_activation="swish",
        fuse_qkv=True,
    )

    def __init__(
        self,
        use_cross_attn: bool = False,
        cross_attn_kv_dim: int = 0,
        use_cross_attn_gate: bool = True,
    ):
        super().__init__()

        # Named constants.
        self.p = 32
        self.o = 128
        self.os = 1024
        self.m = self.o // self.p  # 4
        self.x = self._num_layers  # 20
        self.h = self._xf_kwargs["num_heads"]  # 16
        self.md = self._xf_kwargs["model_dims"]  # 1280
        self.hd = self.md // self.h  # 80
        self.q = len(self._quantiles) + 1  # 10
        self.aridx = 5

        self.use_cross_attn = bool(use_cross_attn)

        xf_kwargs = dict(self._xf_kwargs)
        if self.use_cross_attn:
            xf_kwargs["use_cross_attn"] = True
            xf_kwargs["cross_attn_kv_dim"] = int(cross_attn_kv_dim) or self.md
            xf_kwargs["use_cross_attn_gate"] = bool(use_cross_attn_gate)

        # Layers.
        self.tokenizer = ResidualBlock(**self._tokenizer_kwargs)
        self.stacked_xf = nn.ModuleList([Transformer(**xf_kwargs) for _ in range(self.x)])
        self.output_projection_point = ResidualBlock(**self._output_point_kwargs)
        self.output_projection_quantiles = ResidualBlock(**self._output_quantile_kwargs)

    def update_running_stats(
        self,
        n: torch.Tensor,
        mu: torch.Tensor,
        sigma: torch.Tensor,
        x: torch.Tensor,
        mask: torch.Tensor,
    ):
        """Updates the running stats."""
        is_legit = torch.logical_not(mask)
        inc_n = torch.sum(is_legit.to(x.dtype), dim=-1)

        inc_mu_numerator = torch.sum(x * is_legit, dim=-1)
        inc_n_safe = torch.where(inc_n == 0, 1.0, inc_n)
        inc_mu = inc_mu_numerator / inc_n_safe
        inc_mu = torch.where(inc_n == 0, 0.0, inc_mu)

        inc_var_numerator = torch.sum(((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1)
        inc_var = inc_var_numerator / inc_n_safe
        inc_var = torch.where(inc_n == 0, 0.0, inc_var)
        inc_sigma = torch.sqrt(inc_var)

        new_n = n + inc_n
        new_n_safe = torch.where(new_n == 0, 1.0, new_n)

        new_mu = (n * mu + inc_mu * inc_n) / new_n_safe
        new_mu = torch.where(new_n == 0, 0.0, new_mu)

        term1 = n * sigma.pow(2)
        term2 = inc_n * inc_sigma.pow(2)
        term3 = n * (mu - new_mu).pow(2)
        term4 = inc_n * (inc_mu - new_mu).pow(2)

        new_var = (term1 + term2 + term3 + term4) / new_n_safe
        new_var = torch.where(new_n == 0, 0.0, new_var)
        new_sigma = torch.sqrt(torch.clamp(new_var, min=0.0))

        return (w := (new_n, new_mu, new_sigma), w)

    def forward(
        self,
        inputs: torch.Tensor,
        masks: torch.Tensor,
        decode_caches: list | None = None,
        cross_kv: torch.Tensor | None = None,
    ):
        """Forward pass — history-only path with cross-attention injection."""
        if cross_kv is not None and not getattr(self, "use_cross_attn", False):
            raise ValueError(
                "cross_kv was supplied but the Forecaster module was constructed with"
                " use_cross_attn=False; rebuild the module with use_cross_attn=True."
            )

        input_dtype = inputs.dtype
        model_dtype = next(self.parameters()).dtype
        if inputs.dtype != model_dtype:
            inputs = inputs.to(dtype=model_dtype)

        tokenizer_inputs = torch.cat([inputs, masks.to(inputs.dtype)], dim=-1)
        input_embeddings = self.tokenizer(tokenizer_inputs)

        if cross_kv is not None:
            cross_kv = cross_kv.to(device=input_embeddings.device, dtype=input_embeddings.dtype)

        if decode_caches is None:
            decode_caches = [None] * self.x
        new_decode_caches = []

        output_embeddings = input_embeddings
        token_masks = masks[..., -1]

        for i, layer in enumerate(self.stacked_xf):
            output_embeddings, new_cache = layer(
                output_embeddings,
                token_masks,
                decode_caches[i],
                cross_kv=cross_kv,
            )
            new_decode_caches.append(new_cache)

        output_ts = self.output_projection_point(output_embeddings)
        output_quantile_spread = self.output_projection_quantiles(output_embeddings)

        if output_ts.dtype != input_dtype:
            input_embeddings = input_embeddings.to(dtype=input_dtype)
            output_embeddings = output_embeddings.to(dtype=input_dtype)
            output_ts = output_ts.to(dtype=input_dtype)
            output_quantile_spread = output_quantile_spread.to(dtype=input_dtype)

        return (
            input_embeddings,
            output_embeddings,
            output_ts,
            output_quantile_spread,
        ), new_decode_caches

    def decode(
        self,
        horizon: int,
        inputs,
        masks,
        enable_grad: bool = False,
        cross_kv: torch.Tensor | None = None,
    ):
        """Decodes the time series."""
        with torch.set_grad_enabled(enable_grad):
            batch_size, context = inputs.shape[0], inputs.shape[1]
            num_decode_steps = (horizon - 1) // self.o
            num_input_patches = context // self.p
            decode_cache_size = num_input_patches + num_decode_steps * self.m
            cache_dtype = next(self.parameters()).dtype

            # Prefill
            patched_inputs = torch.reshape(inputs, (batch_size, -1, self.p))
            patched_masks = torch.reshape(masks, (batch_size, -1, self.p))

            # running stats
            n = torch.zeros(batch_size, device=inputs.device)
            mu = torch.zeros(batch_size, device=inputs.device)
            sigma = torch.zeros(batch_size, device=inputs.device)
            patch_mu = []
            patch_sigma = []
            for i in range(num_input_patches):
                (n, mu, sigma), _ = self.update_running_stats(n, mu, sigma, patched_inputs[:, i], patched_masks[:, i])
                patch_mu.append(mu)
                patch_sigma.append(sigma)
            last_n, last_mu, last_sigma = n, mu, sigma
            context_mu = torch.stack(patch_mu, dim=1)
            context_sigma = torch.stack(patch_sigma, dim=1)

            decode_caches = [
                DecodeCache(
                    next_index=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device),
                    num_masked=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device),
                    key=torch.zeros(
                        batch_size,
                        decode_cache_size,
                        self.h,
                        self.hd,
                        device=inputs.device,
                        dtype=cache_dtype,
                    ),
                    value=torch.zeros(
                        batch_size,
                        decode_cache_size,
                        self.h,
                        self.hd,
                        device=inputs.device,
                        dtype=cache_dtype,
                    ),
                )
                for _ in range(self.x)
            ]

            normed_inputs = revin(patched_inputs, context_mu, context_sigma, reverse=False)
            normed_inputs = torch.where(patched_masks, 0.0, normed_inputs)
            (_, _, normed_outputs, normed_quantile_spread), decode_caches = self(
                normed_inputs,
                patched_masks,
                decode_caches,
                cross_kv=cross_kv,
            )
            renormed_outputs = torch.reshape(
                revin(normed_outputs, context_mu, context_sigma, reverse=True),
                (batch_size, -1, self.o, self.q),
            )
            renormed_quantile_spread = torch.reshape(
                revin(normed_quantile_spread, context_mu, context_sigma, reverse=True),
                (batch_size, -1, self.os, self.q),
            )[:, -1, ...]

            # Autoregressive decode
            ar_outputs = []
            last_renormed_output = renormed_outputs[:, -1, :, self.aridx]

            for _ in range(num_decode_steps):
                new_patched_input = torch.reshape(last_renormed_output, (batch_size, self.m, self.p))
                new_mask = torch.zeros_like(new_patched_input, dtype=torch.bool)

                n, mu, sigma = last_n, last_mu, last_sigma
                new_mus, new_sigmas = [], []
                for i in range(self.m):
                    (n, mu, sigma), _ = self.update_running_stats(
                        n, mu, sigma, new_patched_input[:, i], new_mask[:, i]
                    )
                    new_mus.append(mu)
                    new_sigmas.append(sigma)
                last_n, last_mu, last_sigma = n, mu, sigma
                new_mu = torch.stack(new_mus, dim=1)
                new_sigma = torch.stack(new_sigmas, dim=1)

                new_normed_input = revin(new_patched_input, new_mu, new_sigma, reverse=False)
                (_, _, new_normed_output, _), decode_caches = self(
                    new_normed_input,
                    new_mask,
                    decode_caches,
                    cross_kv=cross_kv,
                )

                new_renormed_output = torch.reshape(
                    revin(new_normed_output, new_mu, new_sigma, reverse=True),
                    (batch_size, self.m, self.o, self.q),
                )
                ar_outputs.append(new_renormed_output[:, -1, ...])
                last_renormed_output = new_renormed_output[:, -1, :, self.aridx]

            if num_decode_steps > 0:
                ar_renormed_outputs = torch.stack(ar_outputs, dim=1)
            else:
                ar_renormed_outputs = None

        return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs


@dataclass
class InternS2PreviewTimeSeriesForecasterOutput(ModelOutput):
    """Output of :class:`InternS2PreviewTimeSeriesForecaster`.

    Attributes:
      point_forecast: list of (horizon_i, C_i) per-sample median forecasts.
      quantile_forecast: list of (horizon_i, C_i, 10) per-sample quantile
        forecasts (quantile order: [median, 0.1, 0.2, ..., 0.9]).
      predicted_horizon: (B,) long tensor of horizons predicted by the horizon
        head, or None when the head is disabled / a horizon override was given.
    """

    point_forecast: list[torch.Tensor] | None = None
    quantile_forecast: list[torch.Tensor] | None = None
    predicted_horizon: torch.Tensor | None = None


class InternS2PreviewTimeSeriesForecaster(PreTrainedModel):
    """Standalone TimeOmni_v2 forecaster (cross-attention Forecaster head).

    Inputs (see :meth:`forward`): the raw multi-channel ``history``, plus the two
    precomputed embedding streams ``llm_embedding_input`` and
    ``ts_encoder_embedding_input``. The LLM and TS encoder themselves are NOT part
    of this model.
    """

    config_class = InternS2PreviewTimeSeriesForecasterConfig
    base_model_prefix = "interns2_preview_time_series_forecaster"
    main_input_name = "history"
    supports_gradient_checkpointing = False

    def __init__(self, config: InternS2PreviewTimeSeriesForecasterConfig):
        super().__init__(config)

        self.aligner = Aligner(config)
        self.forecaster = ForecasterBackbone(
            use_cross_attn=True,
            cross_attn_kv_dim=config.cross_attn_kv_dim,
            use_cross_attn_gate=config.use_cross_attn_gate,
        )

        # Normalize max_context / max_horizon to Forecaster patch multiples.
        if config.max_context % self.forecaster.p != 0:
            config.max_context = math.ceil(config.max_context / self.forecaster.p) * self.forecaster.p
        if config.max_horizon % self.forecaster.o != 0:
            config.max_horizon = math.ceil(config.max_horizon / self.forecaster.o) * self.forecaster.o
        if config.max_context + config.max_horizon > self.forecaster.context_limit:
            raise ValueError(
                "Context + horizon must be less than the context limit."
                f" {config.max_context} + {config.max_horizon} > {self.forecaster.context_limit}."
            )
        if config.use_continuous_quantile_head and (config.max_horizon > self.forecaster.os):
            raise ValueError(f"Continuous quantile head is not supported for horizons > {self.forecaster.os}.")

        self._horizon_max_length = int(config.horizon_max_length) or int(config.max_horizon)

        self.post_init()

    def _init_weights(self, module):
        """No-op: every submodule self-initializes in its own constructor.

        Overriding this keeps post_init() from clobbering the carefully chosen
        inits (RMSNorm.scale=0, PerDimScale=0, cross_attn_gate=0, QFormer query
        trunc-normal, ...)."""
        return

    # ------------------------------------------------------------------ #
    # Forecaster forecast orchestration (ported from Forecaster_2p5.forecast /
    # Forecaster_2p5_200M_torch.compile._compiled_decode, cross-attn path).
    # ------------------------------------------------------------------ #

    def _forecaster_compiled_decode(
        self,
        horizon,
        inputs,
        masks,
        enable_grad: bool = False,
        cross_kv=None,
    ):
        """Single-series decode with revin / flip-invariance / quantile post-proc."""
        fc = self.config
        module = self.forecaster
        if horizon > fc.max_horizon:
            raise ValueError(f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}.")

        with torch.set_grad_enabled(enable_grad):
            target_device = next(module.parameters()).device
            inputs = torch.as_tensor(inputs, dtype=torch.float32, device=target_device)
            masks = torch.as_tensor(masks, dtype=torch.bool, device=target_device)
            batch_size = inputs.shape[0]

            cross_kv_t = None
            if cross_kv is not None:
                cross_kv_t = torch.as_tensor(cross_kv, dtype=torch.float32, device=target_device)
                if cross_kv_t.shape[0] != batch_size:
                    raise ValueError(
                        f"Batch size mismatch between cross_kv and inputs: {cross_kv_t.shape[0]} != {batch_size}"
                    )

            if fc.infer_is_positive:
                is_positive = torch.all(inputs >= 0, dim=-1, keepdim=True)
            else:
                is_positive = None

            if fc.normalize_inputs:
                mu = torch.mean(inputs, dim=-1, keepdim=True)
                sigma = torch.std(inputs, dim=-1, keepdim=True)
                inputs = revin(inputs, mu, sigma, reverse=False)
            else:
                mu, sigma = None, None

            pf_outputs, quantile_spreads, ar_outputs = module.decode(
                horizon,
                inputs,
                masks,
                enable_grad=enable_grad,
                cross_kv=cross_kv_t,
            )
            to_cat = [pf_outputs[:, -1, ...]]
            if ar_outputs is not None:
                to_cat.append(ar_outputs.reshape(batch_size, -1, module.q))
            full_forecast = torch.cat(to_cat, dim=1)

            def flip_quantile_fn(x):
                return torch.cat([x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1)

            if fc.force_flip_invariance:
                flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = module.decode(
                    horizon,
                    -inputs,
                    masks,
                    enable_grad=enable_grad,
                    cross_kv=cross_kv_t,
                )
                flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads)
                flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs)
                to_cat = [flipped_pf_outputs[:, -1, ...]]
                if flipped_ar_outputs is not None:
                    to_cat.append(flipped_ar_outputs.reshape(batch_size, -1, module.q))
                flipped_full_forecast = torch.cat(to_cat, dim=1)
                quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2
                pf_outputs = (pf_outputs - flipped_pf_outputs) / 2
                full_forecast = (full_forecast - flipped_full_forecast) / 2

            if fc.use_continuous_quantile_head:
                for quantile_index in [1, 2, 3, 4, 6, 7, 8, 9]:
                    full_forecast[:, :horizon, quantile_index] = (
                        quantile_spreads[:, :horizon, quantile_index]
                        - quantile_spreads[:, :horizon, 5]
                        + full_forecast[:, :horizon, 5]
                    )
            full_forecast = full_forecast[:, :horizon, :]

            if fc.return_backcast:
                full_backcast = pf_outputs[:, :-1, : module.p, :].reshape(batch_size, -1, module.q)
                full_forecast = torch.cat([full_backcast, full_forecast], dim=1)

            if fc.fix_quantile_crossing:
                for i in [4, 3, 2, 1]:
                    full_forecast[:, :, i] = torch.where(
                        full_forecast[:, :, i] < full_forecast[:, :, i + 1],
                        full_forecast[:, :, i],
                        full_forecast[:, :, i + 1],
                    )
                for i in [6, 7, 8, 9]:
                    full_forecast[:, :, i] = torch.where(
                        full_forecast[:, :, i] > full_forecast[:, :, i - 1],
                        full_forecast[:, :, i],
                        full_forecast[:, :, i - 1],
                    )

            if fc.normalize_inputs:
                full_forecast = revin(full_forecast, mu, sigma, reverse=True)

            if is_positive is not None:
                full_forecast = torch.where(
                    is_positive[..., None],
                    torch.maximum(full_forecast, torch.zeros_like(full_forecast)),
                    full_forecast,
                )

        return full_forecast[..., 5], full_forecast

    def strip_leading_nans(self, arr):
        """Removes contiguous NaN values from the beginning of a NumPy array."""
        isnan = np.isnan(arr)
        first_valid_index = np.argmax(~isnan)
        return arr[first_valid_index:]

    def linear_interpolation(self, arr):
        """Linear interpolation to fill NaN values in a 1D numpy array."""
        nans = np.isnan(arr)
        if not np.any(nans):
            return arr

        def x(z):
            return z.nonzero()[0]

        nans_indices = x(nans)
        non_nans_indices = x(~nans)
        non_nans_values = arr[~nans]

        try:
            arr[nans] = np.interp(nans_indices, non_nans_indices, non_nans_values)
        except ValueError:
            if non_nans_values:
                mu = np.nanmean(arr)
            else:
                mu = 0.0
            arr = np.where(np.isfinite(arr), arr, mu)
        return arr

    def _forecaster_forecast(
        self,
        horizon: int,
        inputs: list,
        enable_grad: bool = False,
        cross_kv: list | None = None,
    ):
        """Batched forecast over a list of 1-D series (ported from Forecaster_2p5)."""
        max_context = self.config.max_context
        num_inputs = len(inputs)
        if cross_kv is not None and len(cross_kv) != num_inputs:
            raise ValueError(f"cross_kv length must match inputs length: {len(cross_kv)} != {num_inputs}")

        output_points = []
        output_quantiles = []
        for input_index, each_input in enumerate(inputs):
            value = torch.as_tensor(each_input, dtype=torch.float32).reshape(-1)
            if value.numel() > 0:
                value_np = self.linear_interpolation(self.strip_leading_nans(value.detach().cpu().numpy()))
                value = torch.from_numpy(value_np).to(dtype=torch.float32)

            if value.numel() > max_context:
                value = value[-max_context:]

            context = max(
                self.forecaster.p,
                math.ceil(value.numel() / self.forecaster.p) * self.forecaster.p,
            )
            context = min(context, max_context)

            value_len = value.numel()
            if value_len >= context:
                value = value[-context:]
                mask = torch.zeros(context, dtype=torch.bool)
            else:
                pad_len = context - value_len
                value = torch.cat([torch.zeros(pad_len, dtype=torch.float32), value], dim=0)
                mask = torch.cat(
                    [torch.ones(pad_len, dtype=torch.bool), torch.zeros(value_len, dtype=torch.bool)],
                    dim=0,
                )

            values_t = value.unsqueeze(0)
            masks_t = mask.unsqueeze(0)
            cross_kv_t = None
            if cross_kv is not None:
                cross_kv_t = torch.as_tensor(
                    cross_kv[input_index],
                    dtype=torch.float32,
                    device=value.device,
                ).unsqueeze(0)

            point_forecast, quantile_forecast = self._forecaster_compiled_decode(
                horizon,
                values_t,
                masks_t,
                enable_grad=enable_grad,
                cross_kv=cross_kv_t,
            )
            output_points.append(point_forecast)
            output_quantiles.append(quantile_forecast)

        output_points = torch.cat(output_points, dim=0)
        output_quantiles = torch.cat(output_quantiles, dim=0)
        return output_points[:num_inputs], output_quantiles[:num_inputs]

    @staticmethod
    def _flatten_channelwise(history: list[torch.Tensor], ctx: torch.Tensor):
        """Split each (T, C) history into C single-channel series; replicate the
        per-sample cross-attn context once per channel (Forecaster forecasts a single
        univariate series at a time)."""
        channel_counts = []
        flattened_inputs = []
        flattened_prefix = []
        for sample_idx, ts in enumerate(history):
            if ts.dim() != 2:
                raise ValueError(f"Each history series must be 2D (T, C), got shape {tuple(ts.shape)}")
            channels = ts.shape[1]
            channel_counts.append(channels)
            for channel_idx in range(channels):
                flattened_inputs.append(ts[:, channel_idx])
                flattened_prefix.append(ctx[sample_idx])
        return channel_counts, flattened_inputs, flattened_prefix

    def forward(
        self,
        history: list[torch.Tensor],
        llm_embedding_input: torch.Tensor,
        ts_encoder_embedding_input: torch.Tensor,
        llm_embedding_mask: torch.Tensor | None = None,
        ts_encoder_embedding_mask: torch.Tensor | None = None,
        override_horizon: int | Sequence[int] | None = None,
        return_dict: bool | None = None,
    ) -> InternS2PreviewTimeSeriesForecasterOutput:
        """Forecast multi-channel time series.

        Args:
          history: list of B per-sample 2-D tensors (T_i, C_i) — the numeric
            forecast context fed to Forecaster. T and C may differ across samples.
          llm_embedding_input: (B, T_llm, d_llm) precomputed LLM hidden states.
          ts_encoder_embedding_input: (B, T_ts, d_ts_encoder) precomputed TS-encoder
            hidden states.
          llm_embedding_mask / ts_encoder_embedding_mask: optional (B, T) bool masks
            (True = valid). Default: all valid.
          override_horizon: optional forecast horizon. An int applies to every
            sample; a sequence gives one horizon per sample. When omitted, the
            horizon head predicts it (if enabled), else ``config.default_pred_len``.
          return_dict: kept for HF API compatibility (output is always a
            :class:`TSForecasterOutput`).

        Returns:
          :class:`TSForecasterOutput` with per-sample ``point_forecast`` /
          ``quantile_forecast`` lists and the predicted horizon.
        """
        if not isinstance(history, (list, tuple)):
            raise ValueError("history must be a list/tuple of (T, C) tensors, one per sample")
        batch_size = len(history)
        if llm_embedding_input.size(0) != batch_size or ts_encoder_embedding_input.size(0) != batch_size:
            raise ValueError(
                "Batch size mismatch: len(history)="
                f"{batch_size}, llm_embedding_input={llm_embedding_input.size(0)},"
                f" ts_encoder_embedding_input={ts_encoder_embedding_input.size(0)}"
            )

        # 1. Align modalities -> cross-attention KV stream.
        ctx = self.aligner(
            llm_embedding_input,
            ts_encoder_embedding_input,
            llm_embedding_mask=llm_embedding_mask,
            ts_encoder_embedding_mask=ts_encoder_embedding_mask,
        )

        # 2. Resolve per-sample horizons.
        predicted_horizon = None
        if self.aligner.horizon_head is not None:
            pred_horizon = self.aligner.predict_horizon(
                llm_embedding_input,
                llm_embedding_mask=llm_embedding_mask,
            )
            predicted_horizon = self.aligner.decode_horizon(
                pred_horizon,
                min_h=1,
                max_h=self._horizon_max_length,
            )
        if override_horizon is not None:
            if isinstance(override_horizon, int):
                per_sample_horizons = [int(override_horizon)] * batch_size
            else:
                per_sample_horizons = [int(h) for h in override_horizon]
        elif self.aligner.horizon_head is not None:
            per_sample_horizons = [int(v) for v in predicted_horizon.tolist()]
        else:
            per_sample_horizons = [int(self.config.default_pred_len)] * batch_size

        if len(per_sample_horizons) != batch_size:
            raise ValueError(
                f"per-sample horizon count ({len(per_sample_horizons)}) does not match batch size ({batch_size})"
            )
        # Forecaster decodes a single horizon for the whole flattened batch; use the
        # max and slice each sample back to its own length below.
        horizon = max(per_sample_horizons) if per_sample_horizons else 1

        # 3. Flatten channelwise; replicate the cross-attn context per channel.
        channel_counts, flattened_inputs, flattened_prefix = self._flatten_channelwise(history, ctx)
        # Preprocess the KV stream in float32 (stable revin); the Forecaster module
        # re-casts it to its own dtype internally.
        flattened_prefix = [p.to(dtype=torch.float32) for p in flattened_prefix]

        # 4. Forecaster forecast (cross-attention injection).
        point_flat, quantile_flat = self._forecaster_forecast(
            horizon=horizon,
            inputs=flattened_inputs,
            enable_grad=self.training,
            cross_kv=flattened_prefix,
        )

        # 5. Reshape the flat (Bf, horizon[, Q]) outputs back to per-sample lists.
        point_outputs = []
        quantile_outputs = []
        cursor = 0
        for sample_idx, channels in enumerate(channel_counts):
            sample_horizon = max(1, min(int(per_sample_horizons[sample_idx]), horizon))
            channel_point = []
            channel_quantile = []
            for _ in range(channels):
                channel_point.append(point_flat[cursor, :sample_horizon])
                channel_quantile.append(quantile_flat[cursor, :sample_horizon, :])
                cursor += 1
            point_outputs.append(torch.stack(channel_point, dim=1))
            quantile_outputs.append(torch.stack(channel_quantile, dim=1))

        return InternS2PreviewTimeSeriesForecasterOutput(
            point_forecast=point_outputs,
            quantile_forecast=quantile_outputs,
            predicted_horizon=predicted_horizon,
        )


@dataclass
@auto_docstring(
    custom_intro="""
    Base class for Llava outputs, with hidden states and attentions.
    """
)
class InternS2PreviewModelOutputWithPast(ModelOutput):
    r"""
    Base class for InternS2Preview model outputs with cache and time-series
    encoder states.

    Args:
        last_hidden_state (`torch.FloatTensor`, *optional*):
            Sequence of hidden-states at the output of the last layer.
        past_key_values (`Cache`, *optional*):
            Pre-computed hidden-states that can be used to speed up sequential decoding.
        hidden_states (`tuple[torch.FloatTensor]`, *optional*):
            Hidden-states of the model at the output of each layer.
        attentions (`tuple[torch.FloatTensor]`, *optional*):
            Attention weights after the attention softmax.
        rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
            The rope index difference between sequence length and multimodal rope.
        router_logits (`tuple[torch.FloatTensor]`, *optional*):
            Router logits from MoE layers.
        ts_encoder_embedding (`torch.Tensor`, *optional*):
            Padded time-series encoder embeddings.
        ts_encoder_embedding_mask (`torch.Tensor`, *optional*):
            Boolean mask for valid time-series encoder embeddings.
        ts_history (`list[torch.Tensor]`, *optional*):
            Raw time-series histories used by the forecasting head.
    """

    last_hidden_state: torch.FloatTensor | None = None
    past_key_values: Cache | None = None
    hidden_states: tuple[torch.FloatTensor] | None = None
    attentions: tuple[torch.FloatTensor] | None = None
    rope_deltas: torch.LongTensor | None = None
    router_logits: tuple[torch.FloatTensor] | None = None

    ts_encoder_embedding: torch.Tensor | None = None
    ts_encoder_embedding_mask: torch.Tensor | None = None
    ts_history: list[torch.Tensor] | None = None


@dataclass
@auto_docstring(
    custom_intro="""
    Base class for InternS2Preview causal language model (or autoregressive) outputs.
    """
)
class InternS2PreviewCausalLMOutputWithPast(ModelOutput):
    r"""
    Output type of InternS2Preview causal language modeling with optional
    time-series forecasting outputs.

    Args:
        loss (`torch.FloatTensor` of shape `(1,)`, *optional*):
            Language modeling loss.
        logits (`torch.FloatTensor`, *optional*):
            Prediction scores of the language modeling head.
        past_key_values (`Cache`, *optional*):
            Pre-computed hidden-states that can be used to speed up sequential decoding.
        hidden_states (`tuple[torch.FloatTensor]`, *optional*):
            Hidden-states of the model at the output of each layer.
        attentions (`tuple[torch.FloatTensor]`, *optional*):
            Attention weights after the attention softmax.
        rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
            The rope index difference between sequence length and multimodal rope.
        router_logits (`tuple[torch.FloatTensor]`, *optional*):
            Router logits from MoE layers.
        aux_loss (`torch.FloatTensor`, *optional*):
            Auxiliary router loss.
        ts_point_forecast (`list[torch.Tensor]`, *optional*):
            Point forecasts for each sample.
        ts_quantile_forecast (`list[torch.Tensor]`, *optional*):
            Quantile forecasts for each sample.
        ts_predicted_horizon (`torch.Tensor`, *optional*):
            Predicted forecast horizon.
    """

    loss: torch.FloatTensor | None = None
    logits: torch.FloatTensor | None = None
    past_key_values: Cache | None = None
    hidden_states: tuple[torch.FloatTensor] | None = None
    attentions: tuple[torch.FloatTensor] | None = None
    rope_deltas: torch.LongTensor | None = None
    router_logits: tuple[torch.FloatTensor] | None = None
    aux_loss: torch.FloatTensor | None = None
    ts_point_forecast: list[torch.Tensor] | None = None
    ts_quantile_forecast: list[torch.Tensor] | None = None
    ts_predicted_horizon: torch.Tensor | None = None


@auto_docstring
class InternS2PreviewModel(InternS2PreviewPreTrainedModel):
    base_model_prefix = "model"
    _checkpoint_conversion_mapping = {}
    # Reference: fix gemma3 grad acc #37208
    accepts_loss_kwargs = False

    config: InternS2PreviewConfig
    _no_split_modules = [
        "Qwen3_5MoeTextDecoderLayer",
        "Qwen3_5MoeVisionBlock",
        "InternS2PreviewTimeSeriesEncoderLayer",
    ]

    def __init__(self, config):
        super().__init__(config)
        self.visual = InternS2PreviewVisionModel._from_config(config.vision_config)
        self.language_model = InternS2PreviewTextModel._from_config(config.text_config)
        self.rope_deltas = None  # cache rope_deltas here
        self.time_series = InternS2PreviewTimeSeriesModel._from_config(config.ts_config)

        # Initialize weights and apply final processing
        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 get_rope_index(
        self,
        input_ids: torch.LongTensor | None = None,
        image_grid_thw: torch.LongTensor | None = None,
        video_grid_thw: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Different from the original implementation, InternS2Preview use timestamps rather than absolute time position ids."""

        # Since we use timestamps to separate videos, like <t1> <vision_start> <frame1> <vision_end> <t2> <vision_start> <frame2> <vision_end>, the video_grid_thw should also be split
        if video_grid_thw is not None:
            video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
            video_grid_thw[:, 0] = 1

        image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None
        video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None

        spatial_merge_size = self.config.vision_config.spatial_merge_size
        image_token_id = self.config.image_token_id
        video_token_id = self.config.video_token_id
        vision_start_token_id = self.config.vision_start_token_id
        mrope_position_deltas = []
        total_input_ids = input_ids
        if attention_mask is None:
            attention_mask = torch.ones_like(total_input_ids)
        position_ids = torch.zeros(
            3,
            input_ids.shape[0],
            input_ids.shape[1],
            dtype=input_ids.dtype,
            device=input_ids.device,
        )
        image_index, video_index = 0, 0
        attention_mask = attention_mask.to(total_input_ids.device)
        for i, input_ids in enumerate(total_input_ids):
            input_ids = input_ids[attention_mask[i] == 1]
            image_nums, video_nums = 0, 0
            vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
            vision_tokens = input_ids[vision_start_indices + 1]
            image_nums = (vision_tokens == image_token_id).sum()
            video_nums = (vision_tokens == video_token_id).sum()
            input_tokens = input_ids.tolist()
            llm_pos_ids_list: list = []
            st = 0
            remain_images, remain_videos = image_nums, video_nums
            for _ in range(image_nums + video_nums):
                if image_token_id in input_tokens and remain_images > 0:
                    ed_image = input_tokens.index(image_token_id, st)
                else:
                    ed_image = len(input_tokens) + 1
                if video_token_id in input_tokens and remain_videos > 0:
                    ed_video = input_tokens.index(video_token_id, st)
                else:
                    ed_video = len(input_tokens) + 1
                if ed_image < ed_video:
                    t, h, w = image_grid_thw_list[image_index]
                    image_index += 1
                    remain_images -= 1
                    ed = ed_image
                else:
                    t, h, w = video_grid_thw_list[video_index]
                    video_index += 1
                    remain_videos -= 1
                    ed = ed_video
                llm_grid_t, llm_grid_h, llm_grid_w = (
                    t,
                    h // spatial_merge_size,
                    w // spatial_merge_size,
                )
                text_len = ed - st
                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
                # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos)
                t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
                h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
                w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
                llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
                st = ed + llm_grid_t * llm_grid_h * llm_grid_w
            if st < len(input_tokens):
                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
                text_len = len(input_tokens) - st
                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
            llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
            position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
            mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
        mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
        return position_ids, mrope_position_deltas

    @can_return_tuple
    @auto_docstring
    def get_video_features(
        self,
        pixel_values_videos: torch.FloatTensor,
        video_grid_thw: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple | BaseModelOutputWithPooling:
        r"""
        pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
            The tensors corresponding to the input videos.
        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
            The temporal, height and width of feature shape of each video in LLM.
        """
        # Same implementation as for images
        return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs)

    @can_return_tuple
    @auto_docstring
    def get_image_features(
        self,
        pixel_values: torch.FloatTensor,
        image_grid_thw: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple | BaseModelOutputWithPooling:
        r"""
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
            The tensors corresponding to the input images.
        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
            The temporal, height and width of feature shape of each image in LLM.
        """
        pixel_values = pixel_values.type(self.visual.dtype)
        vision_output: BaseModelOutputWithPooling = self.visual(
            pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs
        )
        image_embeds = vision_output.pooler_output
        split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
        image_embeds = torch.split(image_embeds, split_sizes)
        vision_output.pooler_output = image_embeds

        return vision_output

    def get_placeholder_mask(
        self,
        input_ids: torch.LongTensor,
        inputs_embeds: torch.FloatTensor,
        image_features: torch.FloatTensor | None = None,
        video_features: torch.FloatTensor | None = None,
    ):
        """
        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
        equal to the length of multimodal features. If the lengths are different, an error is raised.
        """
        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)
            )
            special_image_mask = special_image_mask.all(-1)
            special_video_mask = inputs_embeds == self.get_input_embeddings()(
                torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
            )
            special_video_mask = special_video_mask.all(-1)
        else:
            special_image_mask = input_ids == self.config.image_token_id
            special_video_mask = input_ids == self.config.video_token_id

        n_image_tokens = special_image_mask.sum()
        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
        if image_features is not None:
            torch_compilable_check(
                inputs_embeds[special_image_mask].numel() == image_features.numel(),
                f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
            )

        n_video_tokens = special_video_mask.sum()
        special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
        if video_features is not None:
            torch_compilable_check(
                inputs_embeds[special_video_mask].numel() == video_features.numel(),
                f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
            )
        return special_image_mask, special_video_mask

    def compute_3d_position_ids(
        self,
        input_ids: torch.Tensor | None,
        inputs_embeds: torch.Tensor | None,
        image_grid_thw: torch.Tensor | None = None,
        video_grid_thw: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        past_key_values: torch.Tensor | None = None,
    ) -> torch.Tensor | None:
        past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
        can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None)

        if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0):
            position_ids, rope_deltas = self.get_rope_index(
                input_ids,
                image_grid_thw=image_grid_thw,
                video_grid_thw=video_grid_thw,
                attention_mask=attention_mask,
            )
            self.rope_deltas = rope_deltas
        # Use pre-calculated rope-deltas to infer correct 3D position ids
        elif self.rope_deltas is not None:
            batch_size, seq_length, _ = inputs_embeds.shape
            if attention_mask is not None:
                position_ids = attention_mask.long().cumsum(-1) - 1
                position_ids = position_ids.masked_fill(attention_mask == 0, 0)
                position_ids = position_ids.view(1, batch_size, -1).repeat(3, 1, 1).to(inputs_embeds.device)
            else:
                position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length)
                position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device)
            delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0)
            position_ids = position_ids + delta.to(device=position_ids.device)
        else:
            # Can't build correct 3D positions. Let the model infer it from `cache_position`
            position_ids = None
        return position_ids

    @auto_docstring
    @can_return_tuple
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        pixel_values: torch.Tensor | None = None,
        pixel_values_videos: torch.FloatTensor | None = None,
        ts_values: torch.FloatTensor | list[torch.FloatTensor] = None,
        image_grid_thw: torch.LongTensor | None = None,
        video_grid_thw: torch.LongTensor | None = None,
        ts_lens: torch.Tensor | list[torch.Tensor] = None,
        ts_sr: torch.FloatTensor | list[torch.FloatTensor] = None,
        ts_channels: torch.LongTensor | None = None,
        cache_position: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple | InternS2PreviewModelOutputWithPast:
        r"""
        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
            The temporal, height and width of feature shape of each image in LLM.
        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
            The temporal, height and width of feature shape of each video in LLM.
        ts_values (`torch.FloatTensor` of shape `(batch_size, seq_len, num_channels)`, *optional*):
            The tensors corresponding to the input time series signals.
        ts_lens (`torch.Tensor` of shape `(batch_size,)`, *optional*):
            The valid lengths of each time series signal in the batch.
        ts_sr (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
            The sampling rates of each time series signal in the batch.
        ts_channels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            The number of channels of each time series signal in the batch.
        """
        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 None:
            inputs_embeds = self.get_input_embeddings()(input_ids)

        if pixel_values is not None:
            image_outputs: BaseModelOutputWithPooling = self.get_image_features(
                pixel_values, image_grid_thw, return_dict=True
            )
            image_embeds = image_outputs.pooler_output
            image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
            image_mask, _ = self.get_placeholder_mask(
                input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
            )
            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)

        if pixel_values_videos is not None:
            video_outputs: BaseModelOutputWithPooling = self.get_video_features(
                pixel_values_videos, video_grid_thw, return_dict=True
            )
            video_embeds = video_outputs.pooler_output
            video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
            _, video_mask = self.get_placeholder_mask(
                input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
            )
            inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)

        ts_encoder_embedding = None
        ts_encoder_embedding_mask = None
        ts_history = None
        if ts_values is not None:
            ts_embeds, ts_pad_mask, ts_embeds_before_project = self.get_ts_feature(
                ts_values, ts_lens, ts_sr, ts_channels
            )  # [B, T, C], [B, T]
            ts_encoder_embedding = ts_embeds_before_project
            ts_encoder_embedding_mask = ~ts_pad_mask
            ts_features = ts_embeds[~ts_pad_mask].to(
                inputs_embeds.device, inputs_embeds.dtype
            )  # [num_valid_ts_tokens, C]
            B, N, C = inputs_embeds.shape
            input_ids = input_ids.reshape(B * N)
            inputs_embeds = inputs_embeds.reshape(B * N, C)
            # replace ts_token in inputs_embeds and attention_mask
            ts_placeholder = input_ids == self.config.ts_token_id
            n_ts_placeholders = ts_placeholder.sum().item()
            n_ts_tokens = ts_features.size(0)
            assert n_ts_placeholders == n_ts_tokens, (
                f"[ERROR]: Mismatch: <TS_CONTEXT> tokens={n_ts_placeholders}, ts_embeds_valid={n_ts_tokens}"
            )
            input_ids = input_ids.reshape(B, N)

            try:
                # TODO why not scatter?
                inputs_embeds[ts_placeholder] = inputs_embeds[ts_placeholder] * 0.0 + ts_features
            except Exception as e:
                print(
                    f"warning: {e}, inputs_embeds[selected].shape={inputs_embeds[ts_placeholder].shape}, ts_embeds_valid.shape={ts_features.shape}"
                )
                inputs_embeds[ts_placeholder] = inputs_embeds[ts_placeholder] * 0.0 + n_ts_tokens[:n_ts_placeholders]

            inputs_embeds = inputs_embeds.reshape(B, N, C)

            ts_history = [ts_values[i, : ts_lens[i], : ts_channels[i]] for i in range(ts_values.shape[0])]

        if position_ids is None:
            position_ids = self.compute_3d_position_ids(
                input_ids=input_ids,
                image_grid_thw=image_grid_thw,
                video_grid_thw=video_grid_thw,
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
            )

        outputs = self.language_model(
            input_ids=None,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            **kwargs,
        )

        return InternS2PreviewModelOutputWithPast(
            **outputs,
            rope_deltas=self.rope_deltas,
            ts_encoder_embedding=ts_encoder_embedding,
            ts_encoder_embedding_mask=ts_encoder_embedding_mask,
            ts_history=ts_history,
        )

    @can_return_tuple
    @auto_docstring
    def get_ts_feature(
        self,
        ts_values: torch.FloatTensor | list[torch.FloatTensor],
        ts_lens: torch.Tensor | list[torch.Tensor],
        sr: torch.FloatTensor | list[torch.FloatTensor],
        ts_channels: torch.LongTensor | None = None,
    ) -> tuple[torch.FloatTensor, torch.Tensor]:
        r"""
        ts_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, input_size)`):
            The time series values to be fed to a model.
        ts_lens (`torch.Tensor` of shape `(batch_size,)`):
            The length of the time series.
        sr (`torch.FloatTensor` of shape `(batch_size,)`):
            The sampling rate of the time series.
        ts_channels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            The number of channels of the time series.
        """
        ts_embeds, ts_pad_mask, ts_embeds_before_project = self.time_series(
            time_series_signals=ts_values,
            ts_lens=ts_lens,
            sr=sr,
            channels=ts_channels,
            output_hidden_states=False,
            return_dict=True,
        )
        return ts_embeds, ts_pad_mask, ts_embeds_before_project


# NOTE: Cannot inherit from `Qwen3_5MoeForCausalLM` here due to a converter limitation: when the modular
# `__init__` contains statements before `super().__init__()`, `_fix_init_location` unconditionally moves
# `super().__init__()` to the top, which prevents pre-processing `config` (e.g. extracting `config.text_config`)
# before the parent call. The full body is therefore written out explicitly.
class InternS2PreviewForCausalLM(InternS2PreviewPreTrainedModel, GenerationMixin):
    # NOTE: `config` is annotated as `InternS2PreviewConfig` rather than `InternS2PreviewTextConfig` because remote
    # code only exposes a single AutoConfig entry. The auto factory checks `model_class.config_class` against
    # the loaded config type; a mismatch raises ValueError, so the annotation must match `InternS2PreviewConfig`.
    config: InternS2PreviewConfig

    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
    _tp_plan = {"lm_head": "colwise_gather_output"}
    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
    _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"]

    def __init__(self, config: InternS2PreviewConfig | InternS2PreviewTextConfig):
        if isinstance(config, InternS2PreviewConfig):
            config = config.text_config
        super().__init__(config)
        self.model = InternS2PreviewTextModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.router_aux_loss_coef = config.router_aux_loss_coef
        self.num_experts = config.num_experts
        self.num_experts_per_tok = config.num_experts_per_tok

        # Initialize weights and apply final processing
        self.post_init()


def load_balancing_loss_func(
    gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
    num_experts: int | None = None,
    top_k=2,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor | int:
    r"""
    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.

    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
    experts is too unbalanced.

    Args:
        gate_logits:
            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
            shape [batch_size X sequence_length, num_experts].
        num_experts:
            Number of experts
        top_k:
            The number of experts to route per-token, can be also interpreted as the `top-k` routing
            parameter.
        attention_mask (`torch.Tensor`, *optional*):
            The attention_mask used in forward function
            shape [batch_size X sequence_length] if not None.

    Returns:
        The auxiliary loss.
    """
    if gate_logits is None or not isinstance(gate_logits, tuple):
        return 0

    if isinstance(gate_logits, tuple):
        compute_device = gate_logits[0].device
        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)

    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)

    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)

    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)

    if attention_mask is None:
        # Compute the percentage of tokens routed to each experts
        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)

        # Compute the average probability of routing to these experts
        router_prob_per_expert = torch.mean(routing_weights, dim=0)
    else:
        batch_size, sequence_length = attention_mask.shape
        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)

        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
        expert_attention_mask = (
            attention_mask[None, :, :, None, None]
            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
            .reshape(-1, top_k, num_experts)
            .to(compute_device)
        )

        # Compute the percentage of tokens routed to each experts
        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
            expert_attention_mask, dim=0
        )

        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
        router_per_expert_attention_mask = (
            attention_mask[None, :, :, None]
            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
            .reshape(-1, num_experts)
            .to(compute_device)
        )

        # Compute the average probability of routing to these experts
        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
            router_per_expert_attention_mask, dim=0
        )

    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
    return overall_loss * num_experts


class InternS2PreviewForConditionalGeneration(InternS2PreviewPreTrainedModel, GenerationMixin):
    _checkpoint_conversion_mapping = {}
    _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
    # Reference: fix gemma3 grad acc #37208
    accepts_loss_kwargs = False

    config: InternS2PreviewConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = InternS2PreviewModel(config)
        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
        self.time_series_forecaster = InternS2PreviewTimeSeriesForecaster(config.ts_forecaster_config)

        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)

    @auto_docstring
    def get_video_features(
        self,
        pixel_values_videos: torch.FloatTensor,
        video_grid_thw: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple | BaseModelOutputWithPooling:
        r"""
        pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
            The tensors corresponding to the input videos.
        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
            The temporal, height and width of feature shape of each video in LLM.
        """
        return self.model.get_video_features(
            pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs
        )

    @auto_docstring
    def get_image_features(
        self,
        pixel_values: torch.FloatTensor,
        image_grid_thw: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple | BaseModelOutputWithPooling:
        r"""
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
            The tensors corresponding to the input images.
        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
            The temporal, height and width of feature shape of each image in LLM.
        """
        return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs)

    @can_return_tuple
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        pixel_values: torch.Tensor | None = None,
        pixel_values_videos: torch.FloatTensor | None = None,
        ts_values: torch.FloatTensor | list[torch.FloatTensor] = None,
        image_grid_thw: torch.LongTensor | None = None,
        video_grid_thw: torch.LongTensor | None = None,
        ts_lens: torch.Tensor | list[torch.Tensor] = None,
        ts_sr: torch.FloatTensor | list[torch.FloatTensor] = None,
        ts_channels: torch.LongTensor | None = None,
        cache_position: torch.LongTensor | None = None,
        logits_to_keep: int | torch.Tensor = 0,
        ts_forecast: bool = False,
        forecast_horizon: int | list[int] | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple | InternS2PreviewCausalLMOutputWithPast:
        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]`.
        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
            The temporal, height and width of feature shape of each image in LLM.
        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
            The temporal, height and width of feature shape of each video in LLM.

        Example:
        ```python
        >>> from PIL import Image
        >>> import requests
        >>> from transformers import AutoProcessor, Qwen3_5MoeForConditionalGeneration

        >>> model = Qwen3_5MoeForConditionalGeneration.from_pretrained("Qwen/Qwen3.5-35B-A3B-Instruct", dtype="auto", device_map="auto")
        >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3.5-35B-A3B-Instruct")

        >>> messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
                    },
                    {"type": "text", "text": "Describe this image in short."},
                ],
            }
        ]

        >>> # Preparation for inference
        >>> inputs = processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt"
        )
        >>> inputs = inputs.to(model.device)

        >>> # Generate
        >>> generated_ids = model.generate(**inputs, max_new_tokens=128)
        >>> generated_ids_trimmed = [
            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        >>> processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "A woman in a plaid shirt sits on a sandy beach at sunset, smiling as she gives a high-five to a yellow Labrador Retriever wearing a harness. The ocean waves roll in the background."
        ```"""

        outputs = self.model(
            input_ids=input_ids,
            pixel_values=pixel_values,
            pixel_values_videos=pixel_values_videos,
            ts_values=ts_values,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            ts_lens=ts_lens,
            ts_sr=ts_sr,
            ts_channels=ts_channels,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            **kwargs,
        )

        hidden_states = outputs[0]

        # 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, :])

        ts_forecast_outputs = None
        if ts_forecast and ts_values is not None:
            llm_embedding_mask = attention_mask.bool() if attention_mask is not None else None
            ts_forecast_outputs = self.time_series_forecaster(
                history=outputs.ts_history,
                llm_embedding_input=hidden_states,
                ts_encoder_embedding_input=outputs.ts_encoder_embedding,
                llm_embedding_mask=llm_embedding_mask,
                ts_encoder_embedding_mask=outputs.ts_encoder_embedding_mask,
                override_horizon=forecast_horizon,
            )

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)

        aux_loss = None
        if kwargs.get("output_router_logits", False):
            aux_loss = load_balancing_loss_func(
                outputs.router_logits,
                self.config.text_config.num_experts,
                self.config.text_config.num_experts_per_tok,
                attention_mask,
            )
            if labels is not None:
                loss += self.config.text_config.router_aux_loss_coef * aux_loss.to(
                    loss.device
                )  # make sure to reside in the same device

        return InternS2PreviewCausalLMOutputWithPast(
            loss=loss,
            aux_loss=aux_loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            rope_deltas=outputs.rope_deltas,
            router_logits=outputs.router_logits,
            ts_point_forecast=None if ts_forecast_outputs is None else ts_forecast_outputs.point_forecast,
            ts_quantile_forecast=None if ts_forecast_outputs is None else ts_forecast_outputs.quantile_forecast,
            ts_predicted_horizon=None if ts_forecast_outputs is None else ts_forecast_outputs.predicted_horizon,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        cache_position=None,
        position_ids=None,
        use_cache=True,
        pixel_values=None,
        pixel_values_videos=None,
        ts_values=None,
        image_grid_thw=None,
        video_grid_thw=None,
        ts_lens=None,
        ts_sr=None,
        ts_channels=None,
        is_first_iteration=False,
        **kwargs,
    ):
        model_inputs = super().prepare_inputs_for_generation(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            position_ids=position_ids,
            use_cache=use_cache,
            pixel_values=pixel_values,
            pixel_values_videos=pixel_values_videos,
            ts_values=ts_values,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            ts_lens=ts_lens,
            ts_sr=ts_sr,
            ts_channels=ts_channels,
            is_first_iteration=is_first_iteration,
            **kwargs,
        )

        if not is_first_iteration and use_cache:
            model_inputs["ts_values"] = None
            model_inputs["ts_lens"] = None
            model_inputs["ts_sr"] = None
            model_inputs["ts_channels"] = None

        return model_inputs

    def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs):
        # Overwritten -- requires 3D position ids

        text_positions = super()._prepare_position_ids_for_generation(inputs_tensor, model_kwargs)

        # Early exit in case we are continuing generation from past kv
        past_length = 0
        if (cache := model_kwargs.get("past_key_values")) is not None:
            past_length = cache.get_seq_length()
        if past_length != 0 and self.model.rope_deltas is not None:
            text_positions += self.model.rope_deltas
            return text_positions

        # Otherwise compute 3d position ids for vision tokens and concat with text position ids
        if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0:
            inputs_tensor = model_kwargs["input_ids"]

        is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long]
        if is_input_ids and (
            model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None
        ):
            model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"}
            vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs)
            self.model.rope_deltas = rope_deltas
        else:
            vision_positions = text_positions.unsqueeze(0).expand(3, -1, -1)
            self.model.rope_deltas = torch.zeros(
                inputs_tensor.shape[0], 1, dtype=torch.long, device=inputs_tensor.device
            )

        # Concatenate "text + vision" positions into [4, bs, seq-len]
        text_positions = text_positions[None, ...]
        position_ids = torch.cat([text_positions, vision_positions], dim=0)

        return position_ids

    def _get_image_nums_and_video_nums(
        self,
        input_ids: torch.LongTensor | None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Get the number of images and videos for each sample to calculate the separation length of the sample tensor.
        These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications.

        Args:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                Indices of input sequence tokens in the vocabulary.

        Returns:
            image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`)
            video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`)
        """
        image_token_id = self.config.image_token_id
        video_token_id = self.config.video_token_id
        vision_start_token_id = self.config.vision_start_token_id

        if inputs_embeds is not None:
            vision_start_mask = (
                inputs_embeds
                == self.get_input_embeddings()(
                    torch.tensor(vision_start_token_id, dtype=torch.long, device=inputs_embeds.device)
                )
            )[..., 0]
            image_mask = (
                inputs_embeds
                == self.get_input_embeddings()(
                    torch.tensor(image_token_id, dtype=torch.long, device=inputs_embeds.device)
                )
            )[..., 0]
            video_mask = (
                inputs_embeds
                == self.get_input_embeddings()(
                    torch.tensor(video_token_id, dtype=torch.long, device=inputs_embeds.device)
                )
            )[..., 0]
        else:
            vision_start_mask = input_ids == vision_start_token_id
            image_mask = input_ids == image_token_id
            video_mask = input_ids == video_token_id

        vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1)
        image_nums = torch.sum(vision_first_mask & image_mask, dim=1)
        video_nums = torch.sum(vision_first_mask & video_mask, dim=1)

        return image_nums, video_nums

    def _expand_inputs_for_generation(
        self,
        expand_size: int = 1,
        is_encoder_decoder: bool = False,
        input_ids: torch.LongTensor | None = None,
        **model_kwargs,
    ) -> tuple[torch.LongTensor, dict[str, Any]]:
        # Overwritten -- InternS2Preview use timestamps and remove second_per_grid_ts
        # Support for expanding tensors without a batch size dimension
        # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw
        # pixel_values.shape[0] is sum(seqlen_images for samples)
        # image_grid_thw.shape[0] is sum(num_images for samples)

        if expand_size == 1:
            return input_ids, model_kwargs

        visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"]

        def _expand_dict_for_generation_visual(dict_to_expand):
            image_grid_thw = model_kwargs.get("image_grid_thw", None)
            video_grid_thw = model_kwargs.get("video_grid_thw", None)
            image_nums, video_nums = self._get_image_nums_and_video_nums(
                input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None)
            )

            # video_nums: (batch_size,)
            # since video_nums is the number of videos in the input dependent on the input_ids(vision_start),
            # but InternS2Preview append vision_start to each frame of each video, so we need to recover the real video_nums according to video_grid_thw
            if video_grid_thw is not None:
                cumulative_frame_counts = torch.cumsum(video_grid_thw[:, 0], dim=0)
                cumulative_token_video_counts = torch.cumsum(video_nums, dim=0)
                # Find video boundaries in cumulative_frame_counts
                video_boundary_indices = torch.searchsorted(cumulative_frame_counts, cumulative_token_video_counts)
                # example: video_boundary_indices = [3, 5] means video_nums = [4, 2]
                video_nums = torch.diff(torch.cat([-video_boundary_indices.new_ones(1), video_boundary_indices]))

            def _repeat_interleave_samples(x, lengths, repeat_times):
                samples = torch.split(x, lengths)
                repeat_args = [repeat_times] + [1] * (x.dim() - 1)
                result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
                return result

            for key in dict_to_expand:
                if key == "pixel_values":
                    # split images into samples
                    samples = torch.split(image_grid_thw, list(image_nums))
                    # compute the sequence length of images for each sample
                    lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
                    dict_to_expand[key] = _repeat_interleave_samples(
                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
                    )
                elif key == "image_grid_thw":
                    # get the num of images for each sample
                    lengths = list(image_nums)
                    dict_to_expand[key] = _repeat_interleave_samples(
                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
                    )
                elif key == "pixel_values_videos":
                    samples = torch.split(video_grid_thw, list(video_nums))
                    lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
                    dict_to_expand[key] = _repeat_interleave_samples(
                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
                    )
                elif key == "video_grid_thw":
                    lengths = list(video_nums)
                    dict_to_expand[key] = _repeat_interleave_samples(
                        dict_to_expand[key], lengths=lengths, repeat_times=expand_size
                    )
            return dict_to_expand

        def _expand_dict_for_generation(dict_to_expand):
            for key in dict_to_expand:
                if key == "position_ids" and dict_to_expand[key].ndim == 3:
                    dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=1)
                elif (
                    key != "cache_position"
                    and dict_to_expand[key] is not None
                    and isinstance(dict_to_expand[key], torch.Tensor)
                    and key not in visual_keys
                ):
                    dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
            return dict_to_expand

        model_kwargs = _expand_dict_for_generation_visual(model_kwargs)

        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

    @property
    def time_series(self):
        return self.model.time_series

    @auto_docstring
    def get_ts_feature(
        self,
        ts_values: torch.FloatTensor | list[torch.FloatTensor],
        ts_lens: torch.Tensor | list[torch.Tensor],
        sr: torch.FloatTensor | list[torch.FloatTensor],
        ts_channels: torch.LongTensor | None = None,
    ):
        r"""
        ts_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, input_size)`):
            The time series values to be fed to a model.
        ts_lens (`torch.Tensor` of shape `(batch_size,)`):
            The length of the time series.
        sr (`torch.FloatTensor` of shape `(batch_size,)`):
            The sampling rate of the time series.
        ts_channels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            The number of channels of the time series.
        """
        return self.model.get_ts_feature(ts_values, ts_lens, sr, ts_channels)

    def generate_with_time_series_input(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        pixel_values: torch.Tensor | None = None,
        pixel_values_videos: torch.FloatTensor | None = None,
        ts_values: torch.FloatTensor | list[torch.FloatTensor] = None,
        image_grid_thw: torch.LongTensor | None = None,
        video_grid_thw: torch.LongTensor | None = None,
        ts_lens: torch.Tensor | list[torch.Tensor] = None,
        ts_sr: torch.FloatTensor | list[torch.FloatTensor] = None,
        ts_channels: torch.LongTensor | None = None,
        cache_position: torch.LongTensor | None = None,
        forecast_horizon: int | list[int] | None = None,
        enable_forecasting: bool = False,
        logits_to_keep: int | torch.Tensor = 1,
        return_dict: bool | None = True,
        **generate_kwargs,
    ):
        if ts_values is None:
            raise ValueError("No time series was provided.")

        model_inputs = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "past_key_values": past_key_values,
            "inputs_embeds": inputs_embeds,
            "pixel_values": pixel_values,
            "pixel_values_videos": pixel_values_videos,
            "ts_values": ts_values,
            "image_grid_thw": image_grid_thw,
            "video_grid_thw": video_grid_thw,
            "ts_lens": ts_lens,
            "ts_sr": ts_sr,
            "ts_channels": ts_channels,
            "cache_position": cache_position,
        }

        if enable_forecasting:
            return self(
                **model_inputs,
                ts_forecast=True,
                forecast_horizon=forecast_horizon,
                logits_to_keep=logits_to_keep,
                return_dict=return_dict,
            )

        return self.generate(**model_inputs, **generate_kwargs)


__all__ = [
    "InternS2PreviewVisionModel",
    "InternS2PreviewTextModel",
    "InternS2PreviewModel",
    "InternS2PreviewForCausalLM",
    "InternS2PreviewForConditionalGeneration",
    "InternS2PreviewPreTrainedModel",
]