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

from ._auto import (
    AUTO_ATTENTION,
    AttentionCandidate,
    AttentionExecutionContext,
    AttentionResolution,
)
from ._core import (
    LEGACY_CHECKPOINT_ATTENTION_BACKENDS,
    VALID_ATTENTION_BACKENDS,
    AttentionBackend,
    BlockMask,
    FlashPaddingLayout,
    _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,
    canonical_checkpoint_attention_backend,
    clear_flex_attention_caches,
    create_block_mask,
    flash_kernel_unsupported_reason,
    flex_attention,
    get_attention_mask,
    get_attn_implementation,
    get_flash_padding_layout,
    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__ = [
    "AUTO_ATTENTION",
    "FASTPLMS_ATTENTION_FUNCTIONS",
    "FASTPLMS_ATTENTION_MASKS",
    "LEGACY_CHECKPOINT_ATTENTION_BACKENDS",
    "VALID_ATTENTION_BACKENDS",
    "AttentionBackend",
    "AttentionCandidate",
    "AttentionExecutionContext",
    "AttentionResolution",
    "BlockMask",
    "FastPLMsAttentionMixin",
    "FlashPaddingLayout",
    "_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",
    "canonical_checkpoint_attention_backend",
    "clear_flex_attention_caches",
    "create_block_mask",
    "flash_kernel_unsupported_reason",
    "flex_attention",
    "get_attention_mask",
    "get_attn_implementation",
    "get_flash_padding_layout",
    "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",
]