zsfm_tirex/infer/
mod.rs

1use std::io::{BufReader, Read, Seek};
2use std::path::Path;
3
4use anyhow::{Context, Result};
5use candle_core::quantized::gguf_file;
6use candle_core::{DType, Device, Tensor};
7
8use simdeez::prelude::*;
9use simdeez::math::{SimdMathF32Core, SimdMathF32Hyperbolic};
10
11use crate::config::TiRexConfig;
12
13// ---------------------------------------------------------------------------
14// Weight structs
15// ---------------------------------------------------------------------------
16
17struct Block {
18    norm_slstm: Vec<f32>,      // [D]
19    fizo_w: Vec<f32>,          // [NH, 4*DH, DH] — f,i,z,o gates fused at load time
20    slstm_kernel_t: Vec<f32>,  // [NH, NG*DH, DH] = [4, 512, 128] — transposed for SIMD dot
21    slstm_bias: Vec<f32>,      // [NG*NH*DH = 2048] in [NG, NH, DH] order
22    group_norm_w: Vec<f32>,   // [D] (learned offset, applied as 1 + w)
23    norm_ffn: Vec<f32>,       // [D]
24    ffn_gate_w: Tensor,       // [UP, D]
25    ffn_up_w: Tensor,         // [UP, D]
26    ffn_down_w: Tensor,       // [D, UP]
27}
28
29struct EmbedBlock {
30    hidden_w: Tensor,   // [H_DIM, IN_DIM]
31    hidden_b: Vec<f32>, // [H_DIM]
32    output_w: Tensor,   // [OUT_DIM, H_DIM]
33    output_b: Vec<f32>, // [OUT_DIM]
34    residual_w: Tensor, // [OUT_DIM, IN_DIM]
35    residual_b: Vec<f32>, // [OUT_DIM]
36}
37
38pub struct TiRexModel {
39    device: Device,
40    config: TiRexConfig,
41    in_emb: EmbedBlock,
42    blocks: Vec<Block>,
43    out_norm: Vec<f32>, // [D]
44    out_emb: EmbedBlock,
45}
46
47// ---------------------------------------------------------------------------
48// GGUF loading
49// ---------------------------------------------------------------------------
50
51fn load_t(
52    content: &gguf_file::Content,
53    reader: &mut (impl Read + Seek),
54    name: &str,
55    device: &Device,
56) -> Result<Tensor> {
57    zsfm_nn::load_tensor(content, reader, name, device, DType::F32)
58}
59
60fn load_vec(
61    content: &gguf_file::Content,
62    reader: &mut (impl Read + Seek),
63    name: &str,
64    device: &Device,
65) -> Result<Vec<f32>> {
66    Ok(load_t(content, reader, name, device)?.flatten_all()?.to_vec1()?)
67}
68
69fn load_embed_block(
70    content: &gguf_file::Content,
71    reader: &mut (impl Read + Seek),
72    prefix: &str,
73    device: &Device,
74) -> Result<EmbedBlock> {
75    let p = |s: &str| format!("{prefix}.{s}");
76    Ok(EmbedBlock {
77        hidden_w:   load_t(&content, reader, &p("hidden.weight"), device)?,
78        hidden_b:   load_vec(&content, reader, &p("hidden.bias"), device)?,
79        output_w:   load_t(&content, reader, &p("output.weight"), device)?,
80        output_b:   load_vec(&content, reader, &p("output.bias"), device)?,
81        residual_w: load_t(&content, reader, &p("residual.weight"), device)?,
82        residual_b: load_vec(&content, reader, &p("residual.bias"), device)?,
83    })
84}
85
86impl TiRexModel {
87    pub fn load(gguf_path: &Path, config: TiRexConfig) -> Result<Self> {
88        let device = Device::Cpu;
89        let file = std::fs::File::open(gguf_path)
90            .with_context(|| format!("open {}", gguf_path.display()))?;
91        let mut reader = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, file);
92        let content = gguf_file::Content::read(&mut reader).context("parse GGUF header")?;
93
94        let in_emb = load_embed_block(&content, &mut reader, "in_emb", &device)?;
95
96        let nh  = config.num_heads;
97        let dh  = config.head_dim();
98        let wpp = dh * dh; // weights per head per gate
99
100        let mut blocks = Vec::with_capacity(config.num_blocks);
101        for n in 0..config.num_blocks {
102            let p = |s: &str| format!("blk.{n}.{s}");
103
104            // Load the 4 gate weights, concatenate into [NH, 4*DH, DH] at load time
105            let fgate_w = load_vec(&content, &mut reader, &p("fgate.weight"), &device)?;
106            let igate_w = load_vec(&content, &mut reader, &p("igate.weight"), &device)?;
107            let zgate_w = load_vec(&content, &mut reader, &p("zgate.weight"), &device)?;
108            let ogate_w = load_vec(&content, &mut reader, &p("ogate.weight"), &device)?;
109            let mut fizo_w = vec![0.0f32; nh * 4 * wpp];
110            for h in 0..nh {
111                fizo_w[h*4*wpp..       h*4*wpp+wpp  ].copy_from_slice(&fgate_w[h*wpp..(h+1)*wpp]);
112                fizo_w[h*4*wpp+wpp..   h*4*wpp+2*wpp].copy_from_slice(&igate_w[h*wpp..(h+1)*wpp]);
113                fizo_w[h*4*wpp+2*wpp.. h*4*wpp+3*wpp].copy_from_slice(&zgate_w[h*wpp..(h+1)*wpp]);
114                fizo_w[h*4*wpp+3*wpp.. h*4*wpp+4*wpp].copy_from_slice(&ogate_w[h*wpp..(h+1)*wpp]);
115            }
116
117            let raw_kernel = load_vec(&content, &mut reader, &p("slstm_kernel"), &device)?;
118            // Transpose kernel [NH, DH, NG*DH] → [NH, NG*DH, DH] so the inner dot-product
119            // dim (di) is contiguous, enabling simd_dot in compute_ry.
120            let ng = 4usize;
121            let mut slstm_kernel_t = vec![0.0f32; nh * ng * dh * dh];
122            for head in 0..nh {
123                for gate_d in 0..(ng * dh) {
124                    for di in 0..dh {
125                        slstm_kernel_t[head * ng * dh * dh + gate_d * dh + di] =
126                            raw_kernel[head * dh * ng * dh + di * ng * dh + gate_d];
127                    }
128                }
129            }
130
131            blocks.push(Block {
132                norm_slstm:    load_vec(&content, &mut reader, &p("norm_slstm"), &device)?,
133                fizo_w,
134                slstm_kernel_t,
135                slstm_bias:    load_vec(&content, &mut reader, &p("slstm_bias"), &device)?,
136                group_norm_w: load_vec(&content, &mut reader, &p("group_norm"), &device)?,
137                norm_ffn:     load_vec(&content, &mut reader, &p("norm_ffn"), &device)?,
138                ffn_gate_w:   load_t(&content, &mut reader, &p("ffn_gate.weight"), &device)?,
139                ffn_up_w:     load_t(&content, &mut reader, &p("ffn_up.weight"), &device)?,
140                ffn_down_w:   load_t(&content, &mut reader, &p("ffn_down.weight"), &device)?,
141            });
142        }
143
144        let out_norm = load_vec(&content, &mut reader, "out_norm", &device)?;
145        let out_emb = load_embed_block(&content, &mut reader, "out_emb", &device)?;
146
147        Ok(TiRexModel { device, config, in_emb, blocks, out_norm, out_emb })
148    }
149
150    /// Forecast quantiles and the median for a single time series.
151    ///
152    /// Returns `(quantiles, median)` where:
153    /// - `quantiles`: `[prediction_length, num_quantiles]` in the config's quantile order
154    /// - `median`: `[prediction_length]` (the 0.5 quantile row, for convenience — despite the
155    ///   name this function and callers used historically, it is the median, not a distinct
156    ///   mean statistic)
157    pub fn forecast(&self, context: &[f32], prediction_length: usize) -> Result<(Vec<Vec<f32>>, Vec<f32>)> {
158        let cfg = &self.config;
159        let patch_size = cfg.patch_size;
160        let n_quantiles = cfg.num_quantiles;
161        let median_idx = cfg.quantiles.iter().position(|&q| (q - 0.5).abs() < 1e-6).unwrap_or(4);
162
163        // Number of AR steps needed
164        let n_steps = prediction_length.div_ceil(patch_size);
165
166        // Accumulate: [n_steps * patch_size, n_quantiles]
167        let mut all_q: Vec<f32> = Vec::with_capacity(n_steps * patch_size * n_quantiles);
168
169        // Working context (may grow with NaN placeholders between steps)
170        let mut ctx: Vec<f32> = context.to_vec();
171
172        for _ in 0..n_steps {
173            // Pad/truncate to train_ctx_len
174            let full_ctx = adjust_context(&ctx, cfg.train_ctx_len);
175
176            // StandardScaler
177            let (loc, scale) = standard_scaler(&full_ctx);
178
179            // Scale values, zero-out NaN; build mask
180            let s_vals: Vec<f32> = full_ctx.iter().map(|&x| {
181                if x.is_nan() { 0.0f32 } else { (x - loc) / scale }
182            }).collect();
183            let s_mask: Vec<f32> = full_ctx.iter().map(|&x| {
184                if x.is_nan() { 0.0f32 } else { 1.0f32 }
185            }).collect();
186
187            // Patch: [num_patches, patch_size]
188            let num_patches = cfg.train_ctx_len / patch_size;
189            let mut patched_vals = vec![0.0f32; num_patches * patch_size];
190            let mut patched_mask = vec![0.0f32; num_patches * patch_size];
191            for p in 0..num_patches {
192                let start = p * patch_size;
193                patched_vals[p * patch_size..(p + 1) * patch_size]
194                    .copy_from_slice(&s_vals[start..start + patch_size]);
195                patched_mask[p * patch_size..(p + 1) * patch_size]
196                    .copy_from_slice(&s_mask[start..start + patch_size]);
197            }
198
199            // Concatenate [vals | mask] → [num_patches, 2*patch_size]
200            let in_dim = patch_size * 2;
201            let mut x_in = vec![0.0f32; num_patches * in_dim];
202            for p in 0..num_patches {
203                x_in[p * in_dim..p * in_dim + patch_size]
204                    .copy_from_slice(&patched_vals[p * patch_size..(p + 1) * patch_size]);
205                x_in[p * in_dim + patch_size..p * in_dim + in_dim]
206                    .copy_from_slice(&patched_mask[p * patch_size..(p + 1) * patch_size]);
207            }
208
209            // ResidualBlock (input patch embedding): [num_patches, in_dim] → [num_patches, D]
210            let mut hidden = residual_block_forward(
211                &x_in, num_patches, in_dim, cfg.input_ff_dim, cfg.embedding_dim,
212                &self.in_emb, &self.device,
213            )?;
214
215            // Pre-allocate scratch buffers shared across all 12 sLSTM blocks.
216            // Each block previously re-allocated these on every call; pre-allocating once
217            // eliminates 12 × ~(1.5 MB) of heap churn per forecast() call.
218            let ng = 4usize;
219            let d  = cfg.embedding_dim;
220            let mut sc_xg    = vec![0.0f32; num_patches * ng * d]; // fused gate output / x_g
221            let mut sc_hout  = vec![0.0f32; num_patches * d];       // h_out
222            let mut sc_y     = vec![0.0f32; num_patches * d];       // group-norm output y
223            let mut sc_xn    = vec![0.0f32; num_patches * d];       // pre-norm copy x_n
224            let mut sc_raw   = vec![0.0f32; ng * d];                // raw = wx + ry + bias
225            let mut sc_ry_raw = vec![0.0f32; ng * d];               // compute_ry intermediate
226            let mut sc_ry_out = vec![0.0f32; ng * d];               // compute_ry output
227            let mut sc_hnew  = vec![0.0f32; d];                     // h_new (swapped, not re-alloc)
228            let mut sc_cnew  = vec![0.0f32; d];
229            let mut sc_nnew  = vec![0.0f32; d];
230            let mut sc_mnew  = vec![0.0f32; d];
231
232            // 12 sLSTM blocks
233            for block in &self.blocks {
234                hidden = self.forward_slstm_block(
235                    &hidden, num_patches, block,
236                    &mut sc_xg, &mut sc_hout, &mut sc_y, &mut sc_xn,
237                    &mut sc_raw, &mut sc_ry_raw, &mut sc_ry_out,
238                    &mut sc_hnew, &mut sc_cnew, &mut sc_nnew, &mut sc_mnew,
239                )?;
240            }
241
242            // out_norm (RMSNorm)
243            rms_norm_inplace(&mut hidden, &self.out_norm, cfg.embedding_dim, 1e-6);
244
245            // output_patch_embedding: [num_patches, D] → [num_patches, Q*P]
246            let out_dim = cfg.output_dim();
247            let preds = residual_block_forward(
248                &hidden, num_patches, cfg.embedding_dim, cfg.input_ff_dim, out_dim,
249                &self.out_emb, &self.device,
250            )?;
251            // preds: [num_patches, Q*patch_size] = [num_patches, 9*32]
252            // Take last patch: [Q*patch_size = 288]
253            let last = &preds[(num_patches - 1) * out_dim..num_patches * out_dim];
254
255            // Rescale and collect per-quantile for this step
256            for q in 0..n_quantiles {
257                for d in 0..patch_size {
258                    // preds layout: [Q*P] = [q0_p0, q0_p1, ..., q0_p31, q1_p0, ...]
259                    // Wait — the Python output is [Q, patch_size] after unflatten(-1, (Q, P))
260                    // and output_patch_embedding gives [S, Q*P] where Q varies slowest
261                    // i.e., preds[s, q*P + d] = prediction for quantile q, offset d in patch
262                    let v = last[q * patch_size + d] * scale + loc;
263                    all_q.push(v);
264                }
265            }
266
267            // Extend context with NaN patch for next AR step
268            ctx.extend(std::iter::repeat(f32::NAN).take(patch_size));
269        }
270
271        // all_q shape: [n_steps, n_quantiles, patch_size] when indexed as
272        //   all_q[step * n_quantiles * patch_size + q * patch_size + d]
273        // We want [prediction_length, n_quantiles]
274        let total = n_steps * patch_size;
275        let mut quantiles: Vec<Vec<f32>> = vec![Vec::with_capacity(prediction_length); n_quantiles];
276        for t in 0..prediction_length.min(total) {
277            let step = t / patch_size;
278            let d = t % patch_size;
279            for q in 0..n_quantiles {
280                let v = all_q[step * n_quantiles * patch_size + q * patch_size + d];
281                quantiles[q].push(v);
282            }
283        }
284
285        let median: Vec<f32> = (0..prediction_length).map(|t| quantiles[median_idx][t]).collect();
286
287        Ok((quantiles, median))
288    }
289
290    #[allow(clippy::too_many_arguments)]
291    fn forward_slstm_block(
292        &self, x: &[f32], s: usize, blk: &Block,
293        sc_xg:    &mut [f32],  // [s * ng * d]  — fused gate projection output
294        sc_hout:  &mut [f32],  // [s * d]        — sLSTM hidden outputs
295        sc_y:     &mut [f32],  // [s * d]        — group-norm output
296        sc_xn:    &mut [f32],  // [s * d]        — pre-norm scratch
297        sc_raw:   &mut [f32],  // [ng * d]       — wx + ry + bias per step
298        sc_ry_raw: &mut [f32], // [ng * d]       — compute_ry intermediate
299        sc_ry_out: &mut [f32], // [ng * d]       — compute_ry output
300        sc_hnew: &mut Vec<f32>, sc_cnew: &mut Vec<f32>,
301        sc_nnew: &mut Vec<f32>, sc_mnew: &mut Vec<f32>,
302    ) -> Result<Vec<f32>> {
303        let cfg = &self.config;
304        let d  = cfg.embedding_dim;
305        let nh = cfg.num_heads;
306        let dh = cfg.head_dim();
307        let ng = 4usize;
308
309        // Pre-norm: reuse sc_xn scratch to avoid allocation
310        sc_xn[..s * d].copy_from_slice(&x[..s * d]);
311        rms_norm_inplace(&mut sc_xn[..s * d], &blk.norm_slstm, d, 1e-6);
312
313        // Single fused gate projection: [S, D] → [S, ng*D] in one pass over x_n.
314        // Replaces 4 separate headwise_linear_batch calls + interleaving copy loop.
315        headwise_linear_batch_ng(&sc_xn[..s * d], &blk.fizo_w, s, nh, dh, ng, sc_xg);
316
317        // Sequential sLSTM recurrence — state lives on the stack, swapped not re-allocated
318        let mut h = vec![0.0f32; d];
319        let mut c = vec![0.0f32; d];
320        let mut n = vec![0.0f32; d];
321        let mut m = vec![f32::NEG_INFINITY; d];
322
323        for t in 0..s {
324            let wx = &sc_xg[t * ng * d..(t + 1) * ng * d];
325            compute_ry(&h, &blk.slstm_kernel_t, nh, dh, ng, sc_ry_raw, sc_ry_out);
326
327            // Reuse sc_raw to avoid per-step allocation of raw = wx + ry + bias
328            for i in 0..ng * d {
329                sc_raw[i] = wx[i] + sc_ry_out[i] + blk.slstm_bias[i];
330            }
331
332            // is_first ≡ t == 0: n starts as zeros and is never zero again after step 0
333            let is_first = t == 0;
334
335            // gate indices in flat [NG, NH, DH] layout:
336            //   g=0 (offset 0) → iraw (from fgate module, input gate)
337            //   g=1 (offset D) → fraw (from igate module, forget gate)
338            //   g=2 (offset 2D) → zraw (cell gate)
339            //   g=3 (offset 3D) → oraw (output gate)
340            simd_gate_update(
341                &sc_raw[..d], &sc_raw[d..2*d], &sc_raw[2*d..3*d], &sc_raw[3*d..],
342                &c, &n, &m,
343                sc_cnew, sc_nnew, sc_hnew, sc_mnew,
344                is_first,
345            );
346
347            // Swap state vectors — zero allocation, just pointer swap
348            std::mem::swap(&mut h, sc_hnew);
349            std::mem::swap(&mut c, sc_cnew);
350            std::mem::swap(&mut n, sc_nnew);
351            std::mem::swap(&mut m, sc_mnew);
352            sc_hout[t * d..(t + 1) * d].copy_from_slice(&h);
353        }
354
355        // MultiHeadLayerNorm (group norm per head per token) + reshape to [S, D]
356        for t in 0..s {
357            let h_t = &sc_hout[t * d..(t + 1) * d];
358            for head in 0..nh {
359                let start = head * dh;
360                let h_slice = &h_t[start..start + dh];
361                let mean = h_slice.iter().sum::<f32>() / dh as f32;
362                let var = h_slice.iter().map(|&v| (v - mean) * (v - mean)).sum::<f32>() / dh as f32;
363                let inv_std = 1.0 / (var + 1e-5f32).sqrt();
364                let w_slice = &blk.group_norm_w[start..start + dh];
365                for d_i in 0..dh {
366                    sc_y[t * d + start + d_i] =
367                        (h_t[start + d_i] - mean) * inv_std * (1.0 + w_slice[d_i]);
368                }
369            }
370        }
371
372        // Residual: x + y
373        let mut x_out = x.to_vec();
374        for i in 0..s * d {
375            x_out[i] += sc_y[i];
376        }
377
378        // FFN pre-norm (reuse sc_xn)
379        sc_xn[..s * d].copy_from_slice(&x_out[..s * d]);
380        rms_norm_inplace(&mut sc_xn[..s * d], &blk.norm_ffn, d, 1e-6);
381
382        // FFN: SiLU gated: (silu(gate(x)) * up(x)) → down
383        let ffn_out = ffn_forward(&sc_xn[..s * d], s, d, cfg.ffn_up_dim, blk, &self.device)?;
384
385        // Residual: x_out + ffn
386        for i in 0..s * d {
387            x_out[i] += ffn_out[i];
388        }
389
390        Ok(x_out)
391    }
392}
393
394// ---------------------------------------------------------------------------
395// Math helpers
396// ---------------------------------------------------------------------------
397
398#[inline]
399fn sigmoid(x: f32) -> f32 {
400    1.0 / (1.0 + (-x).exp())
401}
402
403#[inline]
404fn log_sigmoid(x: f32) -> f32 {
405    // log(sigmoid(x)) = -log(1 + exp(-x)) for stability
406    if x >= 0.0 {
407        -(1.0 + (-x).exp()).ln()
408    } else {
409        x - (1.0 + x.exp()).ln()
410    }
411}
412
413// ---------------------------------------------------------------------------
414// SIMD kernels (cross-platform: SSE2 / AVX2 / AVX-512 / NEON / WASM / scalar)
415// ---------------------------------------------------------------------------
416
417simd_runtime_generate!(
418    fn simd_sq_sum(row: &[f32]) -> f32 {
419        let mut r = &row[..];
420        let mut acc = S::Vf32::zeroes();
421        while r.len() >= S::Vf32::WIDTH {
422            let v = S::Vf32::load_from_slice(r);
423            acc = v.mul_add(v, acc);
424            r = &r[S::Vf32::WIDTH..];
425        }
426        let mut sum = acc.horizontal_add();
427        for &x in r { sum += x * x; }
428        sum
429    }
430);
431
432simd_runtime_generate!(
433    fn simd_scale_weight(row: &mut [f32], w: &[f32], scale: f32) {
434        let sc = S::Vf32::set1(scale);
435        let mut r = &mut row[..];
436        let mut wv = &w[..];
437        while r.len() >= S::Vf32::WIDTH {
438            let v = S::Vf32::load_from_slice(r);
439            let wi = S::Vf32::load_from_slice(wv);
440            (v * sc * wi).copy_to_slice(r);
441            r = &mut r[S::Vf32::WIDTH..];
442            wv = &wv[S::Vf32::WIDTH..];
443        }
444        for i in 0..r.len() { r[i] *= scale * wv[i]; }
445    }
446);
447
448simd_runtime_generate!(
449    fn simd_dot(a: &[f32], b: &[f32]) -> f32 {
450        let mut aa = &a[..];
451        let mut bb = &b[..];
452        let mut acc = S::Vf32::zeroes();
453        while aa.len() >= S::Vf32::WIDTH {
454            let va = S::Vf32::load_from_slice(aa);
455            let vb = S::Vf32::load_from_slice(bb);
456            acc = va.mul_add(vb, acc);
457            aa = &aa[S::Vf32::WIDTH..];
458            bb = &bb[S::Vf32::WIDTH..];
459        }
460        let mut sum = acc.horizontal_add();
461        for (&x, &y) in aa.iter().zip(bb.iter()) { sum += x * y; }
462        sum
463    }
464);
465
466simd_runtime_generate!(
467    fn simd_gate_update(
468        iraw: &[f32], fraw: &[f32], zraw: &[f32], oraw: &[f32],
469        c: &[f32], n: &[f32], m: &[f32],
470        cnew: &mut [f32], nnew: &mut [f32], hnew: &mut [f32], mnew: &mut [f32],
471        is_first: bool,
472    ) {
473        let clamp = S::Vf32::set1(15.0f32);
474        let one   = S::Vf32::set1(1.0f32);
475        let eps   = S::Vf32::set1(1e-8f32);
476        let zero  = S::Vf32::zeroes();
477
478        let mut ir = &iraw[..]; let mut fr = &fraw[..];
479        let mut zr = &zraw[..]; let mut or_ = &oraw[..];
480        let mut cv = &c[..];    let mut nv = &n[..];    let mut mv = &m[..];
481        let mut cnw = &mut cnew[..]; let mut nnw = &mut nnew[..];
482        let mut hnw = &mut hnew[..]; let mut mnw = &mut mnew[..];
483
484        while ir.len() >= S::Vf32::WIDTH {
485            let iv  = S::Vf32::load_from_slice(ir);
486            let fv  = S::Vf32::load_from_slice(fr).min(clamp);
487            let zv  = S::Vf32::load_from_slice(zr);
488            let ov  = S::Vf32::load_from_slice(or_);
489            let cv_ = S::Vf32::load_from_slice(cv);
490            let nv_ = S::Vf32::load_from_slice(nv);
491            let mv_ = S::Vf32::load_from_slice(mv);
492
493            // log_sigmoid(fv): stable two-branch form, select by sign
494            let ls_pos = -(one + (-fv).exp_u35()).ln_u35();     // fv >= 0: -ln(1+exp(-fv))
495            let ls_neg = fv - (one + fv.exp_u35()).ln_u35();     // fv < 0:  fv - ln(1+exp(fv))
496            let log_sig_f = fv.cmp_lt(zero).blendv(ls_pos, ls_neg);
497
498            let logfplusm = mv_ + log_sig_f;
499            let mnew_v = if is_first { iv } else { iv.max(logfplusm) };
500
501            let ogate = one / (one + (-ov).exp_u35());            // sigmoid(ov)
502            let igate = (iv - mnew_v).exp_u35().min(one);
503            let fgate = (logfplusm - mnew_v).exp_u35().min(one);
504            let zgate = zv.tanh_u35();
505
506            let cnew_v = fgate.mul_add(cv_, igate * zgate);
507            let nnew_v = fgate.mul_add(nv_, igate);
508
509            let mask  = nnew_v.abs().cmp_gt(eps);
510            let hnew_v = mask.blendv(zero, ogate * cnew_v / nnew_v);
511
512            cnew_v.copy_to_slice(cnw); nnew_v.copy_to_slice(nnw);
513            hnew_v.copy_to_slice(hnw); mnew_v.copy_to_slice(mnw);
514
515            ir  = &ir[S::Vf32::WIDTH..];  fr  = &fr[S::Vf32::WIDTH..];
516            zr  = &zr[S::Vf32::WIDTH..];  or_ = &or_[S::Vf32::WIDTH..];
517            cv  = &cv[S::Vf32::WIDTH..];  nv  = &nv[S::Vf32::WIDTH..];
518            mv  = &mv[S::Vf32::WIDTH..];
519            cnw = &mut cnw[S::Vf32::WIDTH..]; nnw = &mut nnw[S::Vf32::WIDTH..];
520            hnw = &mut hnw[S::Vf32::WIDTH..]; mnw = &mut mnw[S::Vf32::WIDTH..];
521        }
522
523        for i in 0..ir.len() {
524            let iv_s = ir[i]; let fv_s = fr[i].min(15.0); let zv_s = zr[i]; let ov_s = or_[i];
525            let c_s = cv[i]; let n_s = nv[i]; let m_s = mv[i];
526            let ls = if fv_s >= 0.0 { -(1.0 + (-fv_s).exp()).ln() } else { fv_s - (1.0 + fv_s.exp()).ln() };
527            let lfpm = m_s + ls;
528            let mn = if is_first { iv_s } else { iv_s.max(lfpm) };
529            let og = 1.0 / (1.0 + (-ov_s).exp());
530            let ig = (iv_s - mn).exp().min(1.0);
531            let fg = (lfpm - mn).exp().min(1.0);
532            let zg = zv_s.tanh();
533            cnw[i] = fg * c_s + ig * zg;
534            nnw[i] = fg * n_s + ig;
535            hnw[i] = if nnw[i].abs() > 1e-8 { og * cnw[i] / nnw[i] } else { 0.0 };
536            mnw[i] = mn;
537        }
538    }
539);
540
541/// StandardScaler: compute (loc, scale) ignoring NaN values.
542fn standard_scaler(x: &[f32]) -> (f32, f32) {
543    let eps = 1e-5f32;
544    let valid: Vec<f32> = x.iter().copied().filter(|v| !v.is_nan()).collect();
545    if valid.is_empty() {
546        return (0.0, 1.0);
547    }
548    let loc = valid.iter().sum::<f32>() / valid.len() as f32;
549    let variance = valid.iter().map(|&v| (v - loc) * (v - loc)).sum::<f32>() / valid.len() as f32;
550    let scale = variance.sqrt();
551    let scale = if scale == 0.0 { loc.abs() + eps } else { scale };
552    (loc, scale)
553}
554
555/// Pad (left with NaN) or truncate (from front) context to target_len.
556fn adjust_context(x: &[f32], target_len: usize) -> Vec<f32> {
557    if x.len() >= target_len {
558        x[x.len() - target_len..].to_vec()
559    } else {
560        let pad = target_len - x.len();
561        let mut out = vec![f32::NAN; target_len];
562        out[pad..].copy_from_slice(x);
563        out
564    }
565}
566
567/// RMSNorm in-place: x = x / rms(x) * w
568fn rms_norm_inplace(x: &mut [f32], w: &[f32], d: usize, eps: f32) {
569    let s = x.len() / d;
570    for t in 0..s {
571        let row = &mut x[t * d..(t + 1) * d];
572        let ss = simd_sq_sum(row);
573        let scale = 1.0 / (ss / d as f32 + eps).sqrt();
574        simd_scale_weight(row, w, scale);
575    }
576}
577
578/// Fused headwise linear for ng gate groups: [S, D] → [S, ng*D] written into `out`.
579///
580/// w layout: [NH, ng*DH, DH] — for head h, gate g: rows g*DH..(g+1)*DH hold w[h,g,:].
581/// Output layout: out[t, g, h, o] at flat index t*ng*d + g*d + h*dh + o.
582/// This matches the x_g layout expected by the sLSTM recurrence, so no interleaving
583/// copy is needed after calling this function.
584fn headwise_linear_batch_ng(x: &[f32], w: &[f32], s: usize, nh: usize, dh: usize, ng: usize, out: &mut [f32]) {
585    let d    = nh * dh;
586    let ngdh = ng * dh;
587    for t in 0..s {
588        for h in 0..nh {
589            let x_h     = &x[t * d + h * dh..t * d + (h + 1) * dh];
590            let w_h_off = h * ngdh * dh;
591            for g in 0..ng {
592                let out_base    = t * ng * d + g * d + h * dh;
593                let w_gate_off  = w_h_off + g * dh * dh;
594                for o in 0..dh {
595                    let w_row = &w[w_gate_off + o * dh..w_gate_off + (o + 1) * dh];
596                    out[out_base + o] = simd_dot(x_h, w_row);
597                }
598            }
599        }
600    }
601}
602
603/// Compute recurrent contribution Ry from h_prev and the sLSTM kernel.
604///
605/// kernel_t layout: [NH, NG*DH, DH] (transposed at load time from [NH, DH, NG*DH])
606/// so the inner dot-product dimension (di) is contiguous, enabling simd_dot.
607/// Output written into `out`: [NG*NH*DH = 2048] in [NG, NH, DH] order.
608fn compute_ry(
609    h: &[f32], kernel_t: &[f32],
610    nh: usize, dh: usize, ng: usize,
611    ry_raw: &mut [f32],
612    out: &mut [f32],
613) {
614    // ry_raw[NH, NG*DH]: for each head, dot h_head with each row of kernel_t[head]
615    for head in 0..nh {
616        let h_h    = &h[head * dh..(head + 1) * dh];
617        let k_head = &kernel_t[head * ng * dh * dh..];
618        for gate_d in 0..ng * dh {
619            let k_row = &k_head[gate_d * dh..(gate_d + 1) * dh]; // contiguous — simd_dot applies
620            ry_raw[head * ng * dh + gate_d] = simd_dot(h_h, k_row);
621        }
622    }
623    // Permute [NH, NG, DH] → [NG, NH, DH]
624    for head in 0..nh {
625        for g in 0..ng {
626            for d in 0..dh {
627                out[g * nh * dh + head * dh + d] = ry_raw[head * ng * dh + g * dh + d];
628            }
629        }
630    }
631}
632
633/// ResidualBlock forward: x → relu(x @ Wh.T + bh) @ Wo.T + bo + x @ Wr.T + br
634fn residual_block_forward(
635    x: &[f32],
636    s: usize,
637    in_dim: usize,
638    h_dim: usize,
639    out_dim: usize,
640    emb: &EmbedBlock,
641    device: &Device,
642) -> Result<Vec<f32>> {
643    // x: [S, in_dim]
644    let xt = Tensor::from_slice(x, (s, in_dim), device)?;
645
646    // hidden: relu(x @ Wh.T + bh)
647    let h = xt.matmul(&emb.hidden_w.t()?)
648        .with_context(|| "in_emb hidden matmul")?;
649    // Add bias
650    let bh = Tensor::from_slice(&emb.hidden_b, (1, h_dim), device)?;
651    let h = h.broadcast_add(&bh)?;
652    let h = h.relu()?;
653
654    // output: h @ Wo.T + bo
655    let out = h.matmul(&emb.output_w.t()?)
656        .with_context(|| "in_emb output matmul")?;
657    let bo = Tensor::from_slice(&emb.output_b, (1, out_dim), device)?;
658    let out = out.broadcast_add(&bo)?;
659
660    // residual: x @ Wr.T + br
661    let res = xt.matmul(&emb.residual_w.t()?)
662        .with_context(|| "in_emb residual matmul")?;
663    let br = Tensor::from_slice(&emb.residual_b, (1, out_dim), device)?;
664    let res = res.broadcast_add(&br)?;
665
666    Ok((out + res)?.flatten_all()?.to_vec1()?)
667}
668
669/// FFN: SiLU(gate(x)) * up(x) → down
670fn ffn_forward(x: &[f32], s: usize, in_dim: usize, up_dim: usize, blk: &Block, device: &Device) -> Result<Vec<f32>> {
671    let xt = Tensor::from_slice(x, (s, in_dim), device)?;
672
673    // gate proj: [S, up_dim]
674    let gate = xt.matmul(&blk.ffn_gate_w.t()?)?;
675    // up proj: [S, up_dim]
676    let up = xt.matmul(&blk.ffn_up_w.t()?)?;
677    // SiLU gated: silu(gate) * up
678    let h = gate.silu()?.mul(&up)?;
679    // down proj: [S, in_dim]
680    let out = h.matmul(&blk.ffn_down_w.t()?)?;
681
682    Ok(out.flatten_all()?.to_vec1()?)
683}
684
685// ---------------------------------------------------------------------------
686// zsfm-core::Forecaster
687// ---------------------------------------------------------------------------
688
689impl zsfm_core::Forecaster for TiRexModel {
690    type Config = TiRexConfig;
691
692    fn load(gguf_path: &Path, config: TiRexConfig) -> Result<Self> {
693        TiRexModel::load(gguf_path, config)
694    }
695
696    /// TiRex is univariate-only; `mask` is unused (missing values aren't currently
697    /// threaded through from the shared JSON envelope).
698    fn forecast(
699        &self,
700        context: &[Vec<f32>],
701        _mask: &[Vec<bool>],
702        horizon: usize,
703    ) -> Result<zsfm_core::QuantileMatrix> {
704        anyhow::ensure!(context.len() == 1, "TiRexModel only supports univariate forecasting (1 variate)");
705        let (quantiles, _median) = TiRexModel::forecast(self, &context[0], horizon)?; // [n_q][horizon]
706        Ok(quantiles.into_iter().map(|row| vec![row]).collect())
707    }
708}
zsfm_tirex/infer/mod.rs

zsfm_tirex/infer/
mod.rs
