zsfm_timesfm/infer/
mod.rs

1mod rope;
2
3use std::collections::HashMap;
4use std::sync::Mutex;
5use std::io::{BufReader, Read, Seek};
6use std::path::Path;
7
8use anyhow::{Context, Result};
9use candle_core::quantized::gguf_file;
10use candle_core::{DType, Device, Tensor, D};
11
12use rope::RopeCache;
13use zsfm_nn::{linear, load_weight, make_causal_mask};
14
15const NORM_EPS: f64 = 1e-6;
16const REVIN_TOL: f32 = 1e-6;
17
18// ---------------------------------------------------------------------------
19// Architecture constants (TimesFM 2.5 200M)
20// ---------------------------------------------------------------------------
21const D_MODEL: usize = 1280;
22const N_HEADS: usize = 16;
23const HEAD_DIM: usize = 80;
24const N_LAYERS: usize = 20;
25const INPUT_PATCH: usize = 32;
26const OUTPUT_PATCH: usize = 128;
27const N_OUTPUTS: usize = 10;
28const DECODE_IDX: usize = 5;
29const M_PATCHES: usize = OUTPUT_PATCH / INPUT_PATCH; // 4
30const ROPE_THETA: f64 = 10000.0;
31const MAX_SEQ: usize = 16384 / INPUT_PATCH + 256; // generous upper bound
32
33// ---------------------------------------------------------------------------
34// Weight structs
35// ---------------------------------------------------------------------------
36
37struct ResidualBlockW {
38    hidden_w: Tensor,
39    hidden_b: Option<Tensor>,
40    output_w: Tensor,
41    output_b: Option<Tensor>,
42    skip_w: Tensor,
43    skip_b: Option<Tensor>,
44}
45
46struct AttnW {
47    qkv_w: Tensor,        // [3*D_MODEL, D_MODEL]
48    out_w: Tensor,        // [D_MODEL, D_MODEL]
49    q_norm_w: Tensor,     // [HEAD_DIM]
50    k_norm_w: Tensor,     // [HEAD_DIM]
51    q_scale: Tensor,      // [HEAD_DIM] cached as Tensor to avoid alloc per attention forward
52    pre_norm_w: Tensor,   // [D_MODEL]
53    post_norm_w: Tensor,  // [D_MODEL]
54}
55
56struct FfnW {
57    up_w: Tensor,         // [D_MODEL, D_MODEL]
58    down_w: Tensor,       // [D_MODEL, D_MODEL]
59    pre_norm_w: Tensor,   // [D_MODEL]
60    post_norm_w: Tensor,  // [D_MODEL]
61}
62
63struct BlockW {
64    attn: AttnW,
65    ffn: FfnW,
66}
67
68/// TimesFM 2.5 200M's architecture is fixed (see the module-level `const`s above) — there is no
69/// `config.json` and hence no `InferConfig`/builder here, unlike every other model in the
70/// workspace. `TimesFMModel::load` takes just the GGUF path.
71pub struct TimesFMModel {
72    device: Device,
73    rope: RopeCache,
74    tokenizer: ResidualBlockW,
75    blocks: Vec<BlockW>,
76    out_point: ResidualBlockW,
77    causal_mask_cache: Mutex<HashMap<usize, Tensor>>,
78}
79
80// ---------------------------------------------------------------------------
81// GGUF loading helpers
82// ---------------------------------------------------------------------------
83
84fn load_tensor(
85    content: &gguf_file::Content,
86    reader: &mut (impl Read + Seek),
87    name: &str,
88    device: &Device,
89) -> Result<Tensor> {
90    zsfm_nn::load_tensor(content, reader, name, device, DType::F32)
91}
92
93fn load_residual_block(
94    content: &gguf_file::Content,
95    reader: &mut (impl Read + Seek),
96    prefix: &str,
97    hidden_d_out: usize,
98    out_d_out: usize,
99    with_bias: bool,
100    device: &Device,
101) -> Result<ResidualBlockW> {
102    let hidden_w = load_weight(content, reader, &format!("{prefix}.hidden.weight"), hidden_d_out, device)?;
103    let hidden_b = if with_bias {
104        Some(load_tensor(content, reader, &format!("{prefix}.hidden.bias"), device)?)
105    } else {
106        None
107    };
108    let output_w = load_weight(content, reader, &format!("{prefix}.output.weight"), out_d_out, device)?;
109    let output_b = if with_bias {
110        Some(load_tensor(content, reader, &format!("{prefix}.output.bias"), device)?)
111    } else {
112        None
113    };
114    let skip_w = load_weight(content, reader, &format!("{prefix}.skip.weight"), out_d_out, device)?;
115    let skip_b = if with_bias {
116        Some(load_tensor(content, reader, &format!("{prefix}.skip.bias"), device)?)
117    } else {
118        None
119    };
120    Ok(ResidualBlockW { hidden_w, hidden_b, output_w, output_b, skip_w, skip_b })
121}
122
123fn softplus(x: f32) -> f32 {
124    if x > 20.0 { x } else { (1.0f32 + x.exp()).ln() }
125}
126
127fn compute_q_scale(raw: Vec<f32>) -> Vec<f32> {
128    let factor = 1.442695041f32 / (HEAD_DIM as f32).sqrt();
129    raw.into_iter().map(|x| factor * softplus(x)).collect()
130}
131
132impl TimesFMModel {
133    pub fn load(gguf_path: &Path) -> Result<Self> {
134        let device = Device::Cpu;
135        let file = std::fs::File::open(gguf_path)
136            .with_context(|| format!("open {}", gguf_path.display()))?;
137        let mut reader = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, file);
138        let content = gguf_file::Content::read(&mut reader).context("parse GGUF header")?;
139
140        // Tokenizer ResidualBlock: in=64, hidden=1280, out=1280, bias=true
141        let tokenizer = load_residual_block(
142            &content, &mut reader, "tokenizer", D_MODEL, D_MODEL, true, &device,
143        )?;
144
145        // 20 transformer blocks
146        let mut blocks = Vec::with_capacity(N_LAYERS);
147        for n in 0..N_LAYERS {
148            let b = format!("blk.{n}");
149            let qkv_raw = load_weight(&content, &mut reader, &format!("{b}.attn_qkv.weight"), 3 * D_MODEL, &device)?;
150            let out_w = load_weight(&content, &mut reader, &format!("{b}.attn_out.weight"), D_MODEL, &device)?;
151            let q_norm_w = load_tensor(&content, &mut reader, &format!("{b}.attn_q_norm.weight"), &device)?;
152            let k_norm_w = load_tensor(&content, &mut reader, &format!("{b}.attn_k_norm.weight"), &device)?;
153            let q_scale_raw = load_tensor(&content, &mut reader, &format!("{b}.attn_q_scale.weight"), &device)?
154                .to_vec1::<f32>()?;
155            let q_scale = Tensor::from_vec(compute_q_scale(q_scale_raw), (HEAD_DIM,), &device)?;
156            let pre_norm_w  = load_tensor(&content, &mut reader, &format!("{b}.pre_attn_norm.weight"), &device)?;
157            let post_norm_w = load_tensor(&content, &mut reader, &format!("{b}.post_attn_norm.weight"), &device)?;
158
159            let up_w   = load_weight(&content, &mut reader, &format!("{b}.ffn_up.weight"),   D_MODEL, &device)?;
160            let down_w = load_weight(&content, &mut reader, &format!("{b}.ffn_down.weight"), D_MODEL, &device)?;
161            let pre_ff_norm_w  = load_tensor(&content, &mut reader, &format!("{b}.pre_ff_norm.weight"),  &device)?;
162            let post_ff_norm_w = load_tensor(&content, &mut reader, &format!("{b}.post_ff_norm.weight"), &device)?;
163
164            blocks.push(BlockW {
165                attn: AttnW {
166                    qkv_w: qkv_raw,
167                    out_w,
168                    q_norm_w,
169                    k_norm_w,
170                    q_scale,
171                    pre_norm_w,
172                    post_norm_w,
173                },
174                ffn: FfnW {
175                    up_w,
176                    down_w,
177                    pre_norm_w: pre_ff_norm_w,
178                    post_norm_w: post_ff_norm_w,
179                },
180            });
181        }
182
183        // Output projection (no bias, hidden=1280, out=1280)
184        let out_point = load_residual_block(
185            &content, &mut reader, "out_point", D_MODEL, D_MODEL, false, &device,
186        )?;
187
188        let rope = RopeCache::new(HEAD_DIM, MAX_SEQ, ROPE_THETA, &device)?;
189
190        Ok(Self { device, rope, tokenizer, blocks, out_point, causal_mask_cache: Mutex::new(HashMap::new()) })
191    }
192
193    // -----------------------------------------------------------------------
194    // Public API
195    // -----------------------------------------------------------------------
196
197    /// Forecast a univariate time series.
198    ///
199    /// Returns `[N_OUTPUTS][prediction_length]` where index 0 is the point
200    /// forecast and indices 1-9 are quantile forecasts (q0.1 … q0.9).
201    pub fn forecast(&self, context: &[f32], prediction_length: usize) -> Result<Vec<Vec<f32>>> {
202        let p = INPUT_PATCH;
203        let o = OUTPUT_PATCH;
204        let q = N_OUTPUTS;
205        let m = M_PATCHES;
206
207        // Pad front so len is divisible by p
208        let len_front = if context.len() % p == 0 { 0 } else { p - context.len() % p };
209        let mut vals = vec![0.0f32; len_front + context.len()];
210        vals[len_front..].copy_from_slice(context);
211        let mut mask = vec![true; len_front]; // True = masked (padding)
212        mask.extend(vec![false; context.len()]);
213
214        let n_ctx = (len_front + context.len()) / p;
215
216        // Compute cumulative running stats per patch
217        let mut patch_mus = Vec::with_capacity(n_ctx);
218        let mut patch_sigmas = Vec::with_capacity(n_ctx);
219        let (mut rs_n, mut rs_mu, mut rs_sigma) = (0.0f32, 0.0f32, 0.0f32);
220        for i in 0..n_ctx {
221            let pv = &vals[i * p..(i + 1) * p];
222            let pm = &mask[i * p..(i + 1) * p];
223            (rs_n, rs_mu, rs_sigma) = update_running_stats(rs_n, rs_mu, rs_sigma, pv, pm);
224            patch_mus.push(rs_mu);
225            patch_sigmas.push(rs_sigma);
226        }
227        let (last_n, last_mu, last_sigma) = (rs_n, rs_mu, rs_sigma);
228
229        // Build normalized tokenizer input for context patches
230        let ctx_input = build_tokenizer_input(&vals, &mask, n_ctx, p, &patch_mus, &patch_sigmas);
231
232        // Prefill: full forward pass collecting per-layer KV cache
233        let num_decode_steps = if prediction_length <= o { 0 } else { (prediction_length - 1) / o };
234        let total_steps = 1 + num_decode_steps;
235        let mut all_outputs: Vec<Vec<[f32; N_OUTPUTS]>> = Vec::with_capacity(total_steps);
236
237        let (ctx_out, mut kv_cache) = self.prefill(&ctx_input, n_ctx)?;
238        let ctx_denorm = denorm_flat(&ctx_out, &patch_mus, &patch_sigmas, o, q)?;
239        all_outputs.push(ctx_denorm[n_ctx - 1].clone());
240
241        // AR decode: process M_PATCHES=4 new patches per step with cached K/V
242        let (mut ar_n, mut ar_mu, mut ar_sigma) = (last_n, last_mu, last_sigma);
243        let mut last_ar: Vec<f32> = ctx_denorm[n_ctx - 1].iter().map(|row| row[DECODE_IDX]).collect();
244
245        for step in 0..num_decode_steps {
246            let new_vals_flat = last_ar.clone();
247            let new_mask_flat = vec![false; o];
248
249            let mut new_mus = Vec::with_capacity(m);
250            let mut new_sigmas = Vec::with_capacity(m);
251            for mi in 0..m {
252                let pv = &new_vals_flat[mi * p..(mi + 1) * p];
253                let pm = &new_mask_flat[mi * p..(mi + 1) * p];
254                (ar_n, ar_mu, ar_sigma) = update_running_stats(ar_n, ar_mu, ar_sigma, pv, pm);
255                new_mus.push(ar_mu);
256                new_sigmas.push(ar_sigma);
257            }
258
259            let new_input = build_tokenizer_input(
260                &new_vals_flat, &new_mask_flat, m, p, &new_mus, &new_sigmas,
261            );
262            let rope_offset = n_ctx + m * step;
263            let new_out = self.decode_chunk(&new_input, &mut kv_cache, rope_offset)?;
264            let last_m_denorm = denorm_flat(&new_out, &new_mus, &new_sigmas, o, q)?;
265
266            let step_out = &last_m_denorm[m - 1];
267            last_ar = step_out.iter().map(|row| row[DECODE_IDX]).collect();
268            all_outputs.push(step_out.clone());
269        }
270
271        // Assemble [N_OUTPUTS][prediction_length]
272        let total_available = all_outputs.len() * o;
273        let n_take = prediction_length.min(total_available);
274        let mut result: Vec<Vec<f32>> = vec![Vec::with_capacity(n_take); q];
275        'outer: for step in &all_outputs {
276            for timestep in step {
277                for qi in 0..q {
278                    if result[qi].len() >= prediction_length { break 'outer; }
279                    result[qi].push(timestep[qi]);
280                }
281            }
282        }
283        Ok(result)
284    }
285
286    // -----------------------------------------------------------------------
287    // Prefill: full forward + collect per-layer KV cache
288    // KV tensors stored as [1, N_HEADS, n_patches, HEAD_DIM]
289    // -----------------------------------------------------------------------
290
291    fn prefill(
292        &self,
293        tokenizer_input: &Tensor,  // [n_patches, 2*INPUT_PATCH]
294        n_patches: usize,
295    ) -> Result<(Tensor, Vec<(Tensor, Tensor)>)> {
296        let x = forward_residual_block(tokenizer_input, &self.tokenizer, true)?;
297        let mut hidden = x.unsqueeze(0)?;
298        let causal = {
299            let mut cache = self.causal_mask_cache.lock().unwrap();
300            if !cache.contains_key(&n_patches) {
301                cache.insert(n_patches, make_causal_mask(n_patches, 0, &self.device)?);
302            }
303            cache[&n_patches].clone()
304        };
305        let mut kv_cache: Vec<(Tensor, Tensor)> = Vec::with_capacity(N_LAYERS);
306        for block in &self.blocks {
307            let (h_out, k, v) = self.prefill_block(hidden, block, n_patches, &causal)?;
308            hidden = h_out;
309            kv_cache.push((k, v));
310        }
311        let out_seq = hidden.squeeze(0)?;
312        let out = forward_residual_block(&out_seq, &self.out_point, false)?;
313        Ok((out, kv_cache))
314    }
315
316    fn prefill_block(
317        &self,
318        x: Tensor,
319        w: &BlockW,
320        n_patches: usize,
321        causal: &Tensor,
322    ) -> Result<(Tensor, Tensor, Tensor)> {
323        let normed = rms_norm(&x, &w.attn.pre_norm_w, NORM_EPS)?;
324        let (attn_out, k, v) = self.prefill_attn(normed, &w.attn, n_patches, causal)?;
325        let attn_out = (rms_norm(&attn_out, &w.attn.post_norm_w, NORM_EPS)? + &x)?;
326        let normed_ff = rms_norm(&attn_out, &w.ffn.pre_norm_w, NORM_EPS)?;
327        let ff_out = self.forward_ffn(normed_ff, &w.ffn)?;
328        let out = (rms_norm(&ff_out, &w.ffn.post_norm_w, NORM_EPS)? + &attn_out)?;
329        Ok((out, k, v))
330    }
331
332    fn prefill_attn(
333        &self,
334        x: Tensor,       // [1, n_patches, D_MODEL] pre-normed
335        w: &AttnW,
336        n_patches: usize,
337        causal: &Tensor,
338    ) -> Result<(Tensor, Tensor, Tensor)> {  // (output, K, V) K/V: [1,N_HEADS,n,HEAD_DIM]
339        let qkv = linear(&x, &w.qkv_w, None)?;
340        let q = qkv.narrow(D::Minus1, 0, D_MODEL)?;
341        let k = qkv.narrow(D::Minus1, D_MODEL, D_MODEL)?;
342        let v = qkv.narrow(D::Minus1, 2 * D_MODEL, D_MODEL)?;
343
344        let q = q.reshape((1, n_patches, N_HEADS, HEAD_DIM))?;
345        let k = k.reshape((1, n_patches, N_HEADS, HEAD_DIM))?;
346        let v = v.reshape((1, n_patches, N_HEADS, HEAD_DIM))?;
347
348        let q = self.rope.apply(&q, 0)?;
349        let k = self.rope.apply(&k, 0)?;
350        let q = rms_norm(&q, &w.q_norm_w, NORM_EPS)?;
351        let k = rms_norm(&k, &w.k_norm_w, NORM_EPS)?;
352
353        let q = q.broadcast_mul(&w.q_scale)?;
354
355        let q = q.permute([0, 2, 1, 3])?.contiguous()?;   // [1, N_HEADS, n, HEAD_DIM]
356        let k = k.permute([0, 2, 1, 3])?.contiguous()?;
357        let v = v.permute([0, 2, 1, 3])?.contiguous()?;
358
359        let scores = q.matmul(&k.transpose(D::Minus1, D::Minus2)?)?;
360        let scores = scores.broadcast_add(causal)?;
361        let attn_w = candle_nn::ops::softmax(&scores, D::Minus1)?;
362        let ctx = attn_w.matmul(&v)?;
363        let ctx = ctx.permute([0, 2, 1, 3])?.contiguous()?.reshape((1, n_patches, D_MODEL))?;
364        Ok((linear(&ctx, &w.out_w, None)?, k, v))
365    }
366
367    // -----------------------------------------------------------------------
368    // KV-cached decode: process M_PATCHES new patches, append to cache
369    // -----------------------------------------------------------------------
370
371    fn decode_chunk(
372        &self,
373        new_input: &Tensor,                    // [M_PATCHES, 2*INPUT_PATCH]
374        kv_cache: &mut Vec<(Tensor, Tensor)>,  // mutated: K/V grow by M_PATCHES per call
375        rope_offset: usize,                    // = n_ctx + M_PATCHES * step
376    ) -> Result<Tensor> {                      // → [M_PATCHES, D_MODEL]
377        let x = forward_residual_block(new_input, &self.tokenizer, true)?;
378        let mut hidden = x.unsqueeze(0)?;
379        let cached_len = rope_offset;
380        let decode_mask = make_decode_mask(M_PATCHES, cached_len, &self.device)?;
381        for (li, block) in self.blocks.iter().enumerate() {
382            // decode_block_kv returns the extended K/V (old cache + new M_PATCHES).
383            // Store directly — no second Tensor::cat needed.
384            let (h_out, k_new, v_new) =
385                self.decode_block_kv(hidden, block, &kv_cache[li], rope_offset, &decode_mask)?;
386            hidden = h_out;
387            kv_cache[li] = (k_new, v_new);
388        }
389        let out_seq = hidden.squeeze(0)?;
390        forward_residual_block(&out_seq, &self.out_point, false)
391    }
392
393    fn decode_block_kv(
394        &self,
395        x: Tensor,
396        w: &BlockW,
397        cache: &(Tensor, Tensor),
398        rope_offset: usize,
399        decode_mask: &Tensor,
400    ) -> Result<(Tensor, Tensor, Tensor)> {
401        let normed = rms_norm(&x, &w.attn.pre_norm_w, NORM_EPS)?;
402        let (attn_out, k_new, v_new) =
403            self.decode_attn_kv(normed, &w.attn, cache, rope_offset, decode_mask)?;
404        let attn_out = (rms_norm(&attn_out, &w.attn.post_norm_w, NORM_EPS)? + &x)?;
405        let normed_ff = rms_norm(&attn_out, &w.ffn.pre_norm_w, NORM_EPS)?;
406        let ff_out = self.forward_ffn(normed_ff, &w.ffn)?;
407        let out = (rms_norm(&ff_out, &w.ffn.post_norm_w, NORM_EPS)? + &attn_out)?;
408        Ok((out, k_new, v_new))
409    }
410
411    fn decode_attn_kv(
412        &self,
413        x: Tensor,                  // [1, M_PATCHES, D_MODEL] pre-normed
414        w: &AttnW,
415        cache: &(Tensor, Tensor),   // ([1,N_HEADS,cached,HEAD_DIM], same)
416        rope_offset: usize,
417        mask: &Tensor,              // [1,1,M_PATCHES,cached+M_PATCHES]
418    ) -> Result<(Tensor, Tensor, Tensor)> {  // (output, K_new, V_new)
419        let qkv = linear(&x, &w.qkv_w, None)?;
420        let q = qkv.narrow(D::Minus1, 0, D_MODEL)?;
421        let k = qkv.narrow(D::Minus1, D_MODEL, D_MODEL)?;
422        let v = qkv.narrow(D::Minus1, 2 * D_MODEL, D_MODEL)?;
423
424        let q = q.reshape((1, M_PATCHES, N_HEADS, HEAD_DIM))?;
425        let k = k.reshape((1, M_PATCHES, N_HEADS, HEAD_DIM))?;
426        let v = v.reshape((1, M_PATCHES, N_HEADS, HEAD_DIM))?;
427
428        let q = self.rope.apply(&q, rope_offset)?;
429        let k = self.rope.apply(&k, rope_offset)?;
430        let q = rms_norm(&q, &w.q_norm_w, NORM_EPS)?;
431        let k = rms_norm(&k, &w.k_norm_w, NORM_EPS)?;
432
433        let q = q.broadcast_mul(&w.q_scale)?;
434
435        let q = q.permute([0, 2, 1, 3])?.contiguous()?;   // [1, N_HEADS, M, HEAD_DIM]
436        let k = k.permute([0, 2, 1, 3])?.contiguous()?;
437        let v = v.permute([0, 2, 1, 3])?.contiguous()?;
438
439        // Extend cache: [1, N_HEADS, cached+M, HEAD_DIM].
440        // Return k_full/v_full so the caller can store them directly without a second cat.
441        let k_full = Tensor::cat(&[&cache.0, &k], 2)?;
442        let v_full = Tensor::cat(&[&cache.1, &v], 2)?;
443
444        let scores = q.matmul(&k_full.transpose(D::Minus1, D::Minus2)?)?;
445        let scores = scores.broadcast_add(mask)?;
446        let attn_w = candle_nn::ops::softmax(&scores, D::Minus1)?;
447        let ctx = attn_w.matmul(&v_full)?;
448        let ctx = ctx.permute([0, 2, 1, 3])?.contiguous()?.reshape((1, M_PATCHES, D_MODEL))?;
449        Ok((linear(&ctx, &w.out_w, None)?, k_full, v_full))
450    }
451
452    fn forward_ffn(&self, x: Tensor, w: &FfnW) -> Result<Tensor> {
453        let h = linear(&x, &w.up_w, None)?;
454        let h = candle_nn::ops::silu(&h)?;
455        linear(&h, &w.down_w, None)
456    }
457}
458
459// ---------------------------------------------------------------------------
460// ResidualBlock forward (tokenizer has bias, output projections do not)
461// ---------------------------------------------------------------------------
462
463fn forward_residual_block(x: &Tensor, w: &ResidualBlockW, _has_bias: bool) -> Result<Tensor> {
464    let h = linear(x, &w.hidden_w, w.hidden_b.as_ref())?;
465    let h = candle_nn::ops::silu(&h)?;
466    let out = linear(&h, &w.output_w, w.output_b.as_ref())?;
467    let skip = linear(x, &w.skip_w, w.skip_b.as_ref())?;
468    Ok((out + skip)?)
469}
470
471// ---------------------------------------------------------------------------
472// Primitive ops
473// ---------------------------------------------------------------------------
474
475fn rms_norm(x: &Tensor, weight: &Tensor, eps: f64) -> Result<Tensor> {
476    zsfm_nn::rms_norm(x, Some(weight), eps)
477}
478
479/// Decode-step causal mask: shape [1,1,new_len,cache_len+new_len].
480/// New query i can attend to all cached tokens plus new tokens 0..=i.
481fn make_decode_mask(new_len: usize, cache_len: usize, device: &Device) -> Result<Tensor> {
482    let total = cache_len + new_len;
483    let data: Vec<f32> = (0..new_len).flat_map(|q_rel| {
484        (0..total).map(move |k_abs| {
485            if k_abs <= cache_len + q_rel { 0.0f32 } else { f32::NEG_INFINITY }
486        })
487    }).collect();
488    Ok(Tensor::from_vec(data, (new_len, total), device)?
489        .unsqueeze(0)?.unsqueeze(0)?)
490}
491
492// ---------------------------------------------------------------------------
493// Running statistics (mirrors Python update_running_stats)
494// ---------------------------------------------------------------------------
495
496fn update_running_stats(n: f32, mu: f32, sigma: f32, vals: &[f32], mask: &[bool]) -> (f32, f32, f32) {
497    let mut inc_n = 0.0f32;
498    let mut sum_x = 0.0f32;
499    for (i, &m) in mask.iter().enumerate() {
500        if !m {
501            inc_n += 1.0;
502            sum_x += vals[i];
503        }
504    }
505    let inc_mu = if inc_n == 0.0 { 0.0 } else { sum_x / inc_n };
506    let inc_var = if inc_n == 0.0 { 0.0 } else {
507        mask.iter().enumerate()
508            .filter(|(_, &m)| !m)
509            .map(|(i, _)| (vals[i] - inc_mu).powi(2))
510            .sum::<f32>() / inc_n
511    };
512    let inc_sigma = inc_var.sqrt();
513
514    let new_n = n + inc_n;
515    let safe_new_n = if new_n == 0.0 { 1.0 } else { new_n };
516    let new_mu = if new_n == 0.0 { 0.0 } else { (n * mu + inc_mu * inc_n) / safe_new_n };
517    let t1 = n * sigma.powi(2);
518    let t2 = inc_n * inc_sigma.powi(2);
519    let t3 = n * (mu - new_mu).powi(2);
520    let t4 = inc_n * (inc_mu - new_mu).powi(2);
521    let new_var = if new_n == 0.0 { 0.0 } else { (t1 + t2 + t3 + t4) / safe_new_n };
522    (new_n, new_mu, new_var.max(0.0).sqrt())
523}
524
525// ---------------------------------------------------------------------------
526// Input preparation
527// ---------------------------------------------------------------------------
528
529/// Build tokenizer input tensor [n_patches, 2*INPUT_PATCH=64].
530///
531/// For each patch: cat([normed_values, mask_as_float], dim=-1).
532/// Masked (padded) positions → value=0.0, mask=1.0.
533fn build_tokenizer_input(
534    vals: &[f32],
535    mask: &[bool],
536    n_patches: usize,
537    p: usize,
538    mus: &[f32],
539    sigmas: &[f32],
540) -> Tensor {
541    let mut data = vec![0.0f32; n_patches * 2 * p];
542    for pi in 0..n_patches {
543        let mu = mus[pi];
544        let sigma = sigmas[pi];
545        let sigma_safe = if sigma < REVIN_TOL { 1.0f32 } else { sigma };
546        for i in 0..p {
547            let idx = pi * p + i;
548            let is_masked = mask[idx];
549            let normed = if is_masked { 0.0 } else { (vals[idx] - mu) / sigma_safe };
550            let base = pi * 2 * p;
551            data[base + i] = normed;                // value channel
552            data[base + p + i] = if is_masked { 1.0 } else { 0.0 }; // mask channel
553        }
554    }
555    Tensor::from_vec(data, (n_patches, 2 * p), &Device::Cpu).expect("build_tokenizer_input")
556}
557
558// ---------------------------------------------------------------------------
559// Output denormalization
560// ---------------------------------------------------------------------------
561
562/// Denorm output [n_patches, D_MODEL] → flat vec of [n_patches][O][Q] triples.
563///
564/// Each patch i is denormed: out * sigma[i] + mu[i].
565/// Then reshaped to [OUTPUT_PATCH, N_OUTPUTS].
566fn denorm_flat(
567    output: &Tensor,          // [n_patches, D_MODEL]
568    mus: &[f32],
569    sigmas: &[f32],
570    o: usize,
571    q: usize,
572) -> Result<Vec<Vec<[f32; N_OUTPUTS]>>> {
573    let flat = output.flatten_all()?.to_vec1::<f32>()?;
574    let n = mus.len();
575    let mut result: Vec<Vec<[f32; N_OUTPUTS]>> = vec![vec![[0.0f32; N_OUTPUTS]; o]; n];
576    for pi in 0..n {
577        let mu = mus[pi];
578        let sigma = sigmas[pi];
579        for ti in 0..o {
580            for qi in 0..q {
581                let raw = flat[pi * (o * q) + ti * q + qi];
582                result[pi][ti][qi] = raw * sigma + mu;
583            }
584        }
585    }
586    Ok(result)
587}
588
589// ---------------------------------------------------------------------------
590// zsfm-core::Forecaster
591// ---------------------------------------------------------------------------
592
593impl zsfm_core::Forecaster for TimesFMModel {
594    /// No config.json exists for this model — architecture is fixed at compile time.
595    type Config = ();
596
597    fn load(gguf_path: &Path, _config: ()) -> Result<Self> {
598        TimesFMModel::load(gguf_path)
599    }
600
601    /// TimesFM's `forecast()` is univariate-only and always treats the context as fully
602    /// observed; `mask` is unused.
603    fn forecast(
604        &self,
605        context: &[Vec<f32>],
606        _mask: &[Vec<bool>],
607        horizon: usize,
608    ) -> Result<zsfm_core::QuantileMatrix> {
609        anyhow::ensure!(context.len() == 1, "TimesFMModel only supports univariate forecasting (1 variate)");
610        let outputs = TimesFMModel::forecast(self, &context[0], horizon)?; // [N_OUTPUTS][horizon]
611        Ok(outputs.into_iter().map(|row| vec![row]).collect())
612    }
613}
zsfm_timesfm/infer/mod.rs

zsfm_timesfm/infer/
mod.rs
