"""NKI kernels for RoPE ViT encoders on AWS Trainium2 (trn2).

Model layers (drop-in forwards for Hugging Face Transformers):
  NeuronDINOv3ViTModel   -- DINOv3ViTModel (S/S+/B/L/H+)
  NeuronVJEPA2Encoder    -- VJEPA2Encoder

Model-agnostic building blocks:
  BlockWeights, pack_block, rope_tables, EncoderRunner   -- pre-norm encoder loop on the kernels
  kernels: attention_block, attention_block_l1, attention_block_l1_ln, res_layernorm_fold,
           mlp, mlp_res_ln, swiglu_stream, swiglu_stream_ln
"""

from .encoder import LNC, BlockWeights, EncoderRunner, calibrate_fp8, pack_block, rope_tables
from .layers import NeuronDINOv3ViTModel, NeuronVJEPA2Encoder
from .nki_kernels import (
    attention_block,
    attention_block_l1,
    attention_block_l1_ln,
    lane_permutation,
    mlp,
    mlp_res_ln,
    pack_out_weights,
    pack_tables,
    pack_v_weights,
    pack_weights,
    res_layernorm_fold,
    swiglu_stream,
    swiglu_stream_ln,
)

__all__ = [
    "NeuronDINOv3ViTModel", "NeuronVJEPA2Encoder",
    "BlockWeights", "pack_block", "rope_tables", "EncoderRunner", "LNC", "calibrate_fp8",
    "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
    "mlp", "mlp_res_ln", "swiglu_stream", "swiglu_stream_ln",
    "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
]