"""Shared attention backends, masks, and optional optimized kernels."""

from ._core import (
    VALID_ATTENTION_BACKENDS,
    AttentionBackend,
    BlockMask,
    _ensure_flash_kernels_loaded,
    _get_flex_attention_fn,
    _get_flex_block_mask,
    _kernels_flash_forward,
    _kernels_flash_varlen_forward,
    _unpad_input,
    bool_to_additive_mask,
    clear_flex_attention_caches,
    create_block_mask,
    flex_attention,
    get_attention_mask,
    get_attn_implementation,
    index_first_axis,
    index_put_first_axis,
    kernels_flash_attention_func,
    pad_input,
    resolve_attention_backend,
    resolve_attention_backend_for_call,
    set_config_attn_implementation,
    warn_attention_backend_fallback,
)
from .interfaces import (
    FASTPLMS_ATTENTION_FUNCTIONS,
    FASTPLMS_ATTENTION_MASKS,
    FastPLMsAttentionMixin,
    validate_transformers_attention_interfaces,
)


__all__ = [
    "FASTPLMS_ATTENTION_FUNCTIONS",
    "FASTPLMS_ATTENTION_MASKS",
    "VALID_ATTENTION_BACKENDS",
    "AttentionBackend",
    "BlockMask",
    "FastPLMsAttentionMixin",
    "_ensure_flash_kernels_loaded",
    "_get_flex_attention_fn",
    "_get_flex_block_mask",
    "_kernels_flash_forward",
    "_kernels_flash_varlen_forward",
    "_unpad_input",
    "bool_to_additive_mask",
    "clear_flex_attention_caches",
    "create_block_mask",
    "flex_attention",
    "get_attention_mask",
    "get_attn_implementation",
    "index_first_axis",
    "index_put_first_axis",
    "kernels_flash_attention_func",
    "pad_input",
    "resolve_attention_backend",
    "resolve_attention_backend_for_call",
    "set_config_attn_implementation",
    "validate_transformers_attention_interfaces",
    "warn_attention_backend_fallback",
]