"""V-JEPA2 ViT-L encoder NKI kernels for AWS Trainium2 (trn2).

Public API:
  NeuronVJEPA2Encoder          -- KernelConfig-mappable layer ("VJEPA2Encoder" -> this)
  res_layernorm_fold, attention_block, mlp   -- the three kernel entry points
  pack_* / build_rope_tables / lane_permutation -- host-side weight packing
"""

from .layers import NeuronVJEPA2Encoder
from .nki_kernels import (
    attention_block,
    build_rope_tables,
    lane_permutation,
    mlp,
    pack_out_weights,
    pack_tables,
    pack_v_weights,
    pack_weights,
    res_layernorm_fold,
)

__all__ = [
    "NeuronVJEPA2Encoder",
    "res_layernorm_fold", "attention_block", "mlp",
    "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights",
    "lane_permutation", "build_rope_tables",
]