Diff Coverage

    key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
    candidate_chunks = []
    for start in range(0, sequence_length, query_chunk_size):
        end = min(start + query_chunk_size, sequence_length)
        scores = torch.matmul(query[:, start:end].float(), key_fp32).relu()
        scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
        if reduce_sum is not None:
            scores = reduce_sum(scores)
        visible = (

            query, key, merge_weight, compress_ratio=compress_ratio,
            sparse_count=sparse_count, topk_blocks=topk_blocks,
            block_size=block_size, query_offset=query_offset,
            reduce_sum=reduce_sum, minimum_key_indices=minimum_key_indices)
        if selected is not None:
            return selected
    batch_size, sequence_length, _, head_dim = query.shape
    if key.ndim != 3 or key.shape[0] != batch_size or key.shape[2] != head_dim:
        raise ValueError(
            "compressed index key must have shape [batch, compressed_sequence, head_dim]"

    selected_chunks = []
    candidate_chunks = []
    for start in range(0, sequence_length, query_chunk_size):
        end = min(start + query_chunk_size, sequence_length)
        scores = torch.matmul(query[:, start:end].float(), key_fp32).relu()
        scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
        if reduce_sum is not None:
            scores = reduce_sum(scores)
        visible = (

            query, key, merge_weight, candidate_blocks,
            compress_ratio=compress_ratio, sparse_count=sparse_count,
            block_size=block_size, query_offset=query_offset,
            reduce_sum=reduce_sum, minimum_key_indices=minimum_key_indices)
        if selected is not None:
            return selected
    batch_size, sequence_length, _, head_dim = query.shape
    if key.ndim != 3 or key.shape[0] != batch_size or key.shape[2] != head_dim:
        raise ValueError(
            "compressed index key must have shape [batch, compressed_sequence, head_dim]"