zsfm_sundial/infer/
mod.rs

1//! Sundial Burn-only inference engine (Burn backend, Flex CPU).
2//!
3//! Patch embedding, transformer blocks with Llama-style RoPE, and the
4//! flow-matching ODE sampler (Heun), with weights loaded straight from GGUF
5//! via [`zsfm_burn::GgufWeights`]. Architecture widths mirror the fixed
6//! constants; layer counts come from GGUF metadata.
7
8use std::path::Path;
9
10use anyhow::{Context, Result};
11use burn::tensor::{activation, Device, Tensor, TensorData};
12use zsfm_burn::linear::burn_linear_nd;
13use zsfm_burn::norm::burn_layer_norm_nd;
14use zsfm_burn::{open_gguf_weights, GgufWeights};
15
16const N_HEADS: usize = 12;
17const HEAD_DIM: usize = 64;
18const EMBED_IN: usize = 32;
19const PATCH_SIZE: usize = 16;
20const TIME_DIM: usize = 256;
21const MAX_SEQ: usize = 4096;
22const ROPE_THETA: f64 = 10000.0;
23
24fn device() -> Device {
25    Device::flex()
26}
27
28fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
29    weights.load_2d(name)
30}
31
32fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
33    weights.load_1d(name)
34}
35
36fn linear<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
37    burn_linear_nd(x, w, b)
38}
39
40fn layer_norm<const D: usize>(x: Tensor<D>, w: Tensor<1>, b: Tensor<1>) -> Tensor<D> {
41    burn_layer_norm_nd(x, w, b, 1e-5)
42}
43
44fn silu<const D: usize>(x: Tensor<D>) -> Tensor<D> {
45    activation::silu(x)
46}
47
48fn softmax_last<const D: usize>(x: Tensor<D>) -> Tensor<D> {
49    activation::softmax(x, D - 1)
50}
51
52struct BurnEmbedW {
53    hidden_w: Tensor<2>,
54    hidden_b: Tensor<1>,
55    output_w: Tensor<2>,
56    output_b: Tensor<1>,
57    skip_w: Tensor<2>,
58    skip_b: Tensor<1>,
59}
60
61struct BurnAttnW {
62    qkv_w: Tensor<2>,
63    qkv_b: Tensor<1>,
64    o_w: Tensor<2>,
65}
66
67struct BurnNormW {
68    w: Tensor<1>,
69    b: Tensor<1>,
70}
71
72struct BurnBlockW {
73    attn: BurnAttnW,
74    attn_norm: BurnNormW,
75    ffn_norm: BurnNormW,
76    gate_w: Tensor<2>,
77    up_w: Tensor<2>,
78    down_w: Tensor<2>,
79}
80
81struct BurnFlowResW {
82    ln_w: Tensor<1>,
83    ln_b: Tensor<1>,
84    mlp1_w: Tensor<2>,
85    mlp1_b: Tensor<1>,
86    mlp2_w: Tensor<2>,
87    mlp2_b: Tensor<1>,
88    adaln_w: Tensor<2>,
89    adaln_b: Tensor<1>,
90}
91
92struct BurnFlowW {
93    t1_w: Tensor<2>,
94    t1_b: Tensor<1>,
95    t2_w: Tensor<2>,
96    t2_b: Tensor<1>,
97    cond_w: Tensor<2>,
98    cond_b: Tensor<1>,
99    in_w: Tensor<2>,
100    in_b: Tensor<1>,
101    res: Vec<BurnFlowResW>,
102    out_w: Tensor<2>,
103    out_b: Tensor<1>,
104    out_adaln_w: Tensor<2>,
105    out_adaln_b: Tensor<1>,
106}
107
108pub struct BurnSundialModel {
109    embed: BurnEmbedW,
110    blocks: Vec<BurnBlockW>,
111    norm_w: Tensor<1>,
112    norm_b: Tensor<1>,
113    flow: BurnFlowW,
114    rope_cos: Tensor<2>,
115    rope_sin: Tensor<2>,
116    n_steps: usize,
117    output_len: usize,
118    t_emb_table: Vec<Tensor<1>>,
119}
120
121fn get_u32(weights: &GgufWeights, key: &str) -> Option<u32> {
122    weights.meta_u32(key)
123}
124
125fn sinusoidal_embed(t: f32, dim: usize) -> Vec<f32> {
126    let half = dim / 2;
127    let freqs: Vec<f32> = (0..half)
128        .map(|i| (-(10000.0f32.ln() * i as f32 / half as f32)).exp())
129        .collect();
130    let args: Vec<f32> = freqs.iter().map(|&f| t * f).collect();
131    let mut emb = Vec::with_capacity(dim);
132    for &a in &args {
133        emb.push(a.cos());
134    }
135    for &a in &args {
136        emb.push(a.sin());
137    }
138    emb
139}
140
141impl BurnSundialModel {
142    pub fn load(path: &Path, steps_override: Option<usize>) -> Result<Self> {
143        let weights =
144            open_gguf_weights(path).with_context(|| format!("open {}", path.display()))?;
145
146        let n_layers = get_u32(&weights, "sundial1.block_count").unwrap_or(12) as usize;
147        let gguf_n_steps =
148            get_u32(&weights, "sundial1.flow.num_sampling_steps").unwrap_or(50) as usize;
149        let n_steps = steps_override.unwrap_or(gguf_n_steps);
150        let flow_depth = get_u32(&weights, "sundial1.flow.depth").unwrap_or(3) as usize;
151        let output_len = get_u32(&weights, "sundial1.output_token_len").unwrap_or(720) as usize;
152
153        let embed = BurnEmbedW {
154            hidden_w: load_t(&weights, "embed.hidden.weight")?,
155            hidden_b: load_v(&weights, "embed.hidden.bias")?,
156            output_w: load_t(&weights, "embed.output.weight")?,
157            output_b: load_v(&weights, "embed.output.bias")?,
158            skip_w: load_t(&weights, "embed.skip.weight")?,
159            skip_b: load_v(&weights, "embed.skip.bias")?,
160        };
161
162        let mut blocks = Vec::with_capacity(n_layers);
163        for n in 0..n_layers {
164            let qw = load_t(&weights, &format!("blk.{n}.attn_q.weight"))?;
165            let kw = load_t(&weights, &format!("blk.{n}.attn_k.weight"))?;
166            let vw = load_t(&weights, &format!("blk.{n}.attn_v.weight"))?;
167            let qb = load_v(&weights, &format!("blk.{n}.attn_q.bias"))?;
168            let kb = load_v(&weights, &format!("blk.{n}.attn_k.bias"))?;
169            let vb = load_v(&weights, &format!("blk.{n}.attn_v.bias"))?;
170            let qkv_w = Tensor::cat(vec![qw, kw, vw], 0);
171            let qkv_b = Tensor::cat(vec![qb, kb, vb], 0);
172            blocks.push(BurnBlockW {
173                attn: BurnAttnW {
174                    qkv_w,
175                    qkv_b,
176                    o_w: load_t(&weights, &format!("blk.{n}.attn_out.weight"))?,
177                },
178                attn_norm: BurnNormW {
179                    w: load_v(&weights, &format!("blk.{n}.attn_norm.weight"))?,
180                    b: load_v(&weights, &format!("blk.{n}.attn_norm.bias"))?,
181                },
182                ffn_norm: BurnNormW {
183                    w: load_v(&weights, &format!("blk.{n}.ffn_norm.weight"))?,
184                    b: load_v(&weights, &format!("blk.{n}.ffn_norm.bias"))?,
185                },
186                gate_w: load_t(&weights, &format!("blk.{n}.ffn_gate.weight"))?,
187                up_w: load_t(&weights, &format!("blk.{n}.ffn_up.weight"))?,
188                down_w: load_t(&weights, &format!("blk.{n}.ffn_down.weight"))?,
189            });
190        }
191
192        let norm_w = load_v(&weights, "norm.weight")?;
193        let norm_b = load_v(&weights, "norm.bias")?;
194
195        let mut res_blocks = Vec::with_capacity(flow_depth);
196        for k in 0..flow_depth {
197            let fp = |s: &str| format!("flow.res.{k}.{s}");
198            res_blocks.push(BurnFlowResW {
199                ln_w: load_v(&weights, &fp("ln.weight"))?,
200                ln_b: load_v(&weights, &fp("ln.bias"))?,
201                mlp1_w: load_t(&weights, &fp("mlp1.weight"))?,
202                mlp1_b: load_v(&weights, &fp("mlp1.bias"))?,
203                mlp2_w: load_t(&weights, &fp("mlp2.weight"))?,
204                mlp2_b: load_v(&weights, &fp("mlp2.bias"))?,
205                adaln_w: load_t(&weights, &fp("adaln.weight"))?,
206                adaln_b: load_v(&weights, &fp("adaln.bias"))?,
207            });
208        }
209        let flow = BurnFlowW {
210            t1_w: load_t(&weights, "flow.t_proj1.weight")?,
211            t1_b: load_v(&weights, "flow.t_proj1.bias")?,
212            t2_w: load_t(&weights, "flow.t_proj2.weight")?,
213            t2_b: load_v(&weights, "flow.t_proj2.bias")?,
214            cond_w: load_t(&weights, "flow.cond.weight")?,
215            cond_b: load_v(&weights, "flow.cond.bias")?,
216            in_w: load_t(&weights, "flow.in_proj.weight")?,
217            in_b: load_v(&weights, "flow.in_proj.bias")?,
218            res: res_blocks,
219            out_w: load_t(&weights, "flow.out_linear.weight")?,
220            out_b: load_v(&weights, "flow.out_linear.bias")?,
221            out_adaln_w: load_t(&weights, "flow.out_adaln.weight")?,
222            out_adaln_b: load_v(&weights, "flow.out_adaln.bias")?,
223        };
224
225        let (rope_cos, rope_sin) = build_rope_tables(HEAD_DIM, MAX_SEQ, ROPE_THETA);
226
227        // Precompute time embeddings for all ODE steps.
228        let t_emb_table: Vec<Tensor<1>> = (0..=n_steps)
229            .map(|i| {
230                let t_scaled = i as f32 / n_steps as f32 * 1000.0;
231                let t_raw = Tensor::<2>::from_data(
232                    TensorData::new(sinusoidal_embed(t_scaled, TIME_DIM), [1, TIME_DIM]),
233                    &device(),
234                );
235                let t_h = silu(linear(t_raw, flow.t1_w.clone(), Some(flow.t1_b.clone())));
236                linear(t_h, flow.t2_w.clone(), Some(flow.t2_b.clone()))
237                    .reshape([flow.t2_b.dims()[0]])
238            })
239            .collect();
240
241        Ok(Self {
242            embed,
243            blocks,
244            norm_w,
245            norm_b,
246            flow,
247            rope_cos,
248            rope_sin,
249            n_steps,
250            output_len,
251            t_emb_table,
252        })
253    }
254
255    pub fn forecast(&self, context: &[f32]) -> Result<Vec<f32>> {
256        let dev = device();
257        let n = context.len();
258        let mean = context.iter().sum::<f32>() / n as f32;
259        let var = context
260            .iter()
261            .map(|&x| (x - mean) * (x - mean))
262            .sum::<f32>()
263            / n as f32;
264        let std = (var + 1e-5).sqrt();
265        let normed: Vec<f32> = context.iter().map(|&x| (x - mean) / std).collect();
266
267        let n_patches = (n + PATCH_SIZE - 1) / PATCH_SIZE;
268        let mut embed_in = vec![0.0f32; n_patches * EMBED_IN];
269        for p in 0..n_patches {
270            let start = p * PATCH_SIZE;
271            let end = (start + PATCH_SIZE).min(n);
272            for i in start..end {
273                embed_in[p * EMBED_IN + (i - start)] = normed[i];
274                embed_in[p * EMBED_IN + PATCH_SIZE + (i - start)] = 1.0;
275            }
276        }
277        let x = Tensor::<3>::from_data(TensorData::new(embed_in, [1, n_patches, EMBED_IN]), &dev);
278        let mut h = self.embed_forward(x);
279        let mask = burn_causal_mask(n_patches);
280        for block in &self.blocks {
281            h = self.block_forward(h, block, mask.clone());
282        }
283        h = layer_norm(h, self.norm_w.clone(), self.norm_b.clone());
284        let cond: Tensor<2> = h.narrow(1, n_patches - 1, 1).squeeze_dim(0);
285        let output = self.flow_sample(cond);
286        let vals: Vec<f32> = output
287            .to_data()
288            .try_to_vec()
289            .map_err(|e| anyhow::anyhow!("{e:?}"))?;
290        Ok(vals.iter().map(|&v| v * std + mean).collect())
291    }
292
293    fn embed_forward(&self, x: Tensor<3>) -> Tensor<3> {
294        let h = silu(linear(
295            x.clone(),
296            self.embed.hidden_w.clone(),
297            Some(self.embed.hidden_b.clone()),
298        ));
299        let out = linear(
300            h,
301            self.embed.output_w.clone(),
302            Some(self.embed.output_b.clone()),
303        );
304        let skip = linear(
305            x,
306            self.embed.skip_w.clone(),
307            Some(self.embed.skip_b.clone()),
308        );
309        out + skip
310    }
311
312    fn block_forward(&self, x: Tensor<3>, blk: &BurnBlockW, mask: Tensor<4>) -> Tensor<3> {
313        let h = layer_norm(x.clone(), blk.attn_norm.w.clone(), blk.attn_norm.b.clone());
314        let h = self.attn_forward(h, &blk.attn, mask);
315        let x = x + h;
316        let h = layer_norm(x.clone(), blk.ffn_norm.w.clone(), blk.ffn_norm.b.clone());
317        let h = self.ffn_forward(h, blk);
318        x + h
319    }
320
321    fn attn_forward(&self, x: Tensor<3>, attn: &BurnAttnW, mask: Tensor<4>) -> Tensor<3> {
322        let [b, seq, _] = x.dims();
323        let d = N_HEADS * HEAD_DIM;
324        let qkv = linear(x, attn.qkv_w.clone(), Some(attn.qkv_b.clone()));
325        let q: Tensor<3> = qkv.clone().narrow(2, 0, d);
326        let k: Tensor<3> = qkv.clone().narrow(2, d, d);
327        let v: Tensor<3> = qkv.narrow(2, 2 * d, d);
328        let pack = |t: Tensor<3>| t.reshape([b, seq, N_HEADS, HEAD_DIM]).permute([0, 2, 1, 3]);
329        let (q, k, v) = (pack(q), pack(k), pack(v));
330        let q = self.rope_apply(q, 0);
331        let k = self.rope_apply(k, 0);
332        let scale = 1.0f32 / (HEAD_DIM as f32).sqrt();
333        let scores = q.matmul(k.transpose()).mul_scalar(scale);
334        let scores = scores + mask;
335        let aw = softmax_last(scores);
336        let out = aw.matmul(v);
337        let out: Tensor<3> = out.permute([0, 2, 1, 3]).reshape([b, seq, d]);
338        linear(out, attn.o_w.clone(), None)
339    }
340
341    fn rope_apply(&self, x: Tensor<4>, start_pos: usize) -> Tensor<4> {
342        let dims = x.dims();
343        let seq = dims[2];
344        let hd = dims[3];
345        let half = hd / 2;
346        let cos = self
347            .rope_cos
348            .clone()
349            .narrow(0, start_pos, seq)
350            .unsqueeze_dim::<3>(0)
351            .unsqueeze_dim::<4>(0);
352        let sin = self
353            .rope_sin
354            .clone()
355            .narrow(0, start_pos, seq)
356            .unsqueeze_dim::<3>(0)
357            .unsqueeze_dim::<4>(0);
358        let x1 = x.clone().narrow(3, 0, half);
359        let x2 = x.clone().narrow(3, half, half);
360        let rotated = Tensor::cat(vec![x2.mul_scalar(-1.0), x1], 3);
361        x * cos + rotated * sin
362    }
363
364    fn ffn_forward(&self, x: Tensor<3>, blk: &BurnBlockW) -> Tensor<3> {
365        let gate = silu(linear(x.clone(), blk.gate_w.clone(), None));
366        let up = linear(x, blk.up_w.clone(), None);
367        linear(gate * up, blk.down_w.clone(), None)
368    }
369
370    fn flow_net_eval(&self, x: Tensor<2>, c_act: Tensor<2>) -> Tensor<2> {
371        let mut h = linear(x, self.flow.in_w.clone(), Some(self.flow.in_b.clone()));
372        for res in &self.flow.res {
373            let adaln = linear(
374                c_act.clone(),
375                res.adaln_w.clone(),
376                Some(res.adaln_b.clone()),
377            );
378            let third = adaln.dims()[1] / 3;
379            let shift = adaln.clone().narrow(1, 0, third);
380            let scale = adaln.clone().narrow(1, third, third);
381            let gate = adaln.narrow(1, 2 * third, third);
382            let h_norm = layer_norm(h.clone(), res.ln_w.clone(), res.ln_b.clone());
383            let h_mod = h_norm * scale.add_scalar(1.0) + shift;
384            let h_mlp = silu(linear(h_mod, res.mlp1_w.clone(), Some(res.mlp1_b.clone())));
385            let h_mlp = linear(h_mlp, res.mlp2_w.clone(), Some(res.mlp2_b.clone()));
386            h = h + gate * h_mlp;
387        }
388        let adaln = linear(
389            c_act,
390            self.flow.out_adaln_w.clone(),
391            Some(self.flow.out_adaln_b.clone()),
392        );
393        let half = adaln.dims()[1] / 2;
394        let shift = adaln.clone().narrow(1, 0, half);
395        let scale = adaln.narrow(1, half, half);
396        let h_norm = burn_layer_norm_nd_no_affine(h);
397        let h_mod = h_norm * scale.add_scalar(1.0) + shift;
398        linear(
399            h_mod,
400            self.flow.out_w.clone(),
401            Some(self.flow.out_b.clone()),
402        )
403    }
404
405    fn flow_sample(&self, cond: Tensor<2>) -> Tensor<2> {
406        let dev = device();
407        let cond_emb = linear(
408            cond,
409            self.flow.cond_w.clone(),
410            Some(self.flow.cond_b.clone()),
411        );
412        let dt = 1.0f32 / self.n_steps as f32;
413        let mut x = Tensor::<2>::zeros([1, self.output_len], &dev);
414        let mut c_cur = silu(self.t_emb_table[0].clone().unsqueeze_dim::<2>(0) + cond_emb.clone());
415        for i in 0..self.n_steps {
416            let k1 = self.flow_net_eval(x.clone(), c_cur.clone());
417            let x_pred = x.clone() + k1.clone().mul_scalar(dt);
418            let c_next =
419                silu(self.t_emb_table[i + 1].clone().unsqueeze_dim::<2>(0) + cond_emb.clone());
420            let k2 = self.flow_net_eval(x_pred, c_next.clone());
421            x = x + (k1 + k2).mul_scalar(dt * 0.5);
422            c_cur = c_next;
423        }
424        x
425    }
426}
427
428fn burn_layer_norm_nd_no_affine(x: Tensor<2>) -> Tensor<2> {
429    zsfm_burn::norm::burn_layer_norm_no_affine(x, 1e-5)
430}
431
432fn build_rope_tables(head_dim: usize, max_seq: usize, theta: f64) -> (Tensor<2>, Tensor<2>) {
433    let dev = device();
434    let half = head_dim / 2;
435    let inv_freq: Vec<f64> = (0..half)
436        .map(|i| 1.0 / theta.powf(2.0 * i as f64 / head_dim as f64))
437        .collect();
438    let mut cos = vec![0.0f32; max_seq * head_dim];
439    let mut sin = vec![0.0f32; max_seq * head_dim];
440    for p in 0..max_seq {
441        for i in 0..half {
442            let angle = (p as f64 * inv_freq[i]) as f32;
443            let (s, c) = angle.sin_cos();
444            cos[p * head_dim + i] = c;
445            cos[p * head_dim + half + i] = c;
446            sin[p * head_dim + i] = s;
447            sin[p * head_dim + half + i] = s;
448        }
449    }
450    (
451        Tensor::<2>::from_data(TensorData::new(cos, [max_seq, head_dim]), &dev),
452        Tensor::<2>::from_data(TensorData::new(sin, [max_seq, head_dim]), &dev),
453    )
454}
455
456fn burn_causal_mask(n: usize) -> Tensor<4> {
457    let dev = device();
458    let data: Vec<f32> = (0..n)
459        .flat_map(|q| (0..n).map(move |k| if k <= q { 0.0f32 } else { f32::NEG_INFINITY }))
460        .collect();
461    Tensor::<2>::from_data(TensorData::new(data, [n, n]), &dev)
462        .unsqueeze_dim::<3>(0)
463        .unsqueeze_dim::<4>(0)
464}
465
466#[cfg(test)]
467mod tests {
468    use super::*;
469    use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
470
471    fn pseudo(n: usize, seed: u64) -> Vec<f32> {
472        let mut x = seed;
473        (0..n)
474            .map(|_| {
475                x = x
476                    .wrapping_mul(6364136223846793005)
477                    .wrapping_add(1442695040888963407);
478                ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
479            })
480            .collect()
481    }
482
483    fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
484        *seed += 1;
485        let data: Vec<f32> = pseudo(rows * cols, *seed)
486            .iter()
487            .map(|v| v * scale)
488            .collect();
489        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
490        w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
491    }
492
493    fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
494        *seed += 1;
495        let bytes: Vec<u8> = pseudo(n, *seed)
496            .iter()
497            .flat_map(|v| v.to_le_bytes())
498            .collect();
499        w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
500    }
501
502    #[test]
503    fn burn_forecast_golden() {
504        // Real widths (768/12x64), 1 layer, depth 1, 2 ODE steps, output 4.
505        let dir = std::env::temp_dir().join(format!("sundial-burn-{}", std::process::id()));
506        let _ = std::fs::remove_dir_all(&dir);
507        std::fs::create_dir_all(&dir).unwrap();
508        let path = dir.join("model.gguf");
509        let mut w = GGUFWriter::new();
510        w.add_metadata(
511            "general.architecture",
512            GGUFMetaValue::String("sundial".into()),
513        );
514        w.add_metadata("sundial1.block_count", GGUFMetaValue::Uint32(1));
515        w.add_metadata("sundial1.flow.num_sampling_steps", GGUFMetaValue::Uint32(2));
516        w.add_metadata("sundial1.flow.depth", GGUFMetaValue::Uint32(1));
517        w.add_metadata("sundial1.output_token_len", GGUFMetaValue::Uint32(4));
518        let mut seed = 700u64;
519        let s = 0.05;
520        put2(&mut w, "embed.hidden.weight", 32, 32, &mut seed, s);
521        put1(&mut w, "embed.hidden.bias", 32, &mut seed);
522        put2(&mut w, "embed.output.weight", 768, 32, &mut seed, s);
523        put1(&mut w, "embed.output.bias", 768, &mut seed);
524        put2(&mut w, "embed.skip.weight", 768, 32, &mut seed, s);
525        put1(&mut w, "embed.skip.bias", 768, &mut seed);
526        for t in ["attn_q", "attn_k", "attn_v"] {
527            put2(&mut w, &format!("blk.0.{t}.weight"), 768, 768, &mut seed, s);
528            put1(&mut w, &format!("blk.0.{t}.bias"), 768, &mut seed);
529        }
530        put2(&mut w, "blk.0.attn_out.weight", 768, 768, &mut seed, s);
531        put1(&mut w, "blk.0.attn_norm.weight", 768, &mut seed);
532        put1(&mut w, "blk.0.attn_norm.bias", 768, &mut seed);
533        put1(&mut w, "blk.0.ffn_norm.weight", 768, &mut seed);
534        put1(&mut w, "blk.0.ffn_norm.bias", 768, &mut seed);
535        put2(&mut w, "blk.0.ffn_gate.weight", 768, 768, &mut seed, s);
536        put2(&mut w, "blk.0.ffn_up.weight", 768, 768, &mut seed, s);
537        put2(&mut w, "blk.0.ffn_down.weight", 768, 768, &mut seed, s);
538        put1(&mut w, "norm.weight", 768, &mut seed);
539        put1(&mut w, "norm.bias", 768, &mut seed);
540        put2(&mut w, "flow.t_proj1.weight", 8, 256, &mut seed, s);
541        put1(&mut w, "flow.t_proj1.bias", 8, &mut seed);
542        put2(&mut w, "flow.t_proj2.weight", 8, 8, &mut seed, s);
543        put1(&mut w, "flow.t_proj2.bias", 8, &mut seed);
544        put2(&mut w, "flow.cond.weight", 8, 768, &mut seed, s);
545        put1(&mut w, "flow.cond.bias", 8, &mut seed);
546        put2(&mut w, "flow.in_proj.weight", 8, 4, &mut seed, s);
547        put1(&mut w, "flow.in_proj.bias", 8, &mut seed);
548        put1(&mut w, "flow.res.0.ln.weight", 8, &mut seed);
549        put1(&mut w, "flow.res.0.ln.bias", 8, &mut seed);
550        put2(&mut w, "flow.res.0.mlp1.weight", 8, 8, &mut seed, s);
551        put1(&mut w, "flow.res.0.mlp1.bias", 8, &mut seed);
552        put2(&mut w, "flow.res.0.mlp2.weight", 8, 8, &mut seed, s);
553        put1(&mut w, "flow.res.0.mlp2.bias", 8, &mut seed);
554        put2(&mut w, "flow.res.0.adaln.weight", 24, 8, &mut seed, s);
555        put1(&mut w, "flow.res.0.adaln.bias", 24, &mut seed);
556        put2(&mut w, "flow.out_linear.weight", 4, 8, &mut seed, s);
557        put1(&mut w, "flow.out_linear.bias", 4, &mut seed);
558        put2(&mut w, "flow.out_adaln.weight", 16, 8, &mut seed, s);
559        put1(&mut w, "flow.out_adaln.bias", 16, &mut seed);
560        let f = std::fs::File::create(&path).unwrap();
561        w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
562
563        let burn = BurnSundialModel::load(&path, None).unwrap();
564        let ctx: Vec<f32> = (0..16).map(|i| 10.0 + 0.5 * i as f32).collect();
565        let t0 = std::time::Instant::now();
566        let a = burn.forecast(&ctx).unwrap();
567        let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
568        let b = burn.forecast(&ctx).unwrap();
569        assert_eq!(a.len(), 4);
570        assert!(a.iter().all(|v| v.is_finite()));
571        assert_eq!(a, b, "forecast must be deterministic");
572        println!("sundial synthetic: burn {burn_ms:.1}ms out={a:?}");
573        // Golden recorded from the verified Burn implementation.
574        let golden: [f32; 4] = [12.953651, 14.637471, 13.848478, 14.506356];
575        assert_eq!(a.len(), golden.len());
576        for (x, y) in a.iter().zip(golden.iter()) {
577            assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
578        }
579        let _ = std::fs::remove_dir_all(&dir);
580    }
581}
zsfm_sundial/infer/mod.rs

zsfm_sundial/infer/
mod.rs
