| """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", | |
| ] | |