ferrox_models/
vision.rs

1//! Kimi K3's MoonViT-V2 vision encoder: patch embedding -> N transformer
2//! encoder layers (RMSNorm, self-attention with 2D RoPE, gated-MLP-free
3//! plain MLP2 feed-forward) -> a patch-merger projector into the text
4//! decoder's embedding space. Genuinely new territory for ferrox, which
5//! has been text-only until now.
6//!
7//! Transcribed directly from real reference source fetched live from the
8//! model repo (not guessed, not derived by analogy to other ViT/RoPE
9//! designs): `moonshotai/Kimi-K3`'s `modeling_kimi_k3.py`
10//! (`MoonVision3dPatchEmbed`, `Rope2DPosEmbRepeated`, `apply_rope`,
11//! `MoonViTEncoderLayer`, `MoonViT3dEncoder`, `tpool_patch_merger`,
12//! `PatchMergerMLPV2`) and its real `config.json`'s `vision_config`
13//! (`norm_type='rmsnorm'`, `activation_func='gelu_pytorch_tanh'`,
14//! `attn_bias=False`, `linear_bias=False`, `patch_embed_proj_bias=False`,
15//! `patch_size=14`, `vt_hidden_size=1024`, `vt_num_attention_heads=12`,
16//! `qkv_hidden_size=1536`, `vt_intermediate_size=4096`,
17//! `vt_num_hidden_layers=27`, `merge_kernel_size=[2,2]`,
18//! `mm_projector_type='patchmergerv2'`, `projector_ln_eps=1e-5`).
19//!
20//! Deliberately scoped, and disclosed as such, to the common case rather
21//! than every real code path:
22//! - Single image, single frame (`t=1`) -- no multi-frame video input,
23//!   so the real code's temporal sincos position-embedding addition and
24//!   the patch-merger's temporal-mean pooling are both no-ops here, not
25//!   implemented as general code paths.
26//! - The caller's patch grid is assumed to exactly match the
27//!   positional-embedding grid it was constructed with, so the real
28//!   code's bicubic/bilinear interpolation branch (`get_rope_shape`, for
29//!   when input resolution differs from the pretrained grid) is never
30//!   exercised and not implemented.
31//! - One image at a time -- no variable-length multi-image batch packing
32//!   (`cu_seqlens`); process a full patch sequence for one image per
33//!   call.
34//! - The real multimodal splicing of vision output into the text
35//!   decoder's input sequence (`KimiK3ForConditionalGeneration.forward`)
36//!   is out of scope here -- this module only implements the vision
37//!   tower itself.
38//!
39//! Two real, non-obvious facts confirmed by reading source rather than
40//! assuming standard ViT/RoPE conventions:
41//!
42//! 1. `MoonVision3dPatchEmbed.forward` reassigns `x` before calling
43//!    `x.size(0)` in the same statement
44//!    (`x = self.proj(x).view(x.size(0), -1)`), but Python evaluates the
45//!    right-hand side of an assignment before rebinding the name, so
46//!    `x.size(0)` refers to the *original*, pre-conv `x` -- meaning the
47//!    real input to this module is already a batch of N individually
48//!    pre-extracted `[C, patch_size, patch_size]` patches (one "image"
49//!    per patch), not one whole image run through a strided conv. A
50//!    `Conv2d(kernel_size=patch_size, stride=patch_size)` applied to an
51//!    input exactly `patch_size` wide/tall is mathematically just one
52//!    linear projection of the flattened patch -- so patch embedding
53//!    here is a single shared `Linear(in_dim*patch_size*patch_size,
54//!    out_dim, bias=false)` applied per patch (`WeightMatrix::apply`),
55//!    not a general strided convolution.
56//! 2. The 2D RoPE (`apply_rope`) rotates *consecutive* real-value pairs
57//!    (`(x[2i], x[2i+1])`, the original RoFormer/GPT-NeoX convention via
58//!    `view_as_complex`), NOT the "rotate half" convention
59//!    (`(x[i], x[i+dim/2])`) `ferrox_core::attention::apply_rope`
60//!    implements for the text decoder's RoPE -- a distinct function is
61//!    needed here, not a reuse of that one. Within each complex pair
62//!    index `j` in `[0, head_dim/2)`, `_precompute_freqs_cis` interleaves
63//!    width- and height-based angles: even `j` uses the patch's width
64//!    coordinate, odd `j` uses its height coordinate (confirmed by
65//!    reading the actual `torch.cat(...).reshape(...)` code, which
66//!    contradicts that same function's own docstring -- code wins).
67//!
68//! Not yet wired into `Decoder` or spliced into the text embedding
69//! sequence. Tested here against synthetic weights, cross-validated
70//! against an independent Python transcription of the same real
71//! algorithm.
72
73use ferrox_core::matmul::rms_norm;
74use ferrox_core::weight_matrix::WeightMatrix;
75
76#[derive(Debug, Clone, Copy)]
77pub struct VisionConfig {
78    pub in_dim: usize,
79    pub patch_size: usize,
80    pub grid_h: usize,
81    pub grid_w: usize,
82    pub hidden_dim: usize,
83    pub num_heads: usize,
84    pub qkv_hidden: usize,
85    pub mlp_dim: usize,
86    pub rms_norm_eps: f32,
87    pub theta_base: f32,
88    pub merge_kh: usize,
89    pub merge_kw: usize,
90    pub projector_ln_eps: f32,
91}
92
93impl VisionConfig {
94    fn head_dim(&self) -> usize {
95        self.qkv_hidden / self.num_heads
96    }
97
98    fn n_patches(&self) -> usize {
99        self.grid_h * self.grid_w
100    }
101}
102
103pub struct VisionEncoderLayerWeights {
104    pub norm0_weight: Vec<f32>,
105    pub wqkv: WeightMatrix, // [3*qkv_hidden, hidden_dim]
106    pub wo: WeightMatrix,   // [hidden_dim, qkv_hidden]
107    pub norm1_weight: Vec<f32>,
108    pub fc0: WeightMatrix, // [mlp_dim, hidden_dim]
109    pub fc1: WeightMatrix, // [hidden_dim, mlp_dim]
110}
111
112pub struct VisionEncoderWeights {
113    pub patch_embed: WeightMatrix, // [hidden_dim, in_dim*patch_size*patch_size]
114    /// Learned per-position embedding, row-major `[grid_h, grid_w,
115    /// hidden_dim]` flattened, added directly to each patch's embedding
116    /// (the "grid matches pretrained size, t=1" scoped case -- see
117    /// module doc comment).
118    pub pos_emb: Vec<f32>,
119    pub layers: Vec<VisionEncoderLayerWeights>,
120    pub final_norm_weight: Vec<f32>,
121}
122
123pub struct VisionMergerWeights {
124    pub proj0: WeightMatrix,        // [merge_hidden, merge_hidden]
125    pub proj1: WeightMatrix,        // [text_hidden, merge_hidden]
126    pub post_norm_weight: Vec<f32>, // [text_hidden]
127}
128
129fn gelu_tanh(x: f32) -> f32 {
130    0.5 * x * (1.0 + (0.797_884_6 * (x + 0.044715 * x * x * x)).tanh())
131}
132
133/// Abramowitz & Stegun 7.1.26 approximation, max absolute error ~1.5e-7
134/// -- std has no built-in `erf`, and this is far more precise than the
135/// 1e-3 tolerance any test here needs.
136fn erf(x: f32) -> f32 {
137    let sign = if x < 0.0 { -1.0 } else { 1.0 };
138    let x = x.abs();
139    let t = 1.0 / (1.0 + 0.3275911 * x);
140    let poly = ((((1.061_405_4 * t - 1.453_152_1) * t) + 1.421_413_8) * t - 0.284_496_72) * t
141        + 0.254_829_6;
142    sign * (1.0 - poly * t * (-x * x).exp())
143}
144
145fn gelu_erf(x: f32) -> f32 {
146    0.5 * x * (1.0 + erf(x / std::f32::consts::SQRT_2))
147}
148
149/// Precomputes `cos`/`sin` for every patch position's `head_dim/2`
150/// rotation angles, per `Rope2DPosEmbRepeated._precompute_freqs_cis`:
151/// row-major `(h, w)` position `p -> (h = p / grid_w, w = p % grid_w)`,
152/// pair index `j` even uses the width coordinate, odd uses height, both
153/// sharing frequency `theta_base^(-4*(j/2)/head_dim)`.
154fn precompute_rope_2d(cfg: &VisionConfig) -> (Vec<f32>, Vec<f32>) {
155    let head_dim = cfg.head_dim();
156    assert_eq!(head_dim % 4, 0, "vision head_dim must be divisible by 4");
157    let num_pairs = head_dim / 2;
158    let n = cfg.n_patches();
159
160    let mut cos_t = vec![0f32; n * num_pairs];
161    let mut sin_t = vec![0f32; n * num_pairs];
162    for p in 0..n {
163        let x_pos = (p % cfg.grid_w) as f32;
164        let y_pos = (p / cfg.grid_w) as f32;
165        for j in 0..num_pairs {
166            let i = j / 2;
167            let freq = cfg.theta_base.powf(-4.0 * i as f32 / head_dim as f32);
168            let angle = if j % 2 == 0 {
169                x_pos * freq
170            } else {
171                y_pos * freq
172            };
173            cos_t[p * num_pairs + j] = angle.cos();
174            sin_t[p * num_pairs + j] = angle.sin();
175        }
176    }
177    (cos_t, sin_t)
178}
179
180/// Applies the 2D RoPE rotation in place to `x` (`[n_patches, num_heads,
181/// head_dim]` flattened), the consecutive-pair convention -- see module
182/// doc comment point 2.
183fn apply_rope_2d(x: &mut [f32], cfg: &VisionConfig, cos_t: &[f32], sin_t: &[f32]) {
184    let head_dim = cfg.head_dim();
185    let num_pairs = head_dim / 2;
186    for p in 0..cfg.n_patches() {
187        for h in 0..cfg.num_heads {
188            let base = (p * cfg.num_heads + h) * head_dim;
189            for j in 0..num_pairs {
190                let c = cos_t[p * num_pairs + j];
191                let s = sin_t[p * num_pairs + j];
192                let a = x[base + 2 * j];
193                let b = x[base + 2 * j + 1];
194                x[base + 2 * j] = a * c - b * s;
195                x[base + 2 * j + 1] = a * s + b * c;
196            }
197        }
198    }
199}
200
201/// Embeds a full image's patches (already extracted as flattened
202/// `in_dim*patch_size*patch_size` pixel vectors, row-major `(h, w)`
203/// order) and adds the learned position embedding. Returns
204/// `[n_patches, hidden_dim]` flattened.
205pub fn embed_patches(
206    weights: &VisionEncoderWeights,
207    cfg: &VisionConfig,
208    patches: &[Vec<f32>],
209) -> Vec<f32> {
210    assert_eq!(patches.len(), cfg.n_patches());
211    let mut out = Vec::with_capacity(cfg.n_patches() * cfg.hidden_dim);
212    for (p, patch) in patches.iter().enumerate() {
213        let embedded = weights.patch_embed.apply(patch);
214        let pos = &weights.pos_emb[p * cfg.hidden_dim..(p + 1) * cfg.hidden_dim];
215        for (e, po) in embedded.iter().zip(pos.iter()) {
216            out.push(e + po);
217        }
218    }
219    out
220}
221
222/// Full (non-causal) self-attention over every patch at once -- the
223/// whole point of a bidirectional vision encoder, unlike the causal
224/// decode-step primitives in `ferrox_core::attention`.
225fn full_self_attention(
226    q: &[f32],
227    k: &[f32],
228    v: &[f32],
229    n_patches: usize,
230    num_heads: usize,
231    head_dim: usize,
232) -> Vec<f32> {
233    let scale = 1.0 / (head_dim as f32).sqrt();
234    let mut out = vec![0f32; n_patches * num_heads * head_dim];
235    for h in 0..num_heads {
236        for i in 0..n_patches {
237            let q_i = &q[(i * num_heads + h) * head_dim..(i * num_heads + h + 1) * head_dim];
238            let mut scores = vec![0f32; n_patches];
239            for t in 0..n_patches {
240                let k_t = &k[(t * num_heads + h) * head_dim..(t * num_heads + h + 1) * head_dim];
241                let dot: f32 = q_i.iter().zip(k_t.iter()).map(|(a, b)| a * b).sum();
242                scores[t] = dot * scale;
243            }
244            let max = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
245            let mut sum = 0f32;
246            for s in scores.iter_mut() {
247                *s = (*s - max).exp();
248                sum += *s;
249            }
250            for s in scores.iter_mut() {
251                *s /= sum;
252            }
253            let out_base = (i * num_heads + h) * head_dim;
254            for t in 0..n_patches {
255                let v_t = &v[(t * num_heads + h) * head_dim..(t * num_heads + h + 1) * head_dim];
256                let w = scores[t];
257                for d in 0..head_dim {
258                    out[out_base + d] += w * v_t[d];
259                }
260            }
261        }
262    }
263    out
264}
265
266fn encoder_layer_forward(
267    layer: &VisionEncoderLayerWeights,
268    cfg: &VisionConfig,
269    x: &[f32],
270    cos_t: &[f32],
271    sin_t: &[f32],
272) -> Vec<f32> {
273    let n = cfg.n_patches();
274    let head_dim = cfg.head_dim();
275
276    let mut residual = x.to_vec();
277    let mut normed = Vec::with_capacity(n * cfg.hidden_dim);
278    for p in 0..n {
279        normed.extend(rms_norm(
280            &x[p * cfg.hidden_dim..(p + 1) * cfg.hidden_dim],
281            &layer.norm0_weight,
282            cfg.rms_norm_eps,
283        ));
284    }
285
286    let mut q = vec![0f32; n * cfg.qkv_hidden];
287    let mut k = vec![0f32; n * cfg.qkv_hidden];
288    let mut v = vec![0f32; n * cfg.qkv_hidden];
289    for p in 0..n {
290        let qkv = layer
291            .wqkv
292            .apply(&normed[p * cfg.hidden_dim..(p + 1) * cfg.hidden_dim]);
293        q[p * cfg.qkv_hidden..(p + 1) * cfg.qkv_hidden].copy_from_slice(&qkv[0..cfg.qkv_hidden]);
294        k[p * cfg.qkv_hidden..(p + 1) * cfg.qkv_hidden]
295            .copy_from_slice(&qkv[cfg.qkv_hidden..2 * cfg.qkv_hidden]);
296        v[p * cfg.qkv_hidden..(p + 1) * cfg.qkv_hidden]
297            .copy_from_slice(&qkv[2 * cfg.qkv_hidden..3 * cfg.qkv_hidden]);
298    }
299
300    apply_rope_2d(&mut q, cfg, cos_t, sin_t);
301    apply_rope_2d(&mut k, cfg, cos_t, sin_t);
302
303    let attn_out = full_self_attention(&q, &k, &v, n, cfg.num_heads, head_dim);
304
305    for p in 0..n {
306        let attn_flat = &attn_out[p * cfg.qkv_hidden..(p + 1) * cfg.qkv_hidden];
307        let projected = layer.wo.apply(attn_flat);
308        let res = &mut residual[p * cfg.hidden_dim..(p + 1) * cfg.hidden_dim];
309        for (r, pr) in res.iter_mut().zip(projected.iter()) {
310            *r += pr;
311        }
312    }
313
314    let mut out = residual.clone();
315    for p in 0..n {
316        let res_p = &residual[p * cfg.hidden_dim..(p + 1) * cfg.hidden_dim];
317        let normed1 = rms_norm(res_p, &layer.norm1_weight, cfg.rms_norm_eps);
318        let mut hidden = layer.fc0.apply(&normed1);
319        for h in hidden.iter_mut() {
320            *h = gelu_tanh(*h);
321        }
322        let mlp_out = layer.fc1.apply(&hidden);
323        let out_p = &mut out[p * cfg.hidden_dim..(p + 1) * cfg.hidden_dim];
324        for (o, m) in out_p.iter_mut().zip(mlp_out.iter()) {
325            *o += m;
326        }
327    }
328
329    out
330}
331
332/// Runs the full encoder (patch embed + every layer + final norm) for
333/// one image. Returns `[n_patches, hidden_dim]` flattened.
334pub fn encoder_forward(
335    weights: &VisionEncoderWeights,
336    cfg: &VisionConfig,
337    patches: &[Vec<f32>],
338) -> Vec<f32> {
339    let (cos_t, sin_t) = precompute_rope_2d(cfg);
340    let mut x = embed_patches(weights, cfg, patches);
341    for layer in &weights.layers {
342        x = encoder_layer_forward(layer, cfg, &x, &cos_t, &sin_t);
343    }
344    let n = cfg.n_patches();
345    let mut out = Vec::with_capacity(n * cfg.hidden_dim);
346    for p in 0..n {
347        out.extend(rms_norm(
348            &x[p * cfg.hidden_dim..(p + 1) * cfg.hidden_dim],
349            &weights.final_norm_weight,
350            cfg.rms_norm_eps,
351        ));
352    }
353    out
354}
355
356/// `tpool_patch_merger` for the single-frame (`t=1`) case: groups
357/// spatially-adjacent `merge_kh x merge_kw` patches (row-major `(h, w)`)
358/// into one merged patch each, without the temporal mean (a no-op at
359/// `t=1`). Returns `[num_merged, merge_kh*merge_kw*hidden_dim]`
360/// flattened -- already contiguous per merged patch, ready for the
361/// projector.
362pub fn patch_merge(encoder_out: &[f32], cfg: &VisionConfig) -> Vec<f32> {
363    assert_eq!(cfg.grid_h % cfg.merge_kh, 0);
364    assert_eq!(cfg.grid_w % cfg.merge_kw, 0);
365    let new_h = cfg.grid_h / cfg.merge_kh;
366    let new_w = cfg.grid_w / cfg.merge_kw;
367    let merge_block = cfg.merge_kh * cfg.merge_kw;
368    let mut out = vec![0f32; new_h * new_w * merge_block * cfg.hidden_dim];
369
370    for nh in 0..new_h {
371        for nw in 0..new_w {
372            for kh in 0..cfg.merge_kh {
373                for kw in 0..cfg.merge_kw {
374                    let orig_h = nh * cfg.merge_kh + kh;
375                    let orig_w = nw * cfg.merge_kw + kw;
376                    let orig_patch = orig_h * cfg.grid_w + orig_w;
377                    let merged_idx = nh * new_w + nw;
378                    let sub_idx = kh * cfg.merge_kw + kw;
379                    let src = &encoder_out
380                        [orig_patch * cfg.hidden_dim..(orig_patch + 1) * cfg.hidden_dim];
381                    let dst_base = (merged_idx * merge_block + sub_idx) * cfg.hidden_dim;
382                    out[dst_base..dst_base + cfg.hidden_dim].copy_from_slice(src);
383                }
384            }
385        }
386    }
387    out
388}
389
390/// `PatchMergerMLPV2`: per merged patch, `Linear -> GELU (erf) -> Linear
391/// -> RMSNorm`. `merged` is `[num_merged, merge_kh*merge_kw*hidden_dim]`
392/// flattened (`patch_merge`'s output); returns `[num_merged,
393/// text_hidden]` flattened.
394pub fn project_merged_patches(
395    weights: &VisionMergerWeights,
396    cfg: &VisionConfig,
397    merged: &[f32],
398    num_merged: usize,
399) -> Vec<f32> {
400    let merge_hidden = cfg.merge_kh * cfg.merge_kw * cfg.hidden_dim;
401    let text_hidden = weights.post_norm_weight.len();
402    let mut out = Vec::with_capacity(num_merged * text_hidden);
403    for m in 0..num_merged {
404        let input = &merged[m * merge_hidden..(m + 1) * merge_hidden];
405        let mut hidden = weights.proj0.apply(input);
406        for h in hidden.iter_mut() {
407            *h = gelu_erf(*h);
408        }
409        let projected = weights.proj1.apply(&hidden);
410        out.extend(rms_norm(
411            &projected,
412            &weights.post_norm_weight,
413            cfg.projector_ln_eps,
414        ));
415    }
416    out
417}
418
419#[cfg(test)]
420mod tests {
421    use super::*;
422    use ferrox_core::tensor::Tensor;
423
424    const IN_DIM: usize = 3;
425    const PATCH_SIZE: usize = 2;
426    const GRID_H: usize = 2;
427    const GRID_W: usize = 2;
428    const HIDDEN_DIM: usize = 8;
429    const NUM_HEADS: usize = 2;
430    const QKV_HIDDEN: usize = 8;
431    const MLP_DIM: usize = 12;
432    const NORM_EPS: f32 = 1e-5;
433    const THETA_BASE: f32 = 10000.0;
434    const MERGE_KH: usize = 2;
435    const MERGE_KW: usize = 2;
436    const TEXT_HIDDEN: usize = 6;
437    const PROJECTOR_LN_EPS: f32 = 1e-5;
438
439    // Generated by an independent Python reference -- do not hand-edit.
440    const VIS_PATCH_0: [f32; 12] = [
441        -0.396561, 0.120286, -0.948163, 0.697886, 0.319147, -0.146024, -0.155975, 0.151918,
442        -0.13383, -0.112954, 0.360034, 0.257353,
443    ];
444    const VIS_PATCH_1: [f32; 12] = [
445        -0.032064, -0.0427383, 0.0804582, -0.307009, -0.201875, 0.27413, -0.0652414, -0.687213,
446        -0.238639, 0.328311, -0.116141, -0.0743664,
447    ];
448    const VIS_PATCH_2: [f32; 12] = [
449        0.320918, 0.912305, -0.356594, 0.674103, -0.615006, 0.0874888, -0.584765, 0.675729,
450        0.416961, 0.568858, -0.442767, 0.342278,
451    ];
452    const VIS_PATCH_3: [f32; 12] = [
453        -0.259507, -0.228693, 0.253269, 0.438359, 0.10221, -0.313994, -0.412908, 0.722158,
454        0.296973, 0.359864, 1.09174, -0.407931,
455    ];
456    const VIS_PATCH_EMBED_W: [f32; 96] = [
457        0.767837,
458        0.945273,
459        0.485524,
460        0.248132,
461        -0.199147,
462        0.298346,
463        -0.132808,
464        -0.00649524,
465        -0.08713,
466        0.085149,
467        0.386423,
468        -0.166675,
469        -0.295621,
470        -0.300887,
471        -0.290483,
472        -0.429331,
473        -0.27388,
474        0.38798,
475        -0.177994,
476        0.0771323,
477        -0.365068,
478        0.0508952,
479        -0.522234,
480        -0.209627,
481        0.676362,
482        -0.174891,
483        0.335993,
484        0.136516,
485        -0.0458957,
486        -0.195632,
487        0.386073,
488        -0.053215,
489        0.458227,
490        -0.215724,
491        0.0172017,
492        0.13965,
493        0.111948,
494        -0.37014,
495        -0.199219,
496        -0.0587941,
497        -0.25611,
498        0.203198,
499        0.176406,
500        -0.587125,
501        -0.541576,
502        -0.384469,
503        0.0351779,
504        0.609952,
505        -0.114707,
506        0.0751952,
507        -0.318934,
508        -0.314051,
509        -0.587168,
510        -0.00850327,
511        0.284165,
512        -0.107014,
513        0.418936,
514        0.0593565,
515        -0.0109237,
516        0.155176,
517        0.146213,
518        0.344342,
519        -0.240586,
520        -0.686415,
521        0.034424,
522        -0.183503,
523        -0.00815316,
524        0.49929,
525        -0.330853,
526        0.229439,
527        0.28376,
528        0.138221,
529        0.335552,
530        -0.137509,
531        -0.204453,
532        0.311695,
533        0.21609,
534        0.417113,
535        0.0636869,
536        0.487069,
537        -0.0846241,
538        -0.318099,
539        -0.611127,
540        -0.33042,
541        0.245551,
542        -0.439479,
543        -0.135709,
544        0.631975,
545        0.254039,
546        0.537295,
547        -0.297469,
548        -0.737993,
549        0.454914,
550        -0.425083,
551        0.0272771,
552        0.065128,
553    ];
554    const VIS_POS_EMB_W: [f32; 32] = [
555        -0.189779,
556        0.310527,
557        0.309868,
558        0.110289,
559        -0.00589813,
560        0.0382874,
561        -0.219182,
562        -0.0478128,
563        -0.0209866,
564        -0.170447,
565        0.169843,
566        -0.285296,
567        -0.0906838,
568        -0.450576,
569        0.0993637,
570        0.148521,
571        0.124258,
572        0.583617,
573        0.187958,
574        -0.188398,
575        0.532708,
576        -0.184064,
577        0.157748,
578        0.148626,
579        -0.175945,
580        0.27907,
581        0.0709187,
582        0.114191,
583        -0.135147,
584        0.260527,
585        -0.172188,
586        -0.355832,
587    ];
588    const VIS_L0_NORM0_W: [f32; 8] = [
589        1.14251, 0.957327, 1.00434, 1.19228, 1.11365, 1.06327, 1.10699, 0.927249,
590    ];
591    const VIS_L0_WQKV: [f32; 192] = [
592        0.225927, 0.332981, 0.226627, -0.0612404, -1.12422, 0.167191, 0.0344629, -0.110866,
593        0.23343, 0.0797109, 0.0611867, -0.445077, -0.437797, 0.111161, 0.0228171, 0.0416223,
594        0.0498789, 0.123437, 0.0108197, 0.0473313, 0.0768343, 0.278823, -0.154407, 0.235919,
595        0.190915, 0.0316206, -0.142534, 0.110742, 0.259192, -0.322493, 0.0186437, 0.0194232,
596        -0.198361, -0.185786, -0.465344, -0.330938, -0.461243, -0.184972, -0.167996, -0.162386,
597        0.111553, -0.05627, -0.218674, -0.216565, -0.339453, -0.0588504, -0.0642164, 0.486887,
598        0.460246, 0.443371, 0.623773, -0.36512, 0.201895, 0.228316, 0.142865, 0.638001, 0.761363,
599        0.29381, 0.195944, -0.336719, -0.336108, 0.210067, -0.0226898, -0.399168, -0.0372073,
600        0.108002, 0.280778, -0.0323038, -0.0807652, 0.00186235, 0.051895, -0.34894, 0.249993,
601        0.657098, 0.24426, -0.434459, 0.0283526, -0.529688, -0.548764, 0.197775, 0.0232947,
602        -0.206133, -0.626234, -0.38068, 0.0486588, 0.00956812, 0.20691, 0.457018, -0.101777,
603        0.14069, -0.0413395, -0.148962, -0.0734372, 0.288226, -0.308526, 0.342312, 0.443186,
604        0.0107199, 0.0764755, -0.42354, 0.70816, 0.293666, 0.0516828, 0.0172313, 0.135292,
605        -0.195371, 0.0849785, -0.277073, 0.149196, -0.203464, 0.268656, -0.0361107, 0.0806238,
606        0.888267, 0.24127, 0.0803401, 0.0133166, -0.311749, 0.325266, 0.480702, 0.0193065,
607        0.503025, -0.0336833, -0.531916, -0.195972, -0.317098, -0.082362, 0.188094, -0.278692,
608        -0.0240735, 0.0598056, -0.219908, -0.272796, 0.0261548, 0.261711, -0.180968, -0.225418,
609        -0.042376, 0.0869846, -0.162905, -0.100244, -0.0693806, -0.325737, -0.0532611, -0.168014,
610        0.475017, -0.155197, 0.137133, 0.113856, -0.33541, -0.374685, -0.133646, 0.198772,
611        -0.171972, -0.196846, -0.870066, -0.119052, 0.0937437, 0.133406, 0.128758, -0.0983533,
612        -0.0681474, 0.249683, -0.0407133, -0.690913, 0.306003, -0.323304, 0.412587, 0.508244,
613        0.228732, 0.271235, 0.55734, -0.330637, 0.0761556, 0.0717198, -0.290133, 0.603544,
614        -0.174208, -0.219374, 0.0803143, 0.309335, -0.180314, -0.17987, 0.151233, -0.427684,
615        -0.231776, -0.251471, -0.1303, -0.0426799, 0.553263, 0.272208, 0.00671946,
616    ];
617    const VIS_L0_WO: [f32; 64] = [
618        -0.0674106,
619        0.262305,
620        0.12852,
621        0.267363,
622        0.414293,
623        -0.302455,
624        -0.675857,
625        -0.00654075,
626        0.0238384,
627        0.0860337,
628        -0.179647,
629        -0.0697861,
630        0.167621,
631        0.183768,
632        0.399294,
633        -0.649884,
634        -0.494598,
635        0.438582,
636        0.18246,
637        0.23803,
638        -0.0854503,
639        -0.142247,
640        0.0599411,
641        -0.281936,
642        -0.459049,
643        -0.333171,
644        0.0349579,
645        0.191157,
646        0.495378,
647        0.111546,
648        0.141755,
649        0.506442,
650        -0.232997,
651        -0.287787,
652        -0.16138,
653        -0.0288672,
654        0.606971,
655        -0.104652,
656        0.392134,
657        -0.141798,
658        0.147425,
659        0.204395,
660        0.159886,
661        0.225314,
662        -0.206712,
663        0.436844,
664        0.294298,
665        -0.210564,
666        0.151729,
667        -0.321948,
668        0.0284701,
669        -0.351326,
670        0.21551,
671        0.889311,
672        0.441252,
673        0.219234,
674        -0.270861,
675        0.114077,
676        0.355367,
677        -0.0265149,
678        0.0387836,
679        -0.508303,
680        -0.256626,
681        -0.554506,
682    ];
683    const VIS_L0_NORM1_W: [f32; 8] = [
684        1.19145, 1.04726, 1.07247, 0.97973, 0.853933, 1.0176, 0.936351, 0.903017,
685    ];
686    const VIS_L0_FC0: [f32; 96] = [
687        0.0166233, 0.0329455, -0.0541478, -0.232281, -0.29576, 0.00189617, 0.366746, -0.481293,
688        -0.197155, -0.0487075, -0.0476214, -0.177829, 0.578552, 0.11994, -0.295615, -0.564639,
689        -0.051645, 0.443841, 0.0578301, -0.694797, -0.413083, 0.0368751, 0.320283, 0.0423223,
690        0.044788, 0.321899, 0.376876, -0.202959, 0.131414, -0.287215, -0.322971, 0.176502,
691        -0.427032, 0.0509592, 0.0942597, 0.258386, -0.345129, 0.114245, -0.0376285, 0.445818,
692        -0.222826, -0.333006, 0.150066, -0.256744, -0.388289, 0.273022, 0.129116, -0.933427,
693        -0.276295, -0.332411, 0.193304, 0.441769, 0.617435, 0.211838, -0.328804, -0.339382,
694        -0.285821, 0.146822, 0.379626, -0.575327, 0.182773, 0.278357, 0.0947908, 0.97087, 0.160575,
695        0.00974755, 0.183156, -0.163253, -0.32978, 0.126654, 0.346088, -0.0704812, 0.277716,
696        -0.112732, 0.489417, 0.0644369, 0.160294, 0.22985, -0.405623, -0.0546382, 0.552158,
697        0.0109478, 0.420782, 0.265227, -0.548319, 0.212211, -0.326486, 0.0783693, -0.12468,
698        0.00652852, -0.0212635, 0.0728004, -0.134494, -0.405292, -0.302986, 0.274872,
699    ];
700    const VIS_L0_FC1: [f32; 96] = [
701        0.517348,
702        0.0697715,
703        0.252689,
704        0.293056,
705        -0.409747,
706        -0.561795,
707        -0.0153769,
708        -0.567693,
709        -0.0817456,
710        -0.443467,
711        -0.187211,
712        -0.225356,
713        0.155999,
714        -0.194173,
715        -0.344326,
716        0.0146415,
717        0.890155,
718        0.147833,
719        -0.583478,
720        0.0273942,
721        0.771329,
722        0.0507425,
723        0.3842,
724        0.333034,
725        0.356512,
726        0.315238,
727        0.0847949,
728        -0.0552949,
729        -0.0835648,
730        0.0360814,
731        0.12811,
732        -0.382702,
733        -0.0841374,
734        -0.0121686,
735        0.0323577,
736        -0.0929058,
737        0.241921,
738        -0.0311971,
739        -0.112632,
740        -0.371058,
741        -0.23322,
742        0.0688837,
743        0.0357527,
744        -0.611612,
745        0.277291,
746        -0.210783,
747        -0.270577,
748        0.334413,
749        0.349005,
750        -0.213423,
751        -0.106685,
752        -0.367216,
753        0.244268,
754        0.271403,
755        -0.131644,
756        0.417505,
757        0.275575,
758        -0.0371131,
759        0.216475,
760        -0.496402,
761        -0.0788385,
762        0.075497,
763        -0.00312696,
764        0.614985,
765        0.0931599,
766        0.109938,
767        -0.398398,
768        -0.0562463,
769        -0.235057,
770        0.70463,
771        -0.124689,
772        0.0357173,
773        0.364742,
774        0.197,
775        -0.423393,
776        0.191556,
777        0.213094,
778        -0.0544942,
779        0.139289,
780        0.0314193,
781        -0.415614,
782        -0.16905,
783        0.00621485,
784        -0.175932,
785        0.161714,
786        0.300591,
787        0.082026,
788        -0.111803,
789        0.408852,
790        0.041258,
791        0.106057,
792        -0.229034,
793        -0.0652944,
794        0.299431,
795        0.223119,
796        -0.0561207,
797    ];
798    const VIS_L1_NORM0_W: [f32; 8] = [
799        1.00082, 1.0199, 0.905881, 0.981468, 0.948844, 0.861423, 1.01522, 0.98562,
800    ];
801    const VIS_L1_WQKV: [f32; 192] = [
802        -0.439142,
803        -0.238502,
804        -0.0261269,
805        -0.867316,
806        0.055201,
807        0.080017,
808        0.19467,
809        -0.178583,
810        -0.426824,
811        0.188486,
812        -0.412493,
813        0.162284,
814        -0.202711,
815        -0.303652,
816        0.0883011,
817        0.336673,
818        -0.047082,
819        -0.0737218,
820        -0.506555,
821        0.249362,
822        0.0218581,
823        -0.0815053,
824        0.392514,
825        -0.470229,
826        -0.151384,
827        0.00145506,
828        -0.0886644,
829        -0.73097,
830        0.104801,
831        -0.00915,
832        0.30323,
833        0.61196,
834        -0.0563681,
835        0.115342,
836        0.075918,
837        0.0509361,
838        -0.0414925,
839        0.115973,
840        0.338449,
841        -0.254719,
842        0.570279,
843        0.151889,
844        0.0908936,
845        -0.674739,
846        -0.00318789,
847        -0.292769,
848        0.172313,
849        0.286671,
850        0.344964,
851        -0.049921,
852        -0.0223788,
853        0.131137,
854        -0.62942,
855        -0.130129,
856        -0.162834,
857        0.156245,
858        -0.137213,
859        0.223168,
860        -0.0578504,
861        0.256585,
862        0.665501,
863        0.467405,
864        0.284134,
865        0.343418,
866        -0.109782,
867        -0.231188,
868        -0.253781,
869        0.0587258,
870        0.44202,
871        -0.218814,
872        0.1406,
873        -0.374162,
874        -0.364904,
875        0.257287,
876        -0.206018,
877        -0.214628,
878        0.592636,
879        0.0206818,
880        -0.0133439,
881        -0.0688817,
882        0.124602,
883        -0.212348,
884        0.017195,
885        0.326011,
886        -0.0961918,
887        -0.713652,
888        -0.196473,
889        0.230423,
890        0.336269,
891        0.109072,
892        0.428225,
893        -0.311886,
894        0.383659,
895        0.0629309,
896        0.20279,
897        -0.858124,
898        0.267745,
899        0.0698875,
900        0.572496,
901        0.0564592,
902        -0.291992,
903        -0.240974,
904        0.121683,
905        0.0923392,
906        -0.190241,
907        -0.320404,
908        0.454806,
909        -0.0716024,
910        -0.020555,
911        -0.160004,
912        0.419512,
913        0.334638,
914        -0.306493,
915        -0.0229742,
916        -0.276633,
917        0.556304,
918        0.132983,
919        -0.0840649,
920        -0.101249,
921        0.163251,
922        -0.0738091,
923        0.462798,
924        -0.27786,
925        -0.370396,
926        0.423145,
927        -0.127708,
928        0.425961,
929        0.263961,
930        -0.407428,
931        -0.182851,
932        0.014795,
933        -0.0661325,
934        0.0908633,
935        -0.183748,
936        -0.0254787,
937        0.300446,
938        0.399875,
939        -0.237013,
940        -0.0695918,
941        -0.0380361,
942        -0.404956,
943        0.281397,
944        0.169411,
945        0.156416,
946        -0.0330841,
947        0.2611,
948        0.289024,
949        0.348009,
950        -0.0510249,
951        0.104812,
952        -0.116636,
953        -0.0968812,
954        0.294442,
955        0.0161214,
956        -0.0461076,
957        0.180963,
958        0.199651,
959        -0.226287,
960        -0.19431,
961        0.21369,
962        0.157819,
963        -0.0628762,
964        0.358757,
965        -0.0297353,
966        -0.579409,
967        -0.238768,
968        -0.442426,
969        0.343682,
970        0.151564,
971        0.0356123,
972        -0.695277,
973        -0.443175,
974        -0.0452226,
975        -0.534062,
976        0.282538,
977        -0.0308858,
978        -0.0429336,
979        -0.019023,
980        -0.285045,
981        -0.312946,
982        -0.569596,
983        0.00704258,
984        0.121878,
985        -0.0712807,
986        -0.0847591,
987        -0.282027,
988        0.0752057,
989        -0.735423,
990        -0.331334,
991        -0.197524,
992        -0.152546,
993        0.0653733,
994    ];
995    const VIS_L1_WO: [f32; 64] = [
996        -0.0577218, 0.0466404, 0.373732, 0.411631, 0.194282, 0.37458, 0.518893, 0.111126,
997        -0.419491, 0.301584, -0.298475, 0.558556, 0.517514, -0.155468, -0.329935, -0.206597,
998        -0.0195243, 0.384379, -0.370628, 0.26457, -0.620897, 0.364109, -0.018589, 0.0599231,
999        -0.0318378, 0.280722, 0.111394, 0.50638, -0.532875, -0.271522, -0.403808, 0.215335,
1000        -0.506274, 0.729307, 0.0275459, 0.366571, -0.289899, 0.0588975, 0.187275, -0.211902,
1001        -0.0683283, 0.467214, 0.379918, 0.173689, 0.591409, -0.58374, 0.08737, -0.114411, 0.427121,
1002        0.332962, 0.350832, -0.179064, 0.16847, 0.324073, 0.389196, -0.277921, -0.204972, 0.073562,
1003        0.418459, 0.186085, 0.412741, -0.58574, -0.0814907, 0.149653,
1004    ];
1005    const VIS_L1_NORM1_W: [f32; 8] = [
1006        1.02696, 0.810868, 1.11009, 1.03453, 0.953349, 1.11978, 0.827073, 0.948868,
1007    ];
1008    const VIS_L1_FC0: [f32; 96] = [
1009        0.265211,
1010        0.112798,
1011        -0.0674063,
1012        -0.0461329,
1013        -0.070063,
1014        0.260172,
1015        -0.0912916,
1016        0.203377,
1017        -0.626122,
1018        0.368826,
1019        0.732467,
1020        0.149894,
1021        -0.0259573,
1022        -0.160522,
1023        0.20993,
1024        -0.195052,
1025        -0.24562,
1026        -0.638671,
1027        0.0727824,
1028        -0.0872477,
1029        -0.0707903,
1030        -0.000692914,
1031        -0.212435,
1032        0.0562055,
1033        -0.323977,
1034        -0.034617,
1035        -0.0844252,
1036        -0.512489,
1037        -0.363591,
1038        -0.177747,
1039        -0.168726,
1040        0.0777597,
1041        -0.121498,
1042        -0.0153486,
1043        -0.122644,
1044        -0.777041,
1045        -0.0370188,
1046        -0.252651,
1047        -0.0479536,
1048        -0.229824,
1049        0.196703,
1050        -0.0137741,
1051        0.212481,
1052        0.234418,
1053        -0.169992,
1054        0.244582,
1055        0.0686207,
1056        -0.317385,
1057        -0.0347921,
1058        0.0282208,
1059        0.380975,
1060        -0.10457,
1061        0.471259,
1062        0.230964,
1063        -0.634277,
1064        -0.310673,
1065        0.175964,
1066        -0.4366,
1067        0.142855,
1068        0.414785,
1069        -0.0804484,
1070        0.422308,
1071        0.39718,
1072        -0.0169027,
1073        -0.0100078,
1074        0.0543716,
1075        0.133468,
1076        -0.467238,
1077        -0.0434751,
1078        -0.501976,
1079        -0.10354,
1080        -0.0658216,
1081        -0.545229,
1082        -0.379061,
1083        -0.119427,
1084        0.648104,
1085        -0.416021,
1086        0.0373657,
1087        -0.268065,
1088        -0.112513,
1089        0.295553,
1090        0.0931424,
1091        0.127734,
1092        0.146628,
1093        0.17265,
1094        -0.64847,
1095        0.115814,
1096        -0.0255255,
1097        0.119235,
1098        0.250913,
1099        0.124699,
1100        0.207121,
1101        -0.351437,
1102        0.134342,
1103        -0.0327808,
1104        -0.0259479,
1105    ];
1106    const VIS_L1_FC1: [f32; 96] = [
1107        0.39004,
1108        -0.146601,
1109        -0.199982,
1110        -0.39407,
1111        -0.0886803,
1112        0.321206,
1113        -0.500337,
1114        0.0753654,
1115        -0.759829,
1116        -0.207078,
1117        0.131248,
1118        -0.161568,
1119        -0.360611,
1120        -0.11047,
1121        0.296733,
1122        0.447581,
1123        0.453663,
1124        0.204687,
1125        0.108174,
1126        0.295752,
1127        -0.346912,
1128        -0.290761,
1129        0.387467,
1130        0.256882,
1131        0.262525,
1132        0.0665477,
1133        -0.247498,
1134        -0.324031,
1135        0.281414,
1136        0.324848,
1137        -0.336142,
1138        -0.221854,
1139        -0.0308679,
1140        0.262391,
1141        -0.102162,
1142        -0.255668,
1143        0.296526,
1144        0.179182,
1145        -0.208277,
1146        -0.197806,
1147        0.332993,
1148        -0.384262,
1149        -0.0878628,
1150        0.134792,
1151        0.0493159,
1152        0.278824,
1153        -0.513304,
1154        0.00528268,
1155        0.0903486,
1156        0.100079,
1157        0.173074,
1158        -0.35006,
1159        0.433543,
1160        0.136213,
1161        0.168189,
1162        -0.17306,
1163        0.054225,
1164        0.399138,
1165        -0.00455385,
1166        0.0586072,
1167        0.0950148,
1168        -0.489918,
1169        -0.0517679,
1170        0.207905,
1171        -0.146564,
1172        -0.0126508,
1173        -0.00982426,
1174        -0.473983,
1175        0.241292,
1176        -0.406213,
1177        0.1268,
1178        -0.040079,
1179        -0.102337,
1180        0.0200342,
1181        0.226104,
1182        0.244038,
1183        -0.129102,
1184        -0.0493391,
1185        -0.257829,
1186        -0.43845,
1187        0.285282,
1188        -0.329176,
1189        -0.22059,
1190        0.113821,
1191        -0.404307,
1192        0.144029,
1193        -0.253124,
1194        -0.340929,
1195        -0.1742,
1196        -0.255316,
1197        0.0708169,
1198        0.425441,
1199        0.202217,
1200        0.139482,
1201        -0.184594,
1202        -0.0723397,
1203    ];
1204    const VIS_FINAL_NORM_W: [f32; 8] = [
1205        1.00924, 0.992412, 1.01246, 0.877821, 1.17394, 0.862922, 0.981503, 1.07865,
1206    ];
1207    const VIS_MERGE_PROJ0: [f32; 1024] = [
1208        -0.236349,
1209        0.302076,
1210        -0.436699,
1211        -0.236767,
1212        -0.719464,
1213        0.295538,
1214        -0.373172,
1215        -0.22321,
1216        0.213612,
1217        0.0864719,
1218        0.302082,
1219        0.176717,
1220        0.249573,
1221        -0.162176,
1222        -0.0439175,
1223        0.203247,
1224        0.207236,
1225        0.265128,
1226        -0.248399,
1227        -0.347564,
1228        -0.395201,
1229        -0.160791,
1230        -0.0491745,
1231        -0.376832,
1232        -0.0740985,
1233        0.170485,
1234        0.277722,
1235        -0.19514,
1236        0.15138,
1237        -0.401887,
1238        -0.217022,
1239        0.0129443,
1240        -0.212729,
1241        -0.252641,
1242        0.0614228,
1243        -0.316426,
1244        0.0141244,
1245        -0.0513168,
1246        0.133343,
1247        0.207095,
1248        0.0205796,
1249        -0.168545,
1250        0.282756,
1251        0.0393554,
1252        -0.283789,
1253        0.0319159,
1254        -0.203635,
1255        -0.777423,
1256        -0.307626,
1257        -0.133649,
1258        -0.157371,
1259        -0.481473,
1260        -0.124941,
1261        0.107979,
1262        -0.362024,
1263        0.0569102,
1264        -0.628088,
1265        0.131413,
1266        -0.190082,
1267        -0.422225,
1268        0.312183,
1269        0.118983,
1270        0.264134,
1271        0.0722628,
1272        0.0832335,
1273        -0.135963,
1274        0.105563,
1275        -0.0461713,
1276        -0.212584,
1277        -0.0778148,
1278        -0.403216,
1279        -0.0291645,
1280        0.110195,
1281        0.0157109,
1282        0.148185,
1283        -0.114513,
1284        -0.0117906,
1285        -0.013211,
1286        0.5547,
1287        -0.410289,
1288        0.581179,
1289        -0.239112,
1290        -0.121669,
1291        0.0614686,
1292        0.151549,
1293        -0.0293001,
1294        0.0450231,
1295        0.393947,
1296        0.753717,
1297        0.101535,
1298        0.187023,
1299        -0.0497537,
1300        0.473786,
1301        -0.0537569,
1302        0.131424,
1303        -0.330517,
1304        -0.81053,
1305        0.655177,
1306        0.297821,
1307        0.253245,
1308        -0.295869,
1309        -0.503227,
1310        -0.0895497,
1311        0.0160632,
1312        0.195476,
1313        -0.0153669,
1314        0.361626,
1315        0.08452,
1316        0.0733994,
1317        0.326634,
1318        -0.216635,
1319        0.0958955,
1320        -0.183129,
1321        -0.21866,
1322        0.430652,
1323        0.306151,
1324        0.140998,
1325        0.370526,
1326        0.0936386,
1327        0.383891,
1328        -0.268531,
1329        -0.0606306,
1330        0.0965548,
1331        -0.315549,
1332        -0.216088,
1333        -0.38473,
1334        0.184519,
1335        -0.153703,
1336        0.0636753,
1337        0.139368,
1338        -0.0627465,
1339        -0.00900365,
1340        0.024932,
1341        0.110283,
1342        0.38189,
1343        -0.341072,
1344        0.0110382,
1345        -0.370518,
1346        0.413105,
1347        -0.557165,
1348        0.0337617,
1349        -0.138768,
1350        -0.185464,
1351        0.19559,
1352        0.141414,
1353        -0.298639,
1354        0.121156,
1355        0.0692044,
1356        -0.236338,
1357        -0.446907,
1358        0.217787,
1359        0.189244,
1360        0.136685,
1361        -0.172439,
1362        0.34132,
1363        1.00006,
1364        0.0300311,
1365        0.13925,
1366        -0.259658,
1367        -0.257112,
1368        0.501067,
1369        0.653685,
1370        -0.0926841,
1371        -0.213994,
1372        -0.29322,
1373        -0.201359,
1374        -0.00520445,
1375        -0.391296,
1376        -0.232173,
1377        0.25891,
1378        -0.0725001,
1379        -0.25552,
1380        -0.372763,
1381        0.293815,
1382        -0.336134,
1383        -0.0673453,
1384        0.272661,
1385        0.1354,
1386        0.0615359,
1387        0.0644006,
1388        -0.307396,
1389        -0.239653,
1390        -0.163541,
1391        -0.88038,
1392        -0.149103,
1393        -0.0547648,
1394        -0.0497764,
1395        -0.0617094,
1396        0.0359132,
1397        0.055568,
1398        0.462535,
1399        0.0217722,
1400        0.230232,
1401        -0.201187,
1402        0.0736835,
1403        0.0634726,
1404        -0.0907999,
1405        -0.532513,
1406        -0.423686,
1407        0.112162,
1408        -0.037038,
1409        -0.0948945,
1410        -0.223742,
1411        0.00604771,
1412        0.150508,
1413        0.535644,
1414        -0.0232303,
1415        -0.0708197,
1416        0.0633776,
1417        0.33081,
1418        0.672843,
1419        0.151136,
1420        -0.445781,
1421        -0.00767956,
1422        -0.354326,
1423        -0.30255,
1424        0.527207,
1425        -0.117645,
1426        -0.226993,
1427        0.397529,
1428        0.0112774,
1429        -0.350571,
1430        0.00682257,
1431        -0.428325,
1432        -0.267424,
1433        0.0460792,
1434        -0.772427,
1435        0.204795,
1436        0.0671937,
1437        0.243194,
1438        -0.125164,
1439        0.11855,
1440        -0.303156,
1441        -0.283589,
1442        0.0239112,
1443        0.00140681,
1444        -0.00935182,
1445        0.398183,
1446        -0.0882406,
1447        0.12846,
1448        0.0970286,
1449        0.0414345,
1450        0.117666,
1451        -0.524796,
1452        -0.171809,
1453        0.324681,
1454        -0.0936075,
1455        0.203437,
1456        0.135693,
1457        -0.917639,
1458        -0.398252,
1459        -0.205078,
1460        -0.209732,
1461        -0.140251,
1462        -0.135229,
1463        -0.250041,
1464        -0.205651,
1465        -0.591604,
1466        0.574865,
1467        0.410002,
1468        0.255556,
1469        -0.0320876,
1470        -0.280266,
1471        0.723274,
1472        -0.290018,
1473        0.421423,
1474        0.0849775,
1475        -0.281182,
1476        -0.252831,
1477        -0.134178,
1478        -0.195518,
1479        -0.432479,
1480        0.118937,
1481        -0.283856,
1482        0.220364,
1483        -0.0620359,
1484        0.0341665,
1485        0.0261671,
1486        -0.124931,
1487        -0.0379213,
1488        0.0432028,
1489        0.68423,
1490        -0.39457,
1491        -0.219883,
1492        -0.125563,
1493        -0.387435,
1494        -0.410432,
1495        -0.476937,
1496        -0.0312687,
1497        -0.175228,
1498        0.0353051,
1499        -0.444102,
1500        -0.283561,
1501        0.257007,
1502        0.208592,
1503        -0.484891,
1504        0.230308,
1505        -0.182565,
1506        -0.350742,
1507        0.199333,
1508        0.250345,
1509        -0.767212,
1510        0.238475,
1511        -0.17154,
1512        -0.180079,
1513        -0.135758,
1514        0.296082,
1515        0.350977,
1516        0.072524,
1517        0.204048,
1518        0.607167,
1519        -0.0381462,
1520        -0.3366,
1521        0.124063,
1522        0.177895,
1523        -0.0799869,
1524        -0.0899268,
1525        0.610806,
1526        -0.109059,
1527        -0.299177,
1528        -0.835396,
1529        0.214294,
1530        0.350627,
1531        -0.380885,
1532        0.144448,
1533        -0.131053,
1534        -0.00930346,
1535        -0.204848,
1536        -0.0314523,
1537        0.231507,
1538        -0.179883,
1539        -0.230056,
1540        -0.328797,
1541        0.310598,
1542        -0.276022,
1543        0.160202,
1544        -0.39869,
1545        -0.664147,
1546        -0.194,
1547        0.0443349,
1548        0.298699,
1549        0.212516,
1550        0.0706104,
1551        -0.771505,
1552        0.118546,
1553        -0.467788,
1554        -0.23413,
1555        0.509708,
1556        0.274984,
1557        -0.347402,
1558        -0.127842,
1559        0.21163,
1560        -0.0675676,
1561        -0.0878888,
1562        -0.358893,
1563        0.137875,
1564        0.177817,
1565        -0.261019,
1566        -0.628945,
1567        0.518781,
1568        -0.231445,
1569        0.103466,
1570        -0.276795,
1571        -0.24368,
1572        0.265653,
1573        -0.115171,
1574        0.210652,
1575        0.217841,
1576        -0.173374,
1577        0.0684118,
1578        0.212281,
1579        0.35257,
1580        -0.0986456,
1581        -0.0154741,
1582        -0.299062,
1583        0.216102,
1584        0.0126834,
1585        -0.415971,
1586        -0.121538,
1587        -0.0510272,
1588        -0.192594,
1589        -0.231441,
1590        -0.504684,
1591        0.0135955,
1592        -0.142346,
1593        0.258746,
1594        -0.0460905,
1595        0.235684,
1596        -0.371481,
1597        0.0983083,
1598        0.202997,
1599        0.207863,
1600        -0.0960896,
1601        -0.278115,
1602        0.485785,
1603        -0.0703955,
1604        -0.36832,
1605        -0.0743555,
1606        0.0874494,
1607        -0.223167,
1608        0.112969,
1609        0.258647,
1610        -0.0219521,
1611        0.293628,
1612        -0.0430369,
1613        -0.148116,
1614        0.465595,
1615        0.180931,
1616        0.0256682,
1617        -0.544123,
1618        -0.419134,
1619        -0.504943,
1620        -0.0196961,
1621        -0.0160084,
1622        0.116454,
1623        0.301819,
1624        -0.463058,
1625        -0.409316,
1626        -0.280946,
1627        0.129024,
1628        0.0789985,
1629        -0.317849,
1630        -0.217392,
1631        0.10596,
1632        0.0127895,
1633        0.295834,
1634        -0.250482,
1635        -0.0489015,
1636        -0.303991,
1637        0.335735,
1638        0.298507,
1639        0.137767,
1640        0.283669,
1641        0.404645,
1642        0.42822,
1643        0.501337,
1644        -0.218792,
1645        -0.242324,
1646        0.380491,
1647        -0.129582,
1648        -0.530074,
1649        -0.13798,
1650        0.303014,
1651        -0.14071,
1652        0.226784,
1653        0.0899447,
1654        -0.403096,
1655        0.134022,
1656        0.00218902,
1657        -0.164883,
1658        -0.428667,
1659        0.47806,
1660        -0.273605,
1661        -0.610717,
1662        0.248782,
1663        -0.134394,
1664        0.353971,
1665        -0.0321702,
1666        -0.0053093,
1667        0.174744,
1668        0.034076,
1669        0.0328364,
1670        0.323228,
1671        -0.271353,
1672        0.165744,
1673        0.0854048,
1674        -0.481041,
1675        0.331567,
1676        0.229811,
1677        0.460195,
1678        -0.130766,
1679        -0.32927,
1680        0.919488,
1681        -0.304845,
1682        -0.237677,
1683        0.166027,
1684        -0.129779,
1685        0.123418,
1686        -0.397317,
1687        -0.509057,
1688        -0.164184,
1689        0.201893,
1690        -0.00989215,
1691        0.358239,
1692        -0.0787409,
1693        0.0564069,
1694        -0.272082,
1695        -0.437823,
1696        -0.363646,
1697        -0.238196,
1698        0.588519,
1699        -0.385448,
1700        -0.264242,
1701        0.410487,
1702        -0.199585,
1703        0.176318,
1704        -0.682871,
1705        -0.291284,
1706        0.15313,
1707        -0.509376,
1708        0.0532867,
1709        0.00989985,
1710        -0.169128,
1711        -0.0829413,
1712        0.210541,
1713        -0.0165126,
1714        0.293903,
1715        0.232969,
1716        0.46007,
1717        -0.128089,
1718        0.0495355,
1719        -0.317411,
1720        0.184258,
1721        -0.214564,
1722        0.183357,
1723        -0.0849151,
1724        -0.0527055,
1725        0.0672831,
1726        -0.223927,
1727        -0.202076,
1728        -0.188411,
1729        -0.16416,
1730        0.184064,
1731        -0.118991,
1732        -0.174886,
1733        0.242811,
1734        0.203539,
1735        -0.0508292,
1736        0.734019,
1737        -0.317723,
1738        0.477404,
1739        0.0731702,
1740        0.296654,
1741        -0.0947508,
1742        0.345185,
1743        -0.300144,
1744        0.234292,
1745        0.0730285,
1746        0.265356,
1747        -0.149729,
1748        0.121972,
1749        -0.239329,
1750        -0.111945,
1751        -0.0787418,
1752        -0.185964,
1753        -0.151469,
1754        -0.370867,
1755        -0.0418644,
1756        0.430025,
1757        0.846542,
1758        -0.254989,
1759        -0.231013,
1760        0.0964694,
1761        0.146091,
1762        -0.39068,
1763        -0.693897,
1764        0.0743391,
1765        0.396942,
1766        0.422646,
1767        -0.00549475,
1768        0.413812,
1769        -0.0535897,
1770        -0.421364,
1771        -0.0428651,
1772        0.307348,
1773        0.0362761,
1774        -0.172384,
1775        0.0135348,
1776        0.0981934,
1777        0.0389974,
1778        0.0883734,
1779        0.110775,
1780        0.317814,
1781        0.0373875,
1782        -0.0478011,
1783        0.105298,
1784        0.175046,
1785        -0.524262,
1786        0.341699,
1787        0.506601,
1788        0.00578107,
1789        0.182254,
1790        -0.223231,
1791        0.200005,
1792        -0.256655,
1793        0.0524954,
1794        0.17099,
1795        -0.237381,
1796        -0.333608,
1797        -0.105021,
1798        -0.170642,
1799        -0.0984785,
1800        -0.33122,
1801        0.323198,
1802        -0.0827793,
1803        0.130607,
1804        -0.489473,
1805        -0.177912,
1806        0.476636,
1807        -0.264911,
1808        0.0834568,
1809        0.624801,
1810        -0.00254831,
1811        0.44937,
1812        -0.270769,
1813        -0.145531,
1814        -0.0201544,
1815        -0.16437,
1816        0.100481,
1817        -0.534917,
1818        -0.320538,
1819        0.389989,
1820        0.0333177,
1821        0.143457,
1822        0.501674,
1823        -0.52644,
1824        -0.898354,
1825        0.527995,
1826        -0.414444,
1827        0.0201641,
1828        -0.492257,
1829        -0.771534,
1830        -0.0286519,
1831        0.133742,
1832        -0.253637,
1833        -0.768402,
1834        0.085024,
1835        0.562326,
1836        -0.526449,
1837        0.221449,
1838        0.0751945,
1839        -0.563033,
1840        0.295024,
1841        0.0474323,
1842        -0.0501224,
1843        0.375451,
1844        0.00226611,
1845        0.376463,
1846        0.778869,
1847        0.0582743,
1848        0.12641,
1849        -0.176428,
1850        0.456384,
1851        -0.371866,
1852        -0.156221,
1853        -0.0189524,
1854        -0.0548232,
1855        -0.214016,
1856        0.263157,
1857        -0.0931925,
1858        -0.645682,
1859        0.348074,
1860        0.519181,
1861        -0.167478,
1862        0.407343,
1863        -0.160051,
1864        -0.298173,
1865        -0.276582,
1866        0.491172,
1867        0.410276,
1868        0.208572,
1869        -0.303514,
1870        0.239872,
1871        -0.325825,
1872        0.187611,
1873        -0.16367,
1874        -0.170864,
1875        -0.156856,
1876        -0.621923,
1877        -0.526686,
1878        -0.0450273,
1879        -0.322389,
1880        0.246644,
1881        0.352417,
1882        -0.467928,
1883        -0.452286,
1884        -0.0509965,
1885        0.583273,
1886        -0.236968,
1887        0.148131,
1888        -0.497616,
1889        -0.275279,
1890        0.141518,
1891        -0.0770699,
1892        0.145007,
1893        0.0319243,
1894        -0.00993607,
1895        -0.016638,
1896        0.405023,
1897        -0.0263409,
1898        0.251191,
1899        -0.0873389,
1900        0.0196543,
1901        -0.330143,
1902        -0.638085,
1903        -0.136588,
1904        0.019404,
1905        -0.153683,
1906        0.152007,
1907        -0.251054,
1908        -0.4297,
1909        -0.454367,
1910        0.183937,
1911        -0.2615,
1912        -0.0117493,
1913        -0.438665,
1914        -0.276094,
1915        -0.0236997,
1916        0.305428,
1917        -0.339538,
1918        -0.455488,
1919        -0.0246618,
1920        0.0245187,
1921        -0.15223,
1922        0.343269,
1923        -0.538411,
1924        -0.270215,
1925        -0.0232314,
1926        -0.286599,
1927        0.0768719,
1928        -0.0225739,
1929        -0.0989169,
1930        0.699144,
1931        -0.321706,
1932        -0.103017,
1933        0.122571,
1934        -0.249993,
1935        0.0163995,
1936        0.232155,
1937        0.0407288,
1938        0.198615,
1939        -0.13675,
1940        -0.10254,
1941        0.552905,
1942        0.11824,
1943        0.266287,
1944        -0.219535,
1945        -0.0890321,
1946        0.32599,
1947        0.00448736,
1948        -0.395851,
1949        -0.292809,
1950        -0.0430488,
1951        0.448302,
1952        0.1378,
1953        -0.251007,
1954        -0.0270933,
1955        -0.00214663,
1956        -0.523369,
1957        -0.0309944,
1958        -0.125962,
1959        -0.0659885,
1960        -0.289194,
1961        0.0919669,
1962        -0.144388,
1963        -0.252316,
1964        0.535763,
1965        -0.132497,
1966        -0.429882,
1967        0.346153,
1968        0.403493,
1969        -0.309738,
1970        -0.1503,
1971        -0.0742551,
1972        -0.141614,
1973        -0.436618,
1974        -0.345322,
1975        0.620119,
1976        -0.204794,
1977        -0.181742,
1978        0.209699,
1979        -0.166686,
1980        -0.223907,
1981        -0.337174,
1982        0.136383,
1983        0.0377023,
1984        0.312557,
1985        -0.151656,
1986        0.366969,
1987        0.0219035,
1988        -0.182791,
1989        0.113654,
1990        0.0807236,
1991        0.0594,
1992        -0.249578,
1993        -0.0533519,
1994        -0.081953,
1995        0.196487,
1996        0.362187,
1997        0.226492,
1998        0.201226,
1999        0.308508,
2000        0.347177,
2001        -0.139138,
2002        -0.371969,
2003        0.0258503,
2004        0.430667,
2005        -0.398186,
2006        -0.309352,
2007        -0.00336896,
2008        -0.187937,
2009        0.276279,
2010        0.00202034,
2011        0.231139,
2012        0.0684979,
2013        0.100552,
2014        -0.374401,
2015        -0.214862,
2016        -0.610311,
2017        -0.252147,
2018        -0.185271,
2019        0.0818731,
2020        -0.366164,
2021        0.512052,
2022        0.345181,
2023        -0.325809,
2024        -0.572269,
2025        -0.500868,
2026        0.0718371,
2027        -0.391286,
2028        -0.107487,
2029        0.148359,
2030        0.0972818,
2031        0.200506,
2032        -0.406185,
2033        0.177848,
2034        0.187755,
2035        0.265422,
2036        0.228593,
2037        -0.401342,
2038        -0.751146,
2039        0.0909925,
2040        0.460462,
2041        -0.249553,
2042        -0.252755,
2043        0.0340475,
2044        0.306479,
2045        0.202305,
2046        -0.372425,
2047        -0.147982,
2048        -0.186426,
2049        0.154744,
2050        0.209422,
2051        -0.281605,
2052        -0.259755,
2053        0.0696664,
2054        -0.0819251,
2055        -0.247159,
2056        -0.151843,
2057        0.164756,
2058        -0.121739,
2059        -0.052189,
2060        -0.248309,
2061        0.162564,
2062        -0.171285,
2063        0.0632468,
2064        0.0353644,
2065        -0.132106,
2066        -0.266058,
2067        -0.0386287,
2068        -0.10764,
2069        0.288536,
2070        -0.269559,
2071        0.212119,
2072        0.347515,
2073        0.182925,
2074        0.209263,
2075        -0.0891225,
2076        0.123076,
2077        0.125977,
2078        -0.493741,
2079        -0.107514,
2080        -0.371816,
2081        -0.567552,
2082        -0.0149517,
2083        0.12388,
2084        -0.206517,
2085        0.148488,
2086        0.165108,
2087        0.0233112,
2088        -0.15338,
2089        0.197413,
2090        0.52503,
2091        -0.61807,
2092        0.215751,
2093        -0.349172,
2094        0.251406,
2095        0.296309,
2096        -0.123011,
2097        -0.130789,
2098        -0.00651585,
2099        0.035497,
2100        -0.0609056,
2101        0.341434,
2102        -0.470598,
2103        -0.150903,
2104        0.195183,
2105        -0.0862897,
2106        -0.0746215,
2107        -0.0906219,
2108        0.285372,
2109        0.208254,
2110        -0.397821,
2111        -0.341474,
2112        -0.124454,
2113        -0.0414253,
2114        0.271967,
2115        -0.0351484,
2116        0.396491,
2117        0.559566,
2118        -0.260349,
2119        0.0515127,
2120        0.309838,
2121        0.0710535,
2122        0.104538,
2123        -0.135632,
2124        -0.163842,
2125        0.063358,
2126        0.300885,
2127        -0.31289,
2128        -0.0402696,
2129        0.542997,
2130        -0.0556195,
2131        0.0532273,
2132        0.243242,
2133        0.175341,
2134        -0.0434567,
2135        -0.0617974,
2136        0.461337,
2137        -0.279505,
2138        -0.402233,
2139        0.338011,
2140        0.383038,
2141        -0.284724,
2142        0.269527,
2143        0.115576,
2144        -0.0259366,
2145        0.041392,
2146        0.563885,
2147        0.23412,
2148        0.188308,
2149        0.168665,
2150        -0.166081,
2151        -0.0822718,
2152        -0.0302363,
2153        -0.250662,
2154        0.172037,
2155        0.358585,
2156        -0.0150482,
2157        -0.232162,
2158        -0.22336,
2159        0.492373,
2160        -0.162311,
2161        -0.352932,
2162        -0.432669,
2163        -0.128367,
2164        0.378123,
2165        0.388758,
2166        0.191049,
2167        -0.208064,
2168        0.261632,
2169        -0.0438233,
2170        0.197892,
2171        0.668999,
2172        -0.0518101,
2173        -0.392415,
2174        0.589557,
2175        -0.0304808,
2176        -0.0361237,
2177        -0.473126,
2178        0.243661,
2179        0.336332,
2180        0.0973627,
2181        0.349706,
2182        -0.096878,
2183        0.427568,
2184        0.255798,
2185        0.553208,
2186        -0.477473,
2187        -0.351159,
2188        0.291254,
2189        0.137889,
2190        -0.24987,
2191        0.0544092,
2192        0.305517,
2193        0.369363,
2194        0.108008,
2195        0.0571747,
2196        0.170541,
2197        0.0154032,
2198        0.603159,
2199        -0.0985079,
2200        0.0345475,
2201        -0.0276445,
2202        0.495077,
2203        -0.130628,
2204        0.29438,
2205        0.0118292,
2206        -0.107588,
2207        -0.598261,
2208        0.00946725,
2209        0.530881,
2210        -0.181808,
2211        -0.851401,
2212        0.113912,
2213        -0.335427,
2214        -0.511092,
2215        -0.136116,
2216        -0.382416,
2217        0.332571,
2218        0.0797383,
2219        0.200826,
2220        -0.175843,
2221        -0.0182147,
2222        -0.156549,
2223        0.0881938,
2224        -0.665213,
2225        -0.00780874,
2226        0.310998,
2227        0.511727,
2228        0.250097,
2229        0.0436686,
2230        -0.511367,
2231        0.0453479,
2232    ];
2233    const VIS_MERGE_PROJ1: [f32; 192] = [
2234        -0.589584,
2235        -0.270176,
2236        -0.0642117,
2237        -0.0839581,
2238        0.263636,
2239        -0.0866331,
2240        -0.38092,
2241        -0.14856,
2242        -0.230682,
2243        0.322097,
2244        -0.444679,
2245        -0.150674,
2246        0.270976,
2247        -0.261266,
2248        0.110946,
2249        0.00495127,
2250        -0.106307,
2251        -0.531284,
2252        0.00520624,
2253        -0.15978,
2254        0.0808836,
2255        -0.440902,
2256        0.153633,
2257        0.194342,
2258        0.417523,
2259        -0.148402,
2260        -0.212159,
2261        0.268179,
2262        -0.317149,
2263        -0.260497,
2264        -0.09544,
2265        -0.343195,
2266        0.429478,
2267        0.285979,
2268        -0.654128,
2269        -0.645155,
2270        -0.196368,
2271        0.148705,
2272        0.203804,
2273        0.422358,
2274        0.0848344,
2275        -0.321419,
2276        0.125365,
2277        -0.34804,
2278        -0.251555,
2279        -0.482738,
2280        -0.271232,
2281        -0.18782,
2282        -0.172891,
2283        0.530903,
2284        0.0288639,
2285        0.104308,
2286        0.0468265,
2287        0.475097,
2288        -0.285449,
2289        0.225047,
2290        -0.256494,
2291        -0.0416107,
2292        0.148517,
2293        -0.405077,
2294        -0.115967,
2295        0.00170611,
2296        0.151352,
2297        -0.0672187,
2298        0.846001,
2299        0.564592,
2300        -0.660677,
2301        0.0353704,
2302        0.166182,
2303        0.102616,
2304        0.417172,
2305        -0.858214,
2306        -0.322372,
2307        0.462712,
2308        0.364483,
2309        0.685643,
2310        0.496457,
2311        -0.38512,
2312        0.310468,
2313        -0.0848315,
2314        0.0626254,
2315        0.0168127,
2316        -0.140008,
2317        -0.151035,
2318        0.232496,
2319        -0.163294,
2320        0.0647018,
2321        0.145547,
2322        -0.159983,
2323        -0.0856119,
2324        0.345313,
2325        -0.509969,
2326        -0.201798,
2327        -0.401557,
2328        0.0651107,
2329        0.811157,
2330        -0.368641,
2331        0.368018,
2332        0.426114,
2333        -0.64137,
2334        -0.402644,
2335        -0.00634645,
2336        -0.296265,
2337        0.0834438,
2338        0.409867,
2339        0.211353,
2340        -0.0590335,
2341        0.170205,
2342        0.105546,
2343        -0.380194,
2344        -0.180618,
2345        0.140349,
2346        -0.0205445,
2347        0.0967417,
2348        0.267025,
2349        -0.725734,
2350        0.444512,
2351        -0.572309,
2352        0.00349554,
2353        0.0459292,
2354        -0.0438933,
2355        0.0563175,
2356        -0.168883,
2357        -0.219364,
2358        -0.0316683,
2359        0.051398,
2360        -0.207059,
2361        0.0949521,
2362        -0.260774,
2363        -0.177944,
2364        -0.0742896,
2365        0.131637,
2366        -0.0958914,
2367        0.114644,
2368        0.480587,
2369        -0.146184,
2370        0.387162,
2371        -0.400214,
2372        0.190505,
2373        -0.0616221,
2374        0.628918,
2375        -0.156545,
2376        -0.253828,
2377        0.311851,
2378        0.119144,
2379        0.0184293,
2380        -0.729126,
2381        0.113713,
2382        -0.0374038,
2383        0.0200473,
2384        -0.239673,
2385        0.385356,
2386        -0.0108929,
2387        -0.0806783,
2388        -0.47495,
2389        -0.0672919,
2390        -0.336395,
2391        0.401548,
2392        0.0655212,
2393        -0.190752,
2394        -0.0640861,
2395        0.425706,
2396        0.0323832,
2397        -0.0136082,
2398        0.54734,
2399        0.10145,
2400        0.261653,
2401        -0.115325,
2402        -0.124108,
2403        0.243261,
2404        -0.0241873,
2405        0.0977518,
2406        -0.0293961,
2407        -0.00992239,
2408        -0.0499667,
2409        -0.404479,
2410        0.381245,
2411        -0.189778,
2412        -0.368962,
2413        0.0141507,
2414        -0.206907,
2415        0.368309,
2416        0.18214,
2417        -0.186088,
2418        -0.275483,
2419        0.183006,
2420        -0.218805,
2421        -0.421012,
2422        -0.508387,
2423        0.201931,
2424        -0.478372,
2425        -0.364119,
2426    ];
2427    const VIS_POST_NORM_W: [f32; 6] = [1.06551, 1.07693, 0.879476, 0.935329, 1.09384, 1.01741];
2428
2429    const VIS_GOLDEN_ENCODER_OUT: [f32; 32] = [
2430        -0.791224, 0.0928006, -1.20306, 0.668509, -1.51065, 0.051747, -1.10594, 1.69022, 0.539921,
2431        1.01363, -1.06878, -0.240584, -1.83794, 0.908773, -0.935426, 1.08535, -0.259518, 1.01078,
2432        -0.0224143, -1.81579, 1.19159, 1.00525, -0.470024, -0.0242557, -1.95959, 0.713632,
2433        -0.145732, -1.19647, -0.051718, 0.905252, -0.80985, -0.243886,
2434    ];
2435    const VIS_GOLDEN_OUTPUT: [f32; 6] = [
2436        -2.50858, 0.0255713, -0.376901, -0.370726, -0.301048, -0.203254,
2437    ];
2438
2439    fn wm(data: &[f32], rows: usize, cols: usize) -> WeightMatrix {
2440        assert_eq!(data.len(), rows * cols);
2441        WeightMatrix::F32(Tensor::new(data.to_vec(), vec![rows, cols]))
2442    }
2443
2444    fn cfg() -> VisionConfig {
2445        VisionConfig {
2446            in_dim: IN_DIM,
2447            patch_size: PATCH_SIZE,
2448            grid_h: GRID_H,
2449            grid_w: GRID_W,
2450            hidden_dim: HIDDEN_DIM,
2451            num_heads: NUM_HEADS,
2452            qkv_hidden: QKV_HIDDEN,
2453            mlp_dim: MLP_DIM,
2454            rms_norm_eps: NORM_EPS,
2455            theta_base: THETA_BASE,
2456            merge_kh: MERGE_KH,
2457            merge_kw: MERGE_KW,
2458            projector_ln_eps: PROJECTOR_LN_EPS,
2459        }
2460    }
2461
2462    fn make_encoder_weights() -> VisionEncoderWeights {
2463        let patch_dim = IN_DIM * PATCH_SIZE * PATCH_SIZE;
2464        VisionEncoderWeights {
2465            patch_embed: wm(&VIS_PATCH_EMBED_W, HIDDEN_DIM, patch_dim),
2466            pos_emb: VIS_POS_EMB_W.to_vec(),
2467            layers: vec![
2468                VisionEncoderLayerWeights {
2469                    norm0_weight: VIS_L0_NORM0_W.to_vec(),
2470                    wqkv: wm(&VIS_L0_WQKV, 3 * QKV_HIDDEN, HIDDEN_DIM),
2471                    wo: wm(&VIS_L0_WO, HIDDEN_DIM, QKV_HIDDEN),
2472                    norm1_weight: VIS_L0_NORM1_W.to_vec(),
2473                    fc0: wm(&VIS_L0_FC0, MLP_DIM, HIDDEN_DIM),
2474                    fc1: wm(&VIS_L0_FC1, HIDDEN_DIM, MLP_DIM),
2475                },
2476                VisionEncoderLayerWeights {
2477                    norm0_weight: VIS_L1_NORM0_W.to_vec(),
2478                    wqkv: wm(&VIS_L1_WQKV, 3 * QKV_HIDDEN, HIDDEN_DIM),
2479                    wo: wm(&VIS_L1_WO, HIDDEN_DIM, QKV_HIDDEN),
2480                    norm1_weight: VIS_L1_NORM1_W.to_vec(),
2481                    fc0: wm(&VIS_L1_FC0, MLP_DIM, HIDDEN_DIM),
2482                    fc1: wm(&VIS_L1_FC1, HIDDEN_DIM, MLP_DIM),
2483                },
2484            ],
2485            final_norm_weight: VIS_FINAL_NORM_W.to_vec(),
2486        }
2487    }
2488
2489    fn make_merger_weights() -> VisionMergerWeights {
2490        let merge_hidden = MERGE_KH * MERGE_KW * HIDDEN_DIM;
2491        VisionMergerWeights {
2492            proj0: wm(&VIS_MERGE_PROJ0, merge_hidden, merge_hidden),
2493            proj1: wm(&VIS_MERGE_PROJ1, TEXT_HIDDEN, merge_hidden),
2494            post_norm_weight: VIS_POST_NORM_W.to_vec(),
2495        }
2496    }
2497
2498    #[test]
2499    fn encoder_matches_independent_python_reference() {
2500        let cfg = cfg();
2501        let weights = make_encoder_weights();
2502        let patches = vec![
2503            VIS_PATCH_0.to_vec(),
2504            VIS_PATCH_1.to_vec(),
2505            VIS_PATCH_2.to_vec(),
2506            VIS_PATCH_3.to_vec(),
2507        ];
2508
2509        let out = encoder_forward(&weights, &cfg, &patches);
2510        assert_eq!(out.len(), VIS_GOLDEN_ENCODER_OUT.len());
2511        for (i, (a, b)) in out.iter().zip(VIS_GOLDEN_ENCODER_OUT.iter()).enumerate() {
2512            assert!((a - b).abs() < 1e-3, "element {i}: rust={a} python={b}");
2513        }
2514    }
2515
2516    #[test]
2517    fn full_pipeline_matches_independent_python_reference() {
2518        let cfg = cfg();
2519        let encoder_weights = make_encoder_weights();
2520        let merger_weights = make_merger_weights();
2521        let patches = vec![
2522            VIS_PATCH_0.to_vec(),
2523            VIS_PATCH_1.to_vec(),
2524            VIS_PATCH_2.to_vec(),
2525            VIS_PATCH_3.to_vec(),
2526        ];
2527
2528        let encoder_out = encoder_forward(&encoder_weights, &cfg, &patches);
2529        let merged = patch_merge(&encoder_out, &cfg);
2530        assert_eq!(merged.len(), MERGE_KH * MERGE_KW * HIDDEN_DIM); // one merged patch
2531
2532        let projected = project_merged_patches(&merger_weights, &cfg, &merged, 1);
2533        assert_eq!(projected.len(), VIS_GOLDEN_OUTPUT.len());
2534        for (i, (a, b)) in projected.iter().zip(VIS_GOLDEN_OUTPUT.iter()).enumerate() {
2535            assert!((a - b).abs() < 1e-3, "element {i}: rust={a} python={b}");
2536        }
2537    }
2538
2539    #[test]
2540    fn patch_merge_groups_spatially_adjacent_patches_row_major() {
2541        // 2x2 grid, merge kernel 2x2 -> one merged patch containing all
2542        // 4 original patches in (kh, kw) order: (0,0),(0,1),(1,0),(1,1).
2543        let hidden_dim = 1;
2544        let cfg = VisionConfig {
2545            in_dim: 1,
2546            patch_size: 1,
2547            grid_h: 2,
2548            grid_w: 2,
2549            hidden_dim,
2550            num_heads: 1,
2551            qkv_hidden: 1,
2552            mlp_dim: 1,
2553            rms_norm_eps: 1e-5,
2554            theta_base: 10000.0,
2555            merge_kh: 2,
2556            merge_kw: 2,
2557            projector_ln_eps: 1e-5,
2558        };
2559        // patch order row-major: (0,0)=10, (0,1)=20, (1,0)=30, (1,1)=40
2560        let encoder_out = vec![10.0, 20.0, 30.0, 40.0];
2561        let merged = patch_merge(&encoder_out, &cfg);
2562        assert_eq!(merged, vec![10.0, 20.0, 30.0, 40.0]);
2563    }
2564
2565    #[test]
2566    fn erf_matches_known_values() {
2567        assert!((erf(0.0)).abs() < 1e-6);
2568        assert!((erf(1.0) - 0.8427008).abs() < 1e-4);
2569        assert!((erf(-1.0) + 0.8427008).abs() < 1e-4);
2570    }
2571}
ferrox_models/vision.rs

ferrox_models/
vision.rs
