code/src/models/fla/ops/utils/__init__.py

from .cumsum import (
    chunk_global_cumsum,
    chunk_global_cumsum_scalar,
    chunk_global_cumsum_vector,
    chunk_local_cumsum,
    chunk_local_cumsum_scalar,
    chunk_local_cumsum_vector,
)
from .index import (
    get_max_num_splits,
    prepare_chunk_indices,
    prepare_chunk_offsets,
    prepare_cu_seqlens_from_lens,
    prepare_cu_seqlens_from_mask,
    prepare_lens,
    prepare_lens_from_cu_seqlens,
    prepare_lens_from_mask,
    prepare_position_ids,
    prepare_sequence_ids,
    prepare_token_indices,
)
from .logsumexp import logsumexp_fwd
from .matmul import addmm, matmul
from .pack import pack_sequence, unpack_sequence
from .pooling import mean_pooling
from .softmax import softmax_bwd, softmax_fwd
from .solve_tril import solve_tril

__all__ = [
    'chunk_global_cumsum',
    'chunk_global_cumsum_scalar',
    'chunk_global_cumsum_vector',
    'chunk_local_cumsum',
    'chunk_local_cumsum_scalar',
    'chunk_local_cumsum_vector',
    'pack_sequence',
    'unpack_sequence',
    'prepare_chunk_indices',
    'prepare_chunk_offsets',
    'prepare_cu_seqlens_from_lens',
    'prepare_cu_seqlens_from_mask',
    'prepare_lens',
    'prepare_lens_from_cu_seqlens',
    'prepare_lens_from_mask',
    'prepare_position_ids',
    'prepare_sequence_ids',
    'prepare_token_indices',
    'logsumexp_fwd',
    'addmm',
    'matmul',
    'mean_pooling',
    'softmax_bwd',
    'softmax_fwd',
    'solve_tril',
    'get_max_num_splits',
]