"""The NKI kernels and their host-side weight packing."""

from .attention_block import (
    attention_block,
    lane_permutation,
    pack_out_weights,
    pack_qk_bias,
    pack_tables,
    pack_v_weights,
    pack_weights,
)
from .attention_block_l1 import attention_block_l1
from .attention_block_l1_ln import attention_block_l1_ln
from .mlp import nki_mlp_k4_xp as mlp
from .mlp import pack_b1
from .mlp_res_ln import nki_mlp_res_ln as mlp_res_ln
from .res_layernorm import nki_res_layernorm_fold as res_layernorm_fold
from .swiglu_stream import nki_swiglu_stream as swiglu_stream
from .swiglu_stream import nki_swiglu_stream_ln as swiglu_stream_ln

__all__ = [
    "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
    "mlp", "mlp_res_ln", "pack_b1", "swiglu_stream", "swiglu_stream_ln",
    "pack_weights", "pack_qk_bias", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
]