def fast_pos_embed_interpolate(self, grid_thw):
    '''
    self.num_grid_per_side 默认为 48
    '''
    # grid_thw shape (bs, t, h, w)
    grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]
    device = grid_thw.device

    # ids_list 是用来存储 用来表征每一个patch 位置的四个角点的位置id
    # weight_list 是用来存储 用来表征每一个patch 位置的四个角点在双线性插值时的权重
    idx_list = [[] for _ in range(4)]
    weight_list = [[] for _ in range(4)]

    # 遍历batch里面的所有samples
    for t, h, w in zip(grid_ts, grid_hs, grid_ws):
        # 在 48 * 48 的板子上均匀采样 h * w 个点，每个点代表一个patch
        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

        # 因为Embedding是从一个整数映射到一个向量过程， 输入是一个整数，所以需要将角点的2维ID (x,y) 转成一个 一位坐标，使用的方法就是 （x, y) -> x *  num_grid_per_side + y 的方式。
        base_h = h_idxs_floor * self.num_grid_per_side # 四个角点中处于上方的角点的行基地址， 对应x *  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(),
        ]

        # 把所有角点的 左上全部放一起， 右上全部放一起 ....
        # idx_list[0][0], idx_list[1][0], idx_list[2][0], idx_list[3][0] 分别对应第一个patch使用的左上, 右上, 左下， 右下四个点
        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

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 shape: (seq_len, hidden_dim_per_head)
    # q: (seq_len, num_head, hidden_dim_per_head)
    # 所以cos和sin 需要unsqueeze
    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 Qwen3VLVisionModel(Qwen3VLPreTrainedModel):
    ...
    def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
        merge_size = self.spatial_merge_size

        max_hw = int(grid_thw[:, 1:].max().item())
        freq_table = self.rotary_pos_emb(max_hw)  # (max_hw, dim // 2)
        device = freq_table.device

        total_tokens = int(torch.prod(grid_thw, dim=1).sum().item())
        pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device)

        offset = 0
        for num_frames, height, width in grid_thw:
            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)
            # 这里的coord 就是 每一个token 对应的 (x, y)
            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

        # pos_ids shape: (seq_len(total visual tokens in batch), 2(x and y coordinate))
        # embeddings shape: (seq_len, 2, dim / 2 = 18), dim = visual_hiddendim = 1152 / num_head = 16 / (x, y) (2) = 36
        embeddings = freq_table[pos_ids]  # lookup rotary embeddings 
        # embeddings shape: (seq_len, dim=36)
        # 此时embeddings 储存的是 每一个 visual token hiddent states 每一维旋转的角度 的一半， 后一半和前面一半相同
        embeddings = embeddings.flatten(1)
        return embeddings

class Qwen3VLVisionModel(Qwen3VLPreTrainedModel):
    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs):
        for layer_num, blk in enumerate(self.blocks):
            # hidden_states shape: (tot_vis_token_in batch, 1152)
            hidden_states = blk(
                hidden_states,
                cu_seqlens=cu_seqlens,
                position_embeddings=position_embeddings,
                **kwargs,
            )
            if layer_num in self.deepstack_visual_indexes:
                # deepstack_feature shape:(tot_vis_token_in batch // 4, 4096)
                deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)](
                    hidden_states
                )
                deepstack_feature_lists.append(deepstack_feature)

def _deepstack_process(
        self, hidden_states: torch.Tensor, visual_pos_masks: torch.Tensor, visual_embeds: torch.Tensor
    ):
        '''
        hidden_states shape: (bs, tot_num_tokens, hidden_dim)
        visual_pos_masks shape: (bs, tot_num_tokens)
        visual_embeds shape: (num_visual_tokens, hidden_dim)

        hidden_states[visual_pos_masks, :] shape: (num_true = num_visual_tokens_in_batch, hidden_dim)
        '''
        visual_pos_masks = visual_pos_masks.to(hidden_states.device)
        visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
        hidden_states = hidden_states.clone()
        # 好神奇，没有将visual_embeds split 开不同的sample， 而是使用mask 
        local_this = hidden_states
        [visual_pos_masks, :] + visual_embeds
        hidden_states[visual_pos_masks, :] = local_this
        return hidden_states

def _preprocess(
    ~ self,
    images: Union[ImageInput, VideoInput],
    do_resize: Optional[bool] = None,
    size: Optional[dict[str, int]] = None,
    resample: Optional[PILImageResampling] = None,
    do_rescale: Optional[bool] = None,
    rescale_factor: Optional[float] = None,
    do_normalize: Optional[bool] = None,
    image_mean: Optional[Union[float, list[float]]] = None,
    image_std: Optional[Union[float, list[float]]] = None,
    patch_size: Optional[int] = None,
    temporal_patch_size: Optional[int] = None,
    merge_size: Optional[int] = None,
    do_convert_rgb: Optional[bool] = None,
    data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
    input_data_format: Optional[Union[str, ChannelDimension]] = None,
):
    """
    Preprocess an image or batch of images. Copy of the `preprocess` method from `CLIPImageProcessor`.

    这个函数针对单个sample 进行处理

    Args:
        images (`ImageInput`):
            Image or batch of images to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`.
        ...
    """

    # make_flat_list_of_images的作用就是将sample中的所有图片都是用一个List装
    # 单图输入 img -> [img], [img] -> [img]
    # 多图输入 [[img1], [img2], ...] -> [img1, img2, ...]
    images = make_flat_list_of_images(images)

    if do_convert_rgb:
        images = [convert_to_rgb(image) for image in images]

    # All transformations expect numpy arrays.
    images = [to_numpy_array(image) for image in images]

    if do_rescale and is_scaled_image(images[0]):
        logger.warning_once(
            "It looks like you are trying to rescale already rescaled images. If the input"
            " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
        )
    if input_data_format is None:
        # We assume that all images have the same channel dimension format.
        input_data_format = infer_channel_dimension_format(images[0])

    height, width = get_image_size(images[0], channel_dim=input_data_format)
    resized_height, resized_width = height, width
    processed_images = []
    for image in images:
        if do_resize:
            # 根据！第一张！图像的分辨率找到最合适的大小模版（height和width来自images[0]
            resized_height, resized_width = smart_resize(
                height,
                width,
                factor=patch_size * merge_size,
                min_pixels=size["shortest_edge"],
                max_pixels=size["longest_edge"],
            )
            image = resize(
                image, size=(resized_height, resized_width), resample=resample, input_data_format=input_data_format
            )

        if do_rescale:
            image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)

        if do_normalize:
            image = self.normalize(
                image=image, mean=image_mean, std=image_std, input_data_format=input_data_format
            )

        image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
        processed_images.append(image)

    patches = np.array(processed_images) # (N, C, H, W)
    if data_format == ChannelDimension.LAST:
        patches = patches.transpose(0, 3, 1, 2)
    '''
    qwen 将image和video 统一成视频进行处理
    对于单图或者数量不为temporal_patch_size倍数的多图样本，复制“最后一张图”（patches[-1]）直至变成temporal_patch_size的倍数
    '''
    if patches.shape[0] % temporal_patch_size != 0:
        repeats = np.repeat(
            patches[-1][np.newaxis], temporal_patch_size - (patches.shape[0] % temporal_patch_size), axis=0
        )
        patches = np.concatenate([patches, repeats], axis=0)
    channel = patches.shape[1]
    grid_t = patches.shape[0] // temporal_patch_size
    grid_h, grid_w = resized_height // patch_size, resized_width // patch_size

    '''
    这一步reshape 和 transpose很关键, 为后面的 “打patch” 和 “merger（4合1）” 做铺垫

    (
        grid_t,
        temporal_patch_size,
        channel,
        grid_h // merge_size,
        merge_size,
        patch_size,
        grid_w // merge_size,
        merge_size,
        patch_size,
    )
    - transpose ->
    (
        grid_t,
        grid_h // merge_size,
        grid_w // merge_size,
        merge_size,
        merge_size, 
        channel,
        temporal_patch_size,    
        patch_size,
        patch_size,
    )
    '''
    patches = patches.reshape(
        grid_t,
        temporal_patch_size,
        channel,
        grid_h // merge_size,
        merge_size,
        patch_size,
        grid_w // merge_size,
        merge_size,
        patch_size,
    )
    patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
    # channel * temporal_patch_size * patch_size * patch_size = 3 * 2 * 16 * 16 = 1156
    flatten_patches = patches.reshape(
        grid_t * grid_h * grid_w, channel * temporal_patch_size * patch_size * patch_size
    ) # 2-dim

    return flatten_patches, (grid_t, grid_h, grid_w)

def preprocess(...):
    ...
    pixel_values, vision_grid_thws = [], []
    for image in images:
        patches, image_grid_thw = self._preprocess(
            image,
            ...
        )
        '''
        注意：这里是extend(), 即把一个batch里面的所有sample的patches拼成一个sample
        patches1 shape: (N1, C)
        patches1 shape: (N2, C)
        pixel_values shape: (\sigma N_i, C)
        '''
        pixel_values.extend(patches)
        vision_grid_thws.append(image_grid_thw)
    pixel_values = np.array(pixel_values)
    vision_grid_thws = np.array(vision_grid_thws)
    data.update({"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws})

    def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):
        """
        Encodes images into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned.

        这里我感觉注释有错：
        pixel_values (`torch.FloatTensor` of shape `(seq_len, hidden_size)`):

        Args:
            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)
        image_embeds, deepstack_image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
        # 在这里根据每一个sample图像的image_grid_thw计算出占用多少个视觉token，然后split 开。
        # image_embeds（tuple） shape: （（N1， C), (N2, C)...)

        # 这里有点没想通，这里把不同sample的视觉token split开，但是后面
        # image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) 又把他concat 起来， 那为什么还要split？
        split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
        image_embeds = torch.split(image_embeds, split_sizes)
        return image_embeds, deepstack_image_embeds

inputs_embeds = torch.zeros(2, 5, 2)

# 掩码: 第0个样本：位置0和2是图像；第1个样本：位置1和3是图像
image_mask = torch.tensor([
    [True, False, True, False, False],   # 样本0
    [False, True, False, True, False]    # 样本1
])  # shape: (2, 5)

# 图像 embeddings: 总共 4 个 tokens (2 + 2), shape = (4, 2)
image_embeds = torch.tensor([
    [1, 2],   # 样本0 的第1个图像token
    [3, 4],   # 样本0 的第2个图像token
    [5, 6],   # 样本1 的第1个图像token
    [7, 8]    # 样本1 的第2个图像token
])  # 注意：顺序必须与 mask 中 True 的出现顺序一致！

# 扩展 mask 到 embedding 维度
mask_expanded = image_mask.unsqueeze(-1).expand_as(inputs_embeds)  # (2, 5, 2)

"""
执行 masked_scatter, 其工作方式是：按顺序遍历 mask 中为 True（或非零）的位置,从 source 中按顺序取出元素，一一对应地填入这些位置
"""
result = inputs_embeds.masked_scatter(mask_expanded, image_embeds)

class Qwen3VLVisionPatchMerger(nn.Module):
    def __init__(self, config: Qwen3VLVisionConfig, 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:
        # 因为图像在preprocess的时候已经组织称 四个四个一组 所以view(-1, self.hidden_size)这个操作相当于把四个token concat 成一列
        x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size)
        # 然后再使用linear 变换，实现了merge token的同时，还实现了维度调整，使得最终输出的visual token 和text token hidden_dim 相同 
        x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
        return x

key_states = repeat_kv(key, module.num_key_value_groups)
value_states = repeat_kv(value, module.num_key_value_groups)

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    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)

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

    def __init__(self, config: Qwen3VLTextConfig, device=None):
        super().__init__()
        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
            self.rope_type = config.rope_scaling.get("rope_type", "default")
        else:
            self.rope_type = "default"
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.original_inv_freq = self.inv_freq

        self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20]) # hidden_dim per head in LLM = 64

    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 [THTHWHTHW...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

    @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, Qwen3VL has different position ids for the grids
        # So we expand the inv_freq to shape (3, ...)
        # position_ids shape: (3, bs, seq_len)
        # Questions: 为什么这里不像vision endcoder 中的2D-MRoPE 进行索引提取, 而是将其扩展呢?
        # Ans: 这时候还没有计算出table
        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) # (3, bs, hidden_dim(64), 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 torch.autocast(device_type=device_type, enabled=False):  # Force float32
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) # this is the freqs table (3, bs,  positions(seq_len), hidden_dim(64))
            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

        # method 2:
        max_thw = position_ids.max().item()
        seq = torch.arange(max_thw+1, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs_table = torch.outer(seq, self.inv_freq) # (max_thw, hidden_dim)
        position_ids = position_ids.int().permute(1,2,0) # (bs, seq_len, 3)
        freqs2 = freqs_table[position_ids].permute(2, 0, 1, 3) # (bs, seq_len, 3, hidden_dim)
        freqs2 = self.apply_interleaved_mrope(freqs2, self.mrope_section)
        emb2 = torch.cat((freqs2, freqs2), dim=-1)

        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)