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