code wiki / (root) / nx_vit_encoder_layer.nx

nx_vit_encoder_layer.nx

buildroot/runtime/nx_vit_encoder_layer.nx

4701 B94 linesdepth 8pulls 20 transitivereach 2 importersview sourcekind library
docsdependenciesstructsconstsfunctions

about

nx_vit_encoder_layer.nx -- the faithful f32 PRE-NORM ViT encoder block (the ViTPose repeated unit), composing the gated blocks: h = x + attn_dense(MHSA(LN_before(x))); out = h + fc2(GELU(fc1(LN_after(h)))). MHSA has explicit qkv_bias and HF [out,in] weights (via nx_f32_matmul_t = x@W^T), 12 heads x head_dim 64, scale 1/sqrt(head_dim), non-causal. Weights are passed as a 16-pointer array to stay under the arg cap. license_tier: ORIGINAL

dependencies 6 imports · 2 importers

nx_syscalls.nx nx_f32.nx nx_f32_matmul_t.nx nx_f32_softmax.nx nx_f32_layernorm.nx nx_f32_gelu.nx nx_vit_encoder_layer.nx nx_vit_encoder_layer_gate.nx nx_vitpose_forward.nx

imports: nx_syscalls.nxnx_f32.nxnx_f32_matmul_t.nxnx_f32_softmax.nxnx_f32_layernorm.nxnx_f32_gelu.nx

imported by: nx_vit_encoder_layer_gate.nxnx_vitpose_forward.nx

structs

none

consts

none

functions

13func vit_bias_add(x: *i64, bias: *i64, n_rows: i64, dim: i64) -> i64
called by 2: vit_mhsavit_encoder_layer calls 1: nx_f32_add
20func vit_mhsa(x: *i64, Wq: *i64, bq: *i64, Wk: *i64, bk: *i64, Wv: *i64, bv: *i64, Wo: *i64, bo: *i64,
called by 2: vit_encoder_layermain calls 7: sys_mmapnx_f32_matmul_tvit_bias_addnx_f32_addnx_f32_mulnx_f32_softmax+1
65func vit_encoder_layer(x: *i64, wts: *i64, T: i64, D: i64, nh: i64, hd: i64, mlp: i64, eps: i64, scale: i64, out: *i64) -> i64
called by 2: mainmain calls 8: sys_mmapnx_f32_layernormvit_mhsanx_f32_addnx_f32_matmul_tvit_bias_add+2