utils

Utility functions for sparse attention module discovery.

Functions

get_named_sparse_attention_modules	Get all sparse attention modules in a model with their names.
get_sparse_attention_modules	Get all sparse attention modules in a model.

get_named_sparse_attention_modules(model)

Get all sparse attention modules in a model with their names.

Parameters:

model (Module) – Model to search for sparse attention modules.

Returns:

List of (name, module) tuples for all SparseAttentionModule instances.

Return type:

list[tuple[str, SparseAttentionModule]]

get_sparse_attention_modules(model)

Get all sparse attention modules in a model.

Parameters:

model (Module) – Model to search for sparse attention modules.

Returns:

List of SparseAttentionModule instances found in the model.

Return type:

list[SparseAttentionModule]