zsfm_chronos/infer/
mod.rs

1//! Chronos-2 inference on the Burn backend (Flex CPU).
2//!
3//! Bit-exact port of `super`: patch embedding with time features,
4//! encoder blocks (TimeSelfAttention with Llama RoPE + GroupSelfAttention +
5//! FFN), quantile output head. Weights load straight from GGUF via
6//! [`zsfm_burn::GgufWeights`].
7
8use std::path::Path;
9
10use anyhow::{bail, Context, Result};
11use burn::tensor::{activation, Device, Tensor, TensorData};
12use zsfm_burn::linear::burn_linear_nd;
13use zsfm_burn::norm::burn_rms_norm_nd;
14use zsfm_burn::rope::{apply_llama_rope, rope_tables};
15use zsfm_burn::{open_gguf_weights, GgufWeights};
16
17use crate::config::Chronos2Config;
18
19#[derive(Clone, Debug)]
20pub struct InferConfig {
21    d_model: usize,
22    d_kv: usize,
23    d_ff: usize,
24    num_layers: usize,
25    num_heads: usize,
26    layer_norm_eps: f64,
27    rope_theta: f64,
28    patch_size: usize,
29    patch_stride: usize,
30    context_length: usize,
31    quantiles: Vec<f32>,
32    use_reg_token: bool,
33    use_arcsinh: bool,
34    time_encoding_scale: usize,
35    dense_act_fn: String,
36}
37
38/// `Chronos2Config` already resolved every default (via serde) or failed to parse if a
39/// genuinely required field was missing, so this mapping is infallible and needs no further
40/// defaulting of its own.
41impl From<&Chronos2Config> for InferConfig {
42    fn from(c2: &Chronos2Config) -> Self {
43        let cc = &c2.chronos_config;
44        InferConfig {
45            d_model: c2.d_model as usize,
46            d_kv: c2.d_kv as usize,
47            d_ff: c2.d_ff as usize,
48            num_layers: c2.num_layers as usize,
49            num_heads: c2.num_heads as usize,
50            layer_norm_eps: c2.layer_norm_epsilon,
51            rope_theta: c2.rope_theta,
52            patch_size: cc.input_patch_size as usize,
53            patch_stride: cc.input_patch_stride as usize,
54            context_length: cc.context_length as usize,
55            quantiles: cc.quantiles.clone(),
56            use_reg_token: cc.use_reg_token,
57            use_arcsinh: cc.use_arcsinh,
58            time_encoding_scale: c2.time_encoding_scale() as usize,
59            dense_act_fn: c2.dense_act_fn().to_string(),
60        }
61    }
62}
63
64impl InferConfig {
65    fn inner_dim(&self) -> usize {
66        self.num_heads * self.d_kv
67    }
68
69    pub fn patch_size(&self) -> usize {
70        self.patch_size
71    }
72    pub fn patch_stride(&self) -> usize {
73        self.patch_stride
74    }
75    pub fn context_length(&self) -> usize {
76        self.context_length
77    }
78    pub fn quantiles(&self) -> &[f32] {
79        &self.quantiles
80    }
81}
82
83fn device() -> Device {
84    Device::flex()
85}
86
87fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
88    weights.load_2d(name)
89}
90
91fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
92    weights.load_1d(name)
93}
94
95fn load_w(weights: &GgufWeights, name: &str, expected_d_out: usize) -> Result<Tensor<2>> {
96    weights.load_weight_2d(name, expected_d_out)
97}
98
99fn linear<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
100    burn_linear_nd(x, w, b)
101}
102
103fn rms<const D: usize>(x: Tensor<D>, w: Tensor<1>, eps: f64) -> Tensor<D> {
104    burn_rms_norm_nd(x, w, eps as f32)
105}
106
107fn softmax_last<const D: usize>(x: Tensor<D>) -> Tensor<D> {
108    activation::softmax(x, D - 1)
109}
110
111fn apply_act<const D: usize>(x: Tensor<D>, name: &str) -> Result<Tensor<D>> {
112    Ok(match name {
113        "relu" => activation::relu(x),
114        "gelu" => activation::gelu(x),
115        "gelu_new" | "gelu_pytorch_tanh" => activation::gelu_approximate(x),
116        "silu" | "swish" => activation::silu(x),
117        other => bail!("unsupported activation function: {other}"),
118    })
119}
120
121struct BurnResidualBlockW {
122    hidden_w: Tensor<2>,
123    hidden_b: Tensor<1>,
124    output_w: Tensor<2>,
125    output_b: Tensor<1>,
126    skip_w: Tensor<2>,
127    skip_b: Tensor<1>,
128}
129
130struct BurnAttnW {
131    qkv_w: Tensor<2>,
132    o_w: Tensor<2>,
133    norm_w: Tensor<1>,
134}
135
136struct BurnFfnW {
137    wi_w: Tensor<2>,
138    wo_w: Tensor<2>,
139    norm_w: Tensor<1>,
140}
141
142struct BurnBlockW {
143    time_attn: BurnAttnW,
144    group_attn: BurnAttnW,
145    ffn: BurnFfnW,
146}
147
148pub struct BurnChronosModel {
149    config: InferConfig,
150    rope_cos: Tensor<2>,
151    rope_sin: Tensor<2>,
152    token_embd: Tensor<2>,
153    input_patch: BurnResidualBlockW,
154    blocks: Vec<BurnBlockW>,
155    enc_norm: Tensor<1>,
156    output_patch: BurnResidualBlockW,
157}
158
159fn load_residual_block(
160    weights: &GgufWeights,
161    prefix: &str,
162    d_out_h: usize,
163    d_out: usize,
164) -> Result<BurnResidualBlockW> {
165    Ok(BurnResidualBlockW {
166        hidden_w: load_w(weights, &format!("{prefix}.hidden.weight"), d_out_h)?,
167        hidden_b: load_v(weights, &format!("{prefix}.hidden.bias"))?,
168        output_w: load_w(weights, &format!("{prefix}.output.weight"), d_out)?,
169        output_b: load_v(weights, &format!("{prefix}.output.bias"))?,
170        skip_w: load_w(weights, &format!("{prefix}.skip.weight"), d_out)?,
171        skip_b: load_v(weights, &format!("{prefix}.skip.bias"))?,
172    })
173}
174
175fn load_attn(
176    weights: &GgufWeights,
177    blk: usize,
178    kind: &str,
179    d_model: usize,
180    inner_dim: usize,
181) -> Result<BurnAttnW> {
182    let qw = load_w(weights, &format!("blk.{blk}.{kind}.q.weight"), inner_dim)?;
183    let kw = load_w(weights, &format!("blk.{blk}.{kind}.k.weight"), inner_dim)?;
184    let vw = load_w(weights, &format!("blk.{blk}.{kind}.v.weight"), inner_dim)?;
185    let qkv_w = Tensor::cat(vec![qw, kw, vw], 0);
186    Ok(BurnAttnW {
187        qkv_w,
188        o_w: load_w(weights, &format!("blk.{blk}.{kind}.o.weight"), d_model)?,
189        norm_w: load_v(weights, &format!("blk.{blk}.{kind}_norm.weight"))?,
190    })
191}
192
193impl BurnChronosModel {
194    pub fn patch_size(&self) -> usize {
195        self.config.patch_size()
196    }
197    pub fn patch_stride(&self) -> usize {
198        self.config.patch_stride()
199    }
200    pub fn context_length(&self) -> usize {
201        self.config.context_length()
202    }
203    pub fn quantiles(&self) -> &[f32] {
204        self.config.quantiles()
205    }
206
207    pub fn load(gguf_path: &Path, config: InferConfig) -> Result<Self> {
208        let weights = open_gguf_weights(gguf_path)
209            .with_context(|| format!("open {}", gguf_path.display()))?;
210
211        let d = config.d_model;
212        let ff = config.d_ff;
213        let id = config.num_heads * config.d_kv;
214        let ps = config.patch_size;
215        let nq = config.quantiles.len();
216
217        let token_embd = load_t(&weights, "token_embd.weight")?;
218        let input_patch = load_residual_block(&weights, "input_patch", ff, d)?;
219
220        let mut blocks = Vec::with_capacity(config.num_layers);
221        for n in 0..config.num_layers {
222            let time_attn = load_attn(&weights, n, "time_attn", d, id)?;
223            let group_attn = load_attn(&weights, n, "group_attn", d, id)?;
224            blocks.push(BurnBlockW {
225                time_attn,
226                group_attn,
227                ffn: BurnFfnW {
228                    wi_w: load_w(&weights, &format!("blk.{n}.ffn.wi.weight"), ff)?,
229                    wo_w: load_w(&weights, &format!("blk.{n}.ffn.wo.weight"), d)?,
230                    norm_w: load_v(&weights, &format!("blk.{n}.ffn_norm.weight"))?,
231                },
232            });
233        }
234
235        let enc_norm = load_v(&weights, "enc_norm.weight")?;
236        let out_d = nq * ps;
237        let output_patch = load_residual_block(&weights, "output_patch", ff, out_d)?;
238
239        let (rope_cos, rope_sin) = rope_tables(config.d_kv, 8192, config.rope_theta, &device())?;
240        Ok(Self {
241            config,
242            rope_cos,
243            rope_sin,
244            token_embd,
245            input_patch,
246            blocks,
247            enc_norm,
248            output_patch,
249        })
250    }
251
252    pub fn forecast(&self, context: &[f32], prediction_length: usize) -> Result<Vec<Vec<f32>>> {
253        let cfg = &self.config;
254        let ps = cfg.patch_size;
255        let stride = cfg.patch_stride;
256        let n_quantiles = cfg.quantiles.len();
257        let n_out_patches = prediction_length.div_ceil(ps);
258
259        let (normalized, loc, scale) = instance_norm(context, cfg.use_arcsinh);
260        let padded = pad_for_patching(&normalized, ps);
261        let n_ctx_patches = (padded.len() - ps) / stride + 1;
262
263        let ctx_features = build_patch_features(
264            &padded,
265            n_ctx_patches,
266            ps,
267            stride,
268            -((n_ctx_patches * ps) as f32),
269            cfg.time_encoding_scale as f32,
270            true,
271        );
272        let fut_features =
273            build_patch_features_future(n_out_patches, ps, 0.0, cfg.time_encoding_scale as f32);
274
275        let dev = device();
276        let ctx_tensor =
277            Tensor::<2>::from_data(TensorData::new(ctx_features, [n_ctx_patches, 3 * ps]), &dev);
278        let fut_tensor =
279            Tensor::<2>::from_data(TensorData::new(fut_features, [n_out_patches, 3 * ps]), &dev);
280        let ctx_embeds = self.forward_residual_block(ctx_tensor, &self.input_patch)?;
281        let fut_embeds = self.forward_residual_block(fut_tensor, &self.input_patch)?;
282
283        let seq = if cfg.use_reg_token {
284            let reg: Tensor<2> = self.token_embd.clone().narrow(0, 1, 1);
285            Tensor::cat(vec![ctx_embeds, reg, fut_embeds], 0)
286        } else {
287            Tensor::cat(vec![ctx_embeds, fut_embeds], 0)
288        };
289        let total_seq = seq.dims()[0];
290        let hidden = seq.unsqueeze_dim::<3>(0);
291        let hidden = self.forward_encoder(hidden, total_seq)?;
292        let forecast_embeds: Tensor<3> = hidden.narrow(1, total_seq - n_out_patches, n_out_patches);
293        let forecast_embeds: Tensor<2> = forecast_embeds.squeeze_dim(0);
294        let quantile_raw = self.forward_residual_block(forecast_embeds, &self.output_patch)?;
295
296        let total_out = n_out_patches * ps;
297        let qmat: Tensor<3> = quantile_raw.reshape([n_out_patches, n_quantiles, ps]);
298        let qmat: Tensor<3> = qmat.permute([1, 0, 2]);
299        let qmat: Tensor<2> = qmat.reshape([n_quantiles, total_out]);
300        let qmat: Tensor<2> = if total_out > prediction_length {
301            qmat.narrow(1, 0, prediction_length)
302        } else {
303            qmat
304        };
305        let data: Vec<f32> = qmat
306            .to_data()
307            .try_to_vec()
308            .map_err(|e| anyhow::anyhow!("{e:?}"))?;
309        let mut result = vec![vec![0.0f32; prediction_length]; n_quantiles];
310        for q in 0..n_quantiles {
311            for t in 0..prediction_length {
312                let v = data[q * prediction_length + t] as f64;
313                let v = if cfg.use_arcsinh { v.sinh() } else { v };
314                result[q][t] = (v as f32) * scale + loc;
315            }
316        }
317        Ok(result)
318    }
319
320    fn forward_encoder(&self, mut x: Tensor<3>, seq_len: usize) -> Result<Tensor<3>> {
321        for blk in &self.blocks {
322            x = self.forward_block(x, blk, seq_len)?;
323        }
324        Ok(rms(x, self.enc_norm.clone(), self.config.layer_norm_eps))
325    }
326
327    fn forward_block(&self, x: Tensor<3>, blk: &BurnBlockW, seq_len: usize) -> Result<Tensor<3>> {
328        let normed = rms(
329            x.clone(),
330            blk.time_attn.norm_w.clone(),
331            self.config.layer_norm_eps,
332        );
333        let attn_out = self.forward_time_attn(normed, &blk.time_attn, seq_len);
334        let x = x + attn_out;
335        let normed = rms(
336            x.clone(),
337            blk.group_attn.norm_w.clone(),
338            self.config.layer_norm_eps,
339        );
340        let grp_out = self.forward_group_attn_univariate(normed, &blk.group_attn);
341        let x = x + grp_out;
342        let normed = rms(
343            x.clone(),
344            blk.ffn.norm_w.clone(),
345            self.config.layer_norm_eps,
346        );
347        let ffn_out = self.forward_ffn(normed, &blk.ffn)?;
348        Ok(x + ffn_out)
349    }
350
351    fn forward_time_attn(&self, x: Tensor<3>, w: &BurnAttnW, seq_len: usize) -> Tensor<3> {
352        let id = self.config.num_heads * self.config.d_kv;
353        let nh = self.config.num_heads;
354        let dkv = self.config.d_kv;
355        let qkv = linear(x, w.qkv_w.clone(), None);
356        let q: Tensor<3> = qkv.clone().narrow(2, 0, id);
357        let k: Tensor<3> = qkv.clone().narrow(2, id, id);
358        let v: Tensor<3> = qkv.narrow(2, 2 * id, id);
359        let pack = |t: Tensor<3>| t.reshape([1, seq_len, nh, dkv]).permute([0, 2, 1, 3]);
360        let (q, k, v) = (pack(q), pack(k), pack(v));
361        let q = apply_llama_rope(q, self.rope_cos.clone(), self.rope_sin.clone(), 0);
362        let k = apply_llama_rope(k, self.rope_cos.clone(), self.rope_sin.clone(), 0);
363        let scores = q.matmul(k.transpose());
364        let attn = softmax_last(scores);
365        let out = attn.matmul(v);
366        let out: Tensor<3> = out.permute([0, 2, 1, 3]).reshape([1, seq_len, id]);
367        linear(out, w.o_w.clone(), None)
368    }
369
370    fn forward_group_attn_univariate(&self, x: Tensor<3>, w: &BurnAttnW) -> Tensor<3> {
371        let id = self.config.num_heads * self.config.d_kv;
372        let qkv = linear(x, w.qkv_w.clone(), None);
373        let v: Tensor<3> = qkv.narrow(2, 2 * id, id);
374        linear(v, w.o_w.clone(), None)
375    }
376
377    fn forward_ffn(&self, x: Tensor<3>, w: &BurnFfnW) -> Result<Tensor<3>> {
378        let h = linear(x, w.wi_w.clone(), None);
379        let h = apply_act(h, &self.config.dense_act_fn)?;
380        Ok(linear(h, w.wo_w.clone(), None))
381    }
382
383    fn forward_residual_block(&self, x: Tensor<2>, w: &BurnResidualBlockW) -> Result<Tensor<2>> {
384        let h = linear(x.clone(), w.hidden_w.clone(), Some(w.hidden_b.clone()));
385        let h = apply_act(h, &self.config.dense_act_fn)?;
386        let out = linear(h, w.output_w.clone(), Some(w.output_b.clone()));
387        let skip = linear(x, w.skip_w.clone(), Some(w.skip_b.clone()));
388        Ok(out + skip)
389    }
390}
391
392fn instance_norm(x: &[f32], use_arcsinh: bool) -> (Vec<f32>, f32, f32) {
393    let n = x.len() as f64;
394    let loc = x.iter().map(|&v| v as f64).sum::<f64>() / n;
395    let var = x.iter().map(|&v| (v as f64 - loc).powi(2)).sum::<f64>() / n;
396    let scale = (var.sqrt() as f32).max(1e-5);
397    let mut out: Vec<f32> = x.iter().map(|&v| (v as f64 - loc) as f32 / scale).collect();
398    if use_arcsinh {
399        for v in &mut out {
400            *v = (*v as f64).asinh() as f32;
401        }
402    }
403    (out, loc as f32, scale)
404}
405
406fn pad_for_patching(x: &[f32], patch_size: usize) -> Vec<f32> {
407    let rem = x.len() % patch_size;
408    if rem == 0 {
409        x.to_vec()
410    } else {
411        let pad_len = patch_size - rem;
412        let mut out = vec![f32::NAN; pad_len];
413        out.extend_from_slice(x);
414        out
415    }
416}
417
418fn build_patch_features(
419    padded: &[f32],
420    n_patches: usize,
421    patch_size: usize,
422    stride: usize,
423    time_start: f32,
424    time_scale: f32,
425    observed: bool,
426) -> Vec<f32> {
427    let feature_dim = 3 * patch_size;
428    let mut out = vec![0.0f32; n_patches * feature_dim];
429    for p in 0..n_patches {
430        let offset = p * stride;
431        let base = p * feature_dim;
432        for i in 0..patch_size {
433            let t = offset + i;
434            let global_t = time_start + t as f32;
435            out[base + i] = global_t / time_scale;
436            let v = padded.get(t).copied().unwrap_or(0.0);
437            let is_obs = observed && v.is_finite();
438            out[base + patch_size + i] = if is_obs { v } else { 0.0 };
439            out[base + 2 * patch_size + i] = if is_obs { 1.0 } else { 0.0 };
440        }
441    }
442    out
443}
444
445fn build_patch_features_future(
446    n_patches: usize,
447    patch_size: usize,
448    time_start: f32,
449    time_scale: f32,
450) -> Vec<f32> {
451    let feature_dim = 3 * patch_size;
452    let mut out = vec![0.0f32; n_patches * feature_dim];
453    for p in 0..n_patches {
454        let offset = p * patch_size;
455        let base = p * feature_dim;
456        for i in 0..patch_size {
457            let global_t = time_start + (offset + i) as f32;
458            out[base + i] = global_t / time_scale;
459        }
460    }
461    out
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467    use crate::config::{Chronos2Config, ChronosInnerConfig};
468    use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
469
470    fn pseudo(n: usize, seed: u64) -> Vec<f32> {
471        let mut x = seed;
472        (0..n)
473            .map(|_| {
474                x = x
475                    .wrapping_mul(6364136223846793005)
476                    .wrapping_add(1442695040888963407);
477                ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
478            })
479            .collect()
480    }
481
482    fn tiny_config() -> InferConfig {
483        let c2 = Chronos2Config {
484            d_model: 16,
485            d_kv: 8,
486            d_ff: 32,
487            num_layers: 1,
488            num_heads: 2,
489            layer_norm_epsilon: 1e-6,
490            rope_theta: 10000.0,
491            feed_forward_proj: "relu".into(),
492            chronos_config: ChronosInnerConfig {
493                context_length: 64,
494                input_patch_size: 4,
495                output_patch_size: 4,
496                input_patch_stride: 2,
497                quantiles: vec![0.1, 0.5, 0.9],
498                use_reg_token: false,
499                use_arcsinh: false,
500                max_output_patches: 1,
501                time_encoding_scale: None,
502            },
503        };
504        InferConfig::from(&c2)
505    }
506
507    fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
508        *seed += 1;
509        let data: Vec<f32> = pseudo(rows * cols, *seed)
510            .iter()
511            .map(|v| v * scale)
512            .collect();
513        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
514        w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
515    }
516
517    fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
518        *seed += 1;
519        let bytes: Vec<u8> = pseudo(n, *seed)
520            .iter()
521            .flat_map(|v| v.to_le_bytes())
522            .collect();
523        w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
524    }
525
526    fn put_res(
527        w: &mut GGUFWriter,
528        prefix: &str,
529        h: usize,
530        d_in: usize,
531        d_out: usize,
532        seed: &mut u64,
533        s: f32,
534    ) {
535        put2(w, &format!("{prefix}.hidden.weight"), h, d_in, seed, s);
536        put1(w, &format!("{prefix}.hidden.bias"), h, seed);
537        put2(w, &format!("{prefix}.output.weight"), d_out, h, seed, s);
538        put1(w, &format!("{prefix}.output.bias"), d_out, seed);
539        put2(w, &format!("{prefix}.skip.weight"), d_out, d_in, seed, s);
540        put1(w, &format!("{prefix}.skip.bias"), d_out, seed);
541    }
542
543    fn put_attn(
544        w: &mut GGUFWriter,
545        blk: usize,
546        kind: &str,
547        d: usize,
548        id: usize,
549        seed: &mut u64,
550        s: f32,
551    ) {
552        for t in ["q", "k", "v"] {
553            put2(w, &format!("blk.{blk}.{kind}.{t}.weight"), id, d, seed, s);
554        }
555        put2(w, &format!("blk.{blk}.{kind}.o.weight"), d, id, seed, s);
556        put1(w, &format!("blk.{blk}.{kind}_norm.weight"), d, seed);
557    }
558
559    #[test]
560    fn burn_forecast_golden() {
561        // d=16, id=16, ff=32, ps=4, stride=2, 3 quantiles, 1 layer.
562        let dir = std::env::temp_dir().join(format!("chronos-burn-{}", std::process::id()));
563        let _ = std::fs::remove_dir_all(&dir);
564        std::fs::create_dir_all(&dir).unwrap();
565        let path = dir.join("model.gguf");
566        let mut w = GGUFWriter::new();
567        w.add_metadata(
568            "general.architecture",
569            GGUFMetaValue::String("chronos".into()),
570        );
571        let mut seed = 900u64;
572        let s = 0.2;
573        put2(&mut w, "token_embd.weight", 4, 16, &mut seed, s);
574        put_res(&mut w, "input_patch", 32, 12, 16, &mut seed, s);
575        put_attn(&mut w, 0, "time_attn", 16, 16, &mut seed, s);
576        put_attn(&mut w, 0, "group_attn", 16, 16, &mut seed, s);
577        put2(&mut w, "blk.0.ffn.wi.weight", 32, 16, &mut seed, s);
578        put2(&mut w, "blk.0.ffn.wo.weight", 16, 32, &mut seed, s);
579        put1(&mut w, "blk.0.ffn_norm.weight", 16, &mut seed);
580        put1(&mut w, "enc_norm.weight", 16, &mut seed);
581        put_res(&mut w, "output_patch", 32, 16, 12, &mut seed, s);
582        let f = std::fs::File::create(&path).unwrap();
583        w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
584
585        let cfg = tiny_config();
586        let burn = BurnChronosModel::load(&path, cfg).unwrap();
587        let ctx: Vec<f32> = (0..8).map(|i| 10.0 + 0.5 * i as f32).collect();
588        let t0 = std::time::Instant::now();
589        let a = burn.forecast(&ctx, 4).unwrap();
590        let ms_a = t0.elapsed().as_secs_f64() * 1000.0;
591        let b = burn.forecast(&ctx, 4).unwrap();
592        assert_eq!(a.len(), 3);
593        for row in &a {
594            assert_eq!(row.len(), 4);
595        }
596        assert!(a.iter().flatten().all(|v| v.is_finite()));
597        assert!(a.iter().flatten().any(|&v| v != 0.0));
598        assert_eq!(a, b, "forecast must be deterministic");
599        println!("chronos synthetic: burn {ms_a:.3}ms out={a:?}");
600        // Golden values recorded from the verified Burn port (ex-candle parity 1e-4).
601        let golden: [[f32; 4]; 3] = [
602            [9.894369, 8.8220005, 9.853922, 9.345992],
603            [9.070834, 9.884812, 9.241523, 9.49992],
604            [9.611428, 10.309262, 10.598283, 9.180033],
605        ];
606        assert_eq!(a.len(), golden.len());
607        for (row, grow) in a.iter().zip(golden.iter()) {
608            assert_eq!(row.len(), grow.len());
609            for (x, y) in row.iter().zip(grow.iter()) {
610                assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
611            }
612        }
613        let _ = std::fs::remove_dir_all(&dir);
614    }
615}
zsfm_chronos/infer/mod.rs

zsfm_chronos/infer/
mod.rs
