openvino.opset17.sparse_attention_indexer

openvino.opset17.sparse_attention_indexer(index_query: Node | int | float | ndarray, block_index_keys: Node | int | float | ndarray, key_length: Node | int | float | ndarray, index_weights: Node | int | float | ndarray | None = None, *, key_block_size: int, token_budget: int, num_selection_heads: int = 1, name: str | None = None) → Node

Return a node which selects the key blocks every query token attends to.

For every query token the operation scores the complete key blocks visible to it and returns the indices of the token_budget / key_block_size best-scoring blocks together with their number.

Parameters:

• index_query – Index query heads of shape [N, Hi, L, Di].

• block_index_keys – One index key per complete key block, shape [N, B, Di].

• key_length – Number of valid key tokens of every sequence, shape [N] or [1].

• index_weights – Weight of every index query head for every query token, shape [N, L, Hi].

• key_block_size – Number of consecutive key tokens forming one key block.

• token_budget – Maximum number of tokens selected from complete blocks for one query.

• num_selection_heads – Number of selections per query token.

• name – The optional name for the new output node.

Returns:

The new node performing SparseAttentionIndexer operation.