zsfm_moirai2/infer/
mod.rs

1//! Moirai-2.0-R-small inference engine.
2//!
3//! Architecture: causal transformer decoder with ResidualBlock in/out projections,
4//! partial interleaved RoPE (first 32 of 64 head dims), BinaryAttentionBias,
5//! PackedStdScaler normalization, and multi-token (4 patches/token) decoding loop.
6
7use std::collections::HashMap;
8use std::sync::Mutex;
9use std::io::{BufReader, Read, Seek};
10use std::path::Path;
11
12use anyhow::{Context, Result};
13use candle_core::quantized::gguf_file;
14use candle_core::{DType, Device, Tensor};
15use rayon::prelude::*;
16
17use crate::config::Moirai2Config;
18
19mod rope;
20use rope::apply_partial_rope;
21
22// ---------------------------------------------------------------------------
23// Weight containers
24// ---------------------------------------------------------------------------
25
26struct ResidualBlockW {
27    hidden_w:   Tensor, // Python [hidden_dim, in_dim]
28    hidden_b:   Tensor, // [hidden_dim]
29    output_w:   Tensor, // Python [out_dim, hidden_dim]
30    output_b:   Tensor, // [out_dim]
31    residual_w: Tensor, // Python [out_dim, in_dim]
32    residual_b: Tensor, // [out_dim]
33}
34
35struct EncoderBlock {
36    norm1_w:         Tensor, // [d_model]
37    norm2_w:         Tensor,
38    attn_qkv_w:      Tensor, // fused [3*d_model, d_model]
39    attn_o_w:        Tensor,
40    attn_qn_w:       Tensor, // [head_dim]  per-head QK norm (shared across heads)
41    attn_kn_w:       Tensor,
42    attn_vbias_t: Tensor,      // [n_heads, 1, 1] same-variate bias, precomputed at load
43    ffn_fc1_w:       Tensor, // [d_ff, d_model]
44    ffn_fc2_w:       Tensor, // [d_model, d_ff]
45    ffn_gate_w:      Tensor, // [d_ff, d_model]
46}
47
48pub struct Moirai2Model {
49    device:   Device,
50    config:   Moirai2Config,
51    in_proj:  ResidualBlockW,
52    blocks:   Vec<EncoderBlock>,
53    norm_f_w: Tensor,
54    out_proj: ResidualBlockW,
55    rope_cos: Vec<f32>,  // [max_pos * half_rope] precomputed cosines
56    rope_sin: Vec<f32>,  // [max_pos * half_rope] precomputed sines
57    causal_mask_cache: Mutex<HashMap<usize, Tensor>>,
58}
59
60// ---------------------------------------------------------------------------
61// GGUF loading
62// ---------------------------------------------------------------------------
63
64fn load_t(
65    content: &gguf_file::Content,
66    reader: &mut (impl Read + Seek),
67    name: &str,
68    device: &Device,
69) -> Result<Tensor> {
70    zsfm_nn::load_tensor(content, reader, name, device, DType::F32)
71}
72
73fn load_residual_block(
74    content: &gguf_file::Content,
75    reader: &mut (impl Read + Seek),
76    prefix: &str,
77    device: &Device,
78) -> Result<ResidualBlockW> {
79    let p = |s: &str| format!("{prefix}.{s}");
80    Ok(ResidualBlockW {
81        hidden_w:   load_t(content, reader, &p("hidden.weight"), device)?,
82        hidden_b:   load_t(content, reader, &p("hidden.bias"), device)?,
83        output_w:   load_t(content, reader, &p("output.weight"), device)?,
84        output_b:   load_t(content, reader, &p("output.bias"), device)?,
85        residual_w: load_t(content, reader, &p("residual.weight"), device)?,
86        residual_b: load_t(content, reader, &p("residual.bias"), device)?,
87    })
88}
89
90impl Moirai2Model {
91    pub fn load(gguf_path: &Path, config: Moirai2Config) -> Result<Self> {
92        let device = Device::Cpu;
93        let file = std::fs::File::open(gguf_path)
94            .with_context(|| format!("open {}", gguf_path.display()))?;
95        let mut reader = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, file);
96        let content = gguf_file::Content::read(&mut reader).context("parse GGUF header")?;
97
98        let in_proj = load_residual_block(&content, &mut reader, "in_proj", &device)?;
99
100        let mut blocks = Vec::with_capacity(config.n_layers);
101        for n in 0..config.n_layers {
102            let p = |s: &str| format!("blk.{n}.{s}");
103            let q_w = load_t(&content, &mut reader, &p("attn_q.weight"), &device)?;
104            let k_w = load_t(&content, &mut reader, &p("attn_k.weight"), &device)?;
105            let v_w = load_t(&content, &mut reader, &p("attn_v.weight"), &device)?;
106            let attn_qkv_w = Tensor::cat(&[&q_w, &k_w, &v_w], 0)
107                .with_context(|| format!("qkv cat blk.{n}"))?;
108            let norm1_w    = load_t(&content, &mut reader, &p("norm1.weight"), &device)?;
109            let norm2_w    = load_t(&content, &mut reader, &p("norm2.weight"), &device)?;
110            let attn_o_w   = load_t(&content, &mut reader, &p("attn_o.weight"), &device)?;
111            let attn_qn_w  = load_t(&content, &mut reader, &p("attn_qn.weight"), &device)?;
112            let attn_kn_w  = load_t(&content, &mut reader, &p("attn_kn.weight"), &device)?;
113            let vbias_raw = load_t(&content, &mut reader, &p("attn_vbias.weight"), &device)?
114                .flatten_all()?.to_vec1::<f32>()?;
115            let n_heads = config.n_heads;
116            let same_var_bias: Vec<f32> = (0..n_heads).map(|h| vbias_raw[n_heads + h]).collect();
117            let attn_vbias_t = Tensor::from_vec(same_var_bias, (n_heads, 1, 1), &device)?;
118            let ffn_fc1_w  = load_t(&content, &mut reader, &p("ffn_fc1.weight"), &device)?;
119            let ffn_fc2_w  = load_t(&content, &mut reader, &p("ffn_fc2.weight"), &device)?;
120            let ffn_gate_w = load_t(&content, &mut reader, &p("ffn_gate.weight"), &device)?;
121            blocks.push(EncoderBlock {
122                norm1_w, norm2_w, attn_qkv_w, attn_o_w, attn_qn_w, attn_kn_w,
123                attn_vbias_t, ffn_fc1_w, ffn_fc2_w, ffn_gate_w,
124            });
125        }
126
127        let norm_f_w = load_t(&content, &mut reader, "norm_f.weight", &device)?;
128        let out_proj = load_residual_block(&content, &mut reader, "out_proj", &device)?;
129
130        let half_rope = config.rope_dim / 2;
131        let max_pos = config.max_ctx_tokens() + 256;
132        let inv_freq: Vec<f32> = (0..half_rope)
133            .map(|i| 1.0_f32 / 10000_f32.powf(2.0 * i as f32 / config.rope_dim as f32))
134            .collect();
135        let mut rope_cos = vec![0.0f32; max_pos * half_rope];
136        let mut rope_sin = vec![0.0f32; max_pos * half_rope];
137        for pos in 0..max_pos {
138            for i in 0..half_rope {
139                let theta = pos as f32 * inv_freq[i];
140                rope_cos[pos * half_rope + i] = theta.cos();
141                rope_sin[pos * half_rope + i] = theta.sin();
142            }
143        }
144
145        Ok(Self { device, config, in_proj, blocks, norm_f_w, out_proj,
146                  rope_cos, rope_sin, causal_mask_cache: Mutex::new(HashMap::new()) })
147    }
148
149    // -----------------------------------------------------------------------
150    // Forecasting
151    // -----------------------------------------------------------------------
152
153    pub fn forecast(&self, context: &[f32], horizon: usize) -> Result<Vec<f32>> {
154        let cfg = &self.config;
155        let ps = cfg.patch_size; // 16
156
157        // --- PackedStdScaler normalization (context only) ---
158        let max_ctx_len = cfg.max_ctx_tokens() * ps;
159        let ctx: &[f32] = if context.len() > max_ctx_len {
160            &context[context.len() - max_ctx_len..]
161        } else {
162            context
163        };
164        let n = ctx.len() as f64;
165        let loc_f64 = ctx.iter().map(|&v| v as f64).sum::<f64>() / n;
166        let var = ctx.iter().map(|&v| (v as f64 - loc_f64).powi(2)).sum::<f64>()
167            / (n - 1.0).max(1.0);
168        let scale = ((var + 1e-5_f64).sqrt()) as f32;
169        let loc   = loc_f64 as f32;
170
171        // --- Normalize and left-pad to multiple of patch_size ---
172        let ctx_norm: Vec<f32> = ctx.iter().map(|&v| (v - loc) / scale).collect();
173        let rem = ctx_norm.len() % ps;
174        let ctx_padded: Vec<f32> = if rem != 0 {
175            let mut padded = vec![0.0f32; ps - rem];
176            padded.extend_from_slice(&ctx_norm);
177            padded
178        } else {
179            ctx_norm
180        };
181        let n_ctx = ctx_padded.len() / ps;
182
183        let ctx_flat: Vec<f32> = (0..n_ctx)
184            .flat_map(|i| {
185                let mut tok = ctx_padded[i * ps..(i + 1) * ps].to_vec();
186                tok.extend(vec![1.0f32; ps]);
187                tok
188            })
189            .collect();
190        let ctx_time_ids: Vec<usize> = (0..n_ctx).collect();
191
192        let num_pt = cfg.num_predict_token; // 4
193        let num_q  = cfg.num_quantiles;     // 9
194        let mq     = cfg.median_quantile;   // 4
195        let n_future_patches = (horizon + ps - 1) / ps;
196
197        // --- Prefill: run context through transformer, collect KV cache ---
198        let input_t = Tensor::from_vec(ctx_flat, (n_ctx, ps * 2), &self.device)?;
199        let h_ctx = residual_block_fwd(&input_t, &self.in_proj)?;
200        let (h_enc, mut kv_cache) = self.prefill_encoder(h_ctx, &ctx_time_ids, n_ctx)?;
201
202        // Apply norm_f + out_proj to last context token only
203        let last_h_norm = zsfm_nn::rms_norm(&h_enc.narrow(0, n_ctx - 1, 1)?, Some(&self.norm_f_w), 1e-6)?;
204        let first_pred: Vec<f32> = residual_block_fwd(&last_h_norm, &self.out_proj)?
205            .flatten_all()?.to_vec1()?;
206
207        let mut collected_patches: Vec<Vec<f32>> = Vec::new();
208        let mut prev_patches: Vec<Vec<f32>> = (0..num_pt)
209            .map(|pt| {
210                let q_start = pt * num_q * ps + mq * ps;
211                first_pred[q_start..q_start + ps].to_vec()
212            })
213            .collect();
214        let n_take = n_future_patches.min(num_pt);
215        for patch in prev_patches.iter().take(n_take) {
216            collected_patches.push(patch.clone());
217        }
218
219        // --- Decode loop: num_pt=4 new tokens per step, cached K/V ---
220        let mut cached_len = n_ctx;
221        while collected_patches.len() < n_future_patches {
222            let new_time_ids: Vec<usize> = (cached_len..cached_len + num_pt).collect();
223            let new_flat: Vec<f32> = prev_patches.iter()
224                .flat_map(|patch| {
225                    let mut tok = patch.clone();
226                    tok.extend(vec![0.0f32; ps]);
227                    tok
228                })
229                .collect();
230            let new_t = Tensor::from_vec(new_flat, (num_pt, ps * 2), &self.device)?;
231            let h_in = residual_block_fwd(&new_t, &self.in_proj)?;
232            let h_dec = self.decode_encoder(h_in, &new_time_ids, &mut kv_cache, cached_len)?;
233
234            let last_h_norm = zsfm_nn::rms_norm(&h_dec.narrow(0, num_pt - 1, 1)?, Some(&self.norm_f_w), 1e-6)?;
235            let pred: Vec<f32> = residual_block_fwd(&last_h_norm, &self.out_proj)?
236                .flatten_all()?.to_vec1()?;
237
238            let new_patches: Vec<Vec<f32>> = (0..num_pt)
239                .map(|pt| {
240                    let q_start = pt * num_q * ps + mq * ps;
241                    pred[q_start..q_start + ps].to_vec()
242                })
243                .collect();
244
245            let need = n_future_patches - collected_patches.len();
246            for patch in new_patches.iter().take(need.min(num_pt)) {
247                collected_patches.push(patch.clone());
248            }
249            cached_len += num_pt;
250            prev_patches = new_patches;
251        }
252
253        // Flatten and denormalize
254        let result: Vec<f32> = collected_patches
255            .iter()
256            .flat_map(|p| p.iter().copied())
257            .take(horizon)
258            .map(|v| v * scale + loc)
259            .collect();
260
261        Ok(result)
262    }
263
264    // -----------------------------------------------------------------------
265    // Prefill: full forward pass collecting per-layer K/V tensors
266    // K/V stored as [n_heads, seq, head_dim]
267    // -----------------------------------------------------------------------
268
269    fn prefill_encoder(
270        &self,
271        mut h: Tensor,
272        time_ids: &[usize],
273        seq_len: usize,
274    ) -> Result<(Tensor, Vec<(Tensor, Tensor)>)> {
275        let mut kv_cache: Vec<(Tensor, Tensor)> = Vec::with_capacity(self.blocks.len());
276        for blk in &self.blocks {
277            let (h_out, k, v) = self.prefill_block(h, blk, time_ids, seq_len)?;
278            h = h_out;
279            kv_cache.push((k, v));
280        }
281        Ok((h, kv_cache))
282    }
283
284    fn prefill_block(
285        &self,
286        h: Tensor,
287        blk: &EncoderBlock,
288        time_ids: &[usize],
289        seq_len: usize,
290    ) -> Result<(Tensor, Tensor, Tensor)> {
291        let h_norm = zsfm_nn::rms_norm(&h, Some(&blk.norm1_w), 1e-6)?;
292        let (attn_out, k, v) = self.prefill_attn(&h_norm, blk, time_ids, seq_len)?;
293        let h = (h + attn_out)?;
294        let h_norm2 = zsfm_nn::rms_norm(&h, Some(&blk.norm2_w), 1e-6)?;
295        let ffn_out = zsfm_nn::swiglu_ffn(&h_norm2, &blk.ffn_fc1_w, &blk.ffn_fc2_w, &blk.ffn_gate_w)?;
296        Ok(((h + ffn_out)?, k, v))
297    }
298
299    fn prefill_attn(
300        &self,
301        h: &Tensor,
302        blk: &EncoderBlock,
303        time_ids: &[usize],
304        seq_len: usize,
305    ) -> Result<(Tensor, Tensor, Tensor)> {
306        let cfg = &self.config;
307        let n_heads  = cfg.n_heads;
308        let head_dim = cfg.head_dim;
309        let d_model  = cfg.d_model;
310        let rope_dim = cfg.rope_dim;
311
312        let qkv = zsfm_nn::linear_nobias(h, &blk.attn_qkv_w)?;
313        let q = qkv.narrow(1, 0, d_model)?;
314        let k = qkv.narrow(1, d_model, d_model)?;
315        let v = qkv.narrow(1, 2 * d_model, d_model)?;
316
317        let q = q.reshape((seq_len, n_heads, head_dim))?;
318        let k = k.reshape((seq_len, n_heads, head_dim))?;
319        let q = qk_norm_heads(&q, &blk.attn_qn_w, seq_len, n_heads, head_dim)?;
320        let k = qk_norm_heads(&k, &blk.attn_kn_w, seq_len, n_heads, head_dim)?;
321
322        let q = q.permute((1, 0, 2))?.contiguous()?;
323        let k = k.permute((1, 0, 2))?.contiguous()?;
324        let v = v.reshape((seq_len, n_heads, head_dim))?.permute((1, 0, 2))?.contiguous()?;
325
326        let q = apply_partial_rope(&q, time_ids, n_heads, head_dim, rope_dim, &self.device, &self.rope_cos, &self.rope_sin)?;
327        let k = apply_partial_rope(&k, time_ids, n_heads, head_dim, rope_dim, &self.device, &self.rope_cos, &self.rope_sin)?;
328
329        let scale = (head_dim as f64).sqrt();
330        let scores = q.matmul(&k.permute((0, 2, 1))?)?;
331        let scores = (scores / scale)?;
332        let causal = {
333            let mut cache = self.causal_mask_cache.lock().unwrap();
334            if !cache.contains_key(&seq_len) {
335                cache.insert(seq_len, make_causal_mask_tensor(seq_len, n_heads, &self.device)?);
336            }
337            cache[&seq_len].clone()
338        };
339        let scores = scores.broadcast_add(&causal)?;
340        let scores = scores.broadcast_add(&blk.attn_vbias_t)?;
341
342        let attn = candle_nn::ops::softmax_last_dim(&scores)?;
343        let out = attn.matmul(&v)?;
344        let out = out.permute((1, 0, 2))?.contiguous()?.reshape((seq_len, d_model))?;
345        Ok((zsfm_nn::linear_nobias(&out, &blk.attn_o_w)?, k, v))
346    }
347
348    // -----------------------------------------------------------------------
349    // KV-cached decode: process num_pt new tokens, append to cache
350    // -----------------------------------------------------------------------
351
352    fn decode_encoder(
353        &self,
354        mut h: Tensor,
355        new_time_ids: &[usize],
356        kv_cache: &mut Vec<(Tensor, Tensor)>,
357        cached_len: usize,
358    ) -> Result<Tensor> {
359        let new_len = new_time_ids.len();
360        for (li, blk) in self.blocks.iter().enumerate() {
361            // decode_block_kv returns the extended K/V (old cache + new tokens).
362            // Store directly — no second Tensor::cat needed.
363            let (h_out, k_new, v_new) =
364                self.decode_block_kv(h, blk, &kv_cache[li], new_time_ids, new_len, cached_len)?;
365            h = h_out;
366            kv_cache[li] = (k_new, v_new);
367        }
368        Ok(h)
369    }
370
371    fn decode_block_kv(
372        &self,
373        h: Tensor,
374        blk: &EncoderBlock,
375        cache: &(Tensor, Tensor),
376        new_time_ids: &[usize],
377        new_len: usize,
378        cached_len: usize,
379    ) -> Result<(Tensor, Tensor, Tensor)> {
380        let h_norm = zsfm_nn::rms_norm(&h, Some(&blk.norm1_w), 1e-6)?;
381        let (attn_out, k_new, v_new) =
382            self.decode_attn_kv(&h_norm, blk, cache, new_time_ids, new_len, cached_len)?;
383        let h = (h + attn_out)?;
384        let h_norm2 = zsfm_nn::rms_norm(&h, Some(&blk.norm2_w), 1e-6)?;
385        let ffn_out = zsfm_nn::swiglu_ffn(&h_norm2, &blk.ffn_fc1_w, &blk.ffn_fc2_w, &blk.ffn_gate_w)?;
386        Ok(((h + ffn_out)?, k_new, v_new))
387    }
388
389    fn decode_attn_kv(
390        &self,
391        h: &Tensor,
392        blk: &EncoderBlock,
393        cache: &(Tensor, Tensor),
394        new_time_ids: &[usize],
395        new_len: usize,
396        cached_len: usize,
397    ) -> Result<(Tensor, Tensor, Tensor)> {
398        let cfg = &self.config;
399        let n_heads  = cfg.n_heads;
400        let head_dim = cfg.head_dim;
401        let d_model  = cfg.d_model;
402        let rope_dim = cfg.rope_dim;
403
404        let qkv = zsfm_nn::linear_nobias(h, &blk.attn_qkv_w)?;
405        let q = qkv.narrow(1, 0, d_model)?;
406        let k = qkv.narrow(1, d_model, d_model)?;
407        let v = qkv.narrow(1, 2 * d_model, d_model)?;
408
409        let q = q.reshape((new_len, n_heads, head_dim))?;
410        let k = k.reshape((new_len, n_heads, head_dim))?;
411        let q = qk_norm_heads(&q, &blk.attn_qn_w, new_len, n_heads, head_dim)?;
412        let k = qk_norm_heads(&k, &blk.attn_kn_w, new_len, n_heads, head_dim)?;
413
414        let q = q.permute((1, 0, 2))?.contiguous()?;
415        let k = k.permute((1, 0, 2))?.contiguous()?;
416        let v = v.reshape((new_len, n_heads, head_dim))?.permute((1, 0, 2))?.contiguous()?;
417
418        let q = apply_partial_rope(&q, new_time_ids, n_heads, head_dim, rope_dim, &self.device, &self.rope_cos, &self.rope_sin)?;
419        let k = apply_partial_rope(&k, new_time_ids, n_heads, head_dim, rope_dim, &self.device, &self.rope_cos, &self.rope_sin)?;
420
421        // Extend cache: [n_heads, cached+new_len, head_dim].
422        // Return k_full/v_full so the caller can store them directly without a second cat.
423        let k_full = Tensor::cat(&[&cache.0, &k], 1)?;
424        let v_full = Tensor::cat(&[&cache.1, &v], 1)?;
425
426        let scale = (head_dim as f64).sqrt();
427        let scores = q.matmul(&k_full.permute((0, 2, 1))?)?; // [n_heads, new_len, cached+new_len]
428        let scores = (scores / scale)?;
429        let scores = add_decode_causal_mask(&scores, new_len, cached_len, n_heads, &self.device)?;
430        let scores = scores.broadcast_add(&blk.attn_vbias_t)?;
431
432        let attn = candle_nn::ops::softmax_last_dim(&scores)?;
433        let out = attn.matmul(&v_full)?; // [n_heads, new_len, head_dim]
434        let out = out.permute((1, 0, 2))?.contiguous()?.reshape((new_len, d_model))?;
435        Ok((zsfm_nn::linear_nobias(&out, &blk.attn_o_w)?, k_full, v_full))
436    }
437}
438
439// ---------------------------------------------------------------------------
440// Ops
441// ---------------------------------------------------------------------------
442
443fn residual_block_fwd(x: &Tensor, w: &ResidualBlockW) -> Result<Tensor> {
444    let hidden   = zsfm_nn::linear_bias(x, &w.hidden_w, &w.hidden_b)?.silu()?;
445    let output   = zsfm_nn::linear_bias(&hidden, &w.output_w, &w.output_b)?;
446    let residual = zsfm_nn::linear_bias(x, &w.residual_w, &w.residual_b)?;
447    Ok((output + residual)?)
448}
449
450fn qk_norm_heads(
451    x: &Tensor,
452    weight: &Tensor,
453    seq_len: usize,
454    n_heads: usize,
455    head_dim: usize,
456) -> Result<Tensor> {
457    let x_flat = x.reshape((seq_len * n_heads, head_dim))?;
458    let normed  = zsfm_nn::rms_norm(&x_flat, Some(weight), 1e-6)?;
459    Ok(normed.reshape((seq_len, n_heads, head_dim))?)
460}
461
462/// Build lower-triangular causal mask [1, seq, seq] — broadcast_add handles head dim.
463fn make_causal_mask_tensor(seq_len: usize, n_heads: usize, device: &Device) -> Result<Tensor> {
464    let _ = n_heads;
465    let mut mask = vec![0.0f32; seq_len * seq_len];
466    for i in 0..seq_len {
467        for j in (i + 1)..seq_len {
468            mask[i * seq_len + j] = f32::NEG_INFINITY;
469        }
470    }
471    Tensor::from_vec(mask, (1usize, seq_len, seq_len), device).map_err(anyhow::Error::from)
472}
473
474/// Decode causal mask: [n_heads, new_len, cached+new_len].
475/// New token q_rel attends to all cached tokens and new tokens 0..=q_rel.
476fn add_decode_causal_mask(
477    scores: &Tensor,
478    new_len: usize,
479    cached_len: usize,
480    n_heads: usize,
481    device: &Device,
482) -> Result<Tensor> {
483    let total = cached_len + new_len;
484    let mut mask = vec![0.0f32; new_len * total];
485    for q_rel in 0..new_len {
486        for k_abs in (cached_len + q_rel + 1)..total {
487            mask[q_rel * total + k_abs] = f32::NEG_INFINITY;
488        }
489    }
490    let mask_t = Tensor::from_vec(mask, (1usize, new_len, total), device)?;
491    Ok(scores.broadcast_add(&mask_t)?)
492}
493
494// ---------------------------------------------------------------------------
495// zsfm-core::Forecaster
496// ---------------------------------------------------------------------------
497
498impl zsfm_core::Forecaster for Moirai2Model {
499    type Config = Moirai2Config;
500
501    fn load(gguf_path: &Path, config: Moirai2Config) -> Result<Self> {
502        Moirai2Model::load(gguf_path, config)
503    }
504
505    /// Moirai-2 is channel-independent (each variate forecast separately) and
506    /// point-forecast only; `mask` is unused. Variates are forecast in parallel via `rayon`
507    /// — each is an independent forward pass through the same read-only weights (the only
508    /// shared mutable state, `causal_mask_cache`, is behind a `Mutex`).
509    fn forecast(
510        &self,
511        context: &[Vec<f32>],
512        _mask: &[Vec<bool>],
513        horizon: usize,
514    ) -> Result<zsfm_core::QuantileMatrix> {
515        anyhow::ensure!(!context.is_empty(), "context must have at least one variate");
516        let variates: Vec<Vec<f32>> = context
517            .par_iter()
518            .enumerate()
519            .map(|(vi, ctx)| -> Result<Vec<f32>> {
520                anyhow::ensure!(!ctx.is_empty(), "variate {vi} context must not be empty");
521                Moirai2Model::forecast(self, ctx, horizon).with_context(|| format!("forecast variate {vi}"))
522            })
523            .collect::<Result<Vec<_>>>()?;
524        Ok(vec![variates])
525    }
526}
zsfm_moirai2/infer/mod.rs

zsfm_moirai2/infer/
mod.rs
