File size: 536 Bytes
eadc767 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | """The four NKI kernels of the V-JEPA2 encoder plus their host-side weight packing."""
from .attention_block import (
attention_block,
lane_permutation,
pack_out_weights,
pack_tables,
pack_v_weights,
pack_weights,
)
from .mlp import mlp
from .res_layernorm import res_layernorm_fold
from .rope_tables import build_rope_tables
__all__ = [
"res_layernorm_fold", "attention_block", "mlp",
"pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights",
"lane_permutation", "build_rope_tables",
]
|