Timestral-209M on LiteRT: float16 graphs for the GPU and the NPU, int8 graphs for the CPU, host code in Python and Kotlin
11f5cc3 verified | """Host side of the Timestral LiteRT port: the official forecast() around one graph call. | |
| NumPy only (no torch), univariate series, the official defaults: symmetric averaging, | |
| non-negative clipping, quantile cap k = 1e4, blocks of up to 32 patches with the advancing | |
| refine frame, max_context 32768. Each stage keeps the official dtype: the causal statistics | |
| and the refine frames are float64; normalization, denormalization, the cap and the averaging | |
| are float32. asinh / sinh are evaluated in float64 and rounded to float32 (the official | |
| float32 kernels agree to about 1 ulp); this is the reference for the Kotlin port. | |
| A runner turns one decode block into normalized quantiles: | |
| runner(scaled [R, T] f32, observed [R, T] bool, horizon_patches h) -> [R, h*32, 21] f32 | |
| where scaled / observed are the normalized context of R rows that share `observed`. | |
| GraphInputs builds the five graph inputs of timestral_graph.TimestralGraph for a bucket. | |
| """ | |
| import math | |
| import os | |
| import numpy as np | |
| PATCH = 32 | |
| D_MODEL = 768 | |
| N_Q = 21 | |
| MEDIAN = 10 | |
| MIN_SCALE = 1e-6 | |
| MAX_CONTEXT = 32768 | |
| BLOCK = 32 | |
| CAP_K = 1e4 | |
| EXTREME = 1e30 | |
| # Additive mask for blocked pairs. Timestral-209M's attention logits reach ~1.6e5 and a blocked key can | |
| # sit 4e4 above the best allowed one (measured on the parity windows), so -1e4 does not zero it; -1e30 | |
| # does in fp32 (and becomes -inf in fp16, where exp still gives 0). | |
| NEG = -1.0e30 | |
| # Graphs exported with attn_k = k (file names with `_att`) add the mask to logits already divided by | |
| # 2^k, so the finite fp16 value -3e4 (= -4.8e5 logit units at k = 4, far below the 4.2e4 gap) blocks in | |
| # fp32 too, and it keeps -inf out of the row max: with -1e30 (-inf in fp16) Metal's fp16 kernels gave | |
| # NaN rows (desktop GPU, LiteRT 2.2.0). TS_NEG overrides the value for every graph. | |
| NEG_ATTN_K = -3.0e4 | |
| if os.environ.get("TS_NEG"): | |
| NEG = NEG_ATTN_K = float(os.environ["TS_NEG"]) | |
| def mask_value(graph_name): | |
| """The additive mask value the host uses for a graph file.""" | |
| return NEG_ATTN_K if "_att" in os.path.basename(str(graph_name)) else NEG | |
| F32 = np.float32 | |
| F64 = np.float64 | |
| F32_MAX = float(np.finfo(np.float32).max) | |
| F32_EPS = float(np.finfo(np.float32).eps) | |
| F64_TINY = float(np.finfo(np.float64).tiny) | |
| # ------------------------------------------------------------------ normalization | |
| def causal_statistics(x, w): | |
| """Official compute_causal_statistics: float64 running mean / Bessel std, floored.""" | |
| x = x.astype(F64) | |
| w = w.astype(F64) | |
| weighted = w * x | |
| cum_w = np.cumsum(w, axis=-1) | |
| cum_x = np.cumsum(weighted, axis=-1) | |
| den = np.maximum(cum_w, 1.0) | |
| means = cum_x / den | |
| shifted = np.zeros_like(means) | |
| shifted[..., 1:] = means[..., :-1] | |
| inc = (x - shifted) * (x - means) * w | |
| m2 = np.cumsum(inc, axis=-1) | |
| var = m2 / np.maximum(den - 1.0, 1.0) | |
| scales = np.maximum(np.sqrt(var), MIN_SCALE) | |
| return means.astype(F32), scales.astype(F32) | |
| def patch_loc_scale(x, observed): | |
| """Per-step (loc, scale) f32: each patch takes the statistics at its last step.""" | |
| means, scales = causal_statistics(x, observed) | |
| shp = x.shape[:-1] + (x.shape[-1] // PATCH, PATCH) | |
| loc = np.repeat(means.reshape(shp)[..., -1], PATCH, axis=-1) | |
| scale = np.repeat(scales.reshape(shp)[..., -1], PATCH, axis=-1) | |
| return loc, scale | |
| def normalize(x, loc, scale, observed): | |
| """asinh((x - loc) / max(scale, eps)) on observed steps (f32), 0 elsewhere.""" | |
| y = (x.astype(F32) - loc) / np.maximum(scale, F32(F32_EPS)) | |
| y = np.where(observed, y, F32(0.0)) | |
| return np.where(observed, np.arcsinh(y.astype(F64)).astype(F32), y) | |
| def clamp_nonfinite(v): | |
| """Official clamp_nonfinite along the last axis (NaN untouched).""" | |
| fin = np.isfinite(v) | |
| hi = np.where(fin, v, -np.inf).max(axis=-1, keepdims=True) | |
| lo = np.where(fin, v, np.inf).min(axis=-1, keepdims=True) | |
| out = np.where(v == np.inf, hi, np.where(v == -np.inf, lo, v)) | |
| return np.where(np.isinf(out), F32(0.0), out).astype(F32) | |
| def denormalize(preds, loc, scale): | |
| """clamp_nonfinite(sinh(preds) * scale + loc) in f32; loc / scale broadcast [R, 1, 1].""" | |
| with np.errstate(over="ignore", invalid="ignore"): | |
| s = np.sinh(preds.astype(F64)).astype(F32) | |
| return clamp_nonfinite(s * scale + loc) | |
| # ------------------------------------------------------------------ refine frame | |
| def welford_combine(w_a, mean_a, m2_a, w_b, mean_b, m2_b): | |
| w = w_a + w_b | |
| w_safe = np.maximum(w, F64_TINY) | |
| delta = mean_b - mean_a | |
| mean = mean_a + delta * w_b / w_safe | |
| m2 = m2_a + m2_b + np.square(delta) * w_a * w_b / w_safe | |
| zero = w == 0 | |
| return w, np.where(zero, 0.0, mean), np.where(zero, 0.0, m2) | |
| def cumulative_welford(w, mean, m2): | |
| n = w.shape[1] | |
| d = 1 | |
| while d < n: | |
| def shift(a): | |
| out = np.zeros_like(a) | |
| out[:, d:] = a[:, : n - d] | |
| return out | |
| w, mean, m2 = welford_combine(shift(w), shift(mean), shift(m2), w, mean, m2) | |
| d *= 2 | |
| return w, mean, m2 | |
| def advancing_block_frames(context, observed, sinh_median): | |
| """Official advancing_block_frames: per-patch float64 (loc, scale), each [R, H].""" | |
| r, n_pred, ps = sinh_median.shape | |
| blank = np.zeros((r, n_pred, ps), F64) | |
| values = np.concatenate([context.reshape(r, -1, ps).astype(F64), blank], axis=1) | |
| weights = np.concatenate([observed.reshape(r, -1, ps).astype(F64), blank], axis=1) | |
| cnt = weights.sum(-1) | |
| mean = (values * weights).sum(-1) / np.maximum(cnt, 1.0) | |
| m2 = (weights * np.square(values - mean[..., None])).sum(-1) | |
| zero = cnt == 0 | |
| mean = np.where(zero, 0.0, mean) | |
| m2 = np.where(zero, 0.0, m2) | |
| cnt, mean, m2 = (s[:, -n_pred - 1] for s in cumulative_welford(cnt, mean, m2)) | |
| offset = mean | |
| clip = F32_MAX ** 0.5 / 1e3 | |
| locs, scales = [], [] | |
| for h in range(n_pred): | |
| var = np.maximum(m2 / np.maximum(cnt - 1.0, 1.0), 0.0) | |
| scale = np.maximum(np.sqrt(var), MIN_SCALE) | |
| locs.append(mean) | |
| scales.append(scale) | |
| folded = np.clip(sinh_median[:, h] * scale[:, None] + mean[:, None] - offset[:, None], -clip, clip) | |
| folded = folded + offset[:, None] | |
| b_mean = folded.mean(-1) | |
| b_m2 = np.square(folded - b_mean[:, None]).sum(-1) | |
| cnt, mean, m2 = welford_combine(cnt, mean, m2, np.full_like(cnt, ps), b_mean, b_m2) | |
| return np.stack(locs, axis=1), np.stack(scales, axis=1) | |
| def refine_block(quantiles, context, observed): | |
| """Official _refine_block for target rows: anchor frame -> advancing refine frame.""" | |
| n_pred = quantiles.shape[1] // PATCH | |
| if n_pred <= 1: | |
| return quantiles | |
| loc, scale = patch_loc_scale(context, observed) | |
| loc0, scale0 = loc[:, -1], scale[:, -1] | |
| median = quantiles[..., MEDIAN].reshape(-1, n_pred, PATCH).astype(F64) | |
| sinh_median = (median - loc0[:, None, None].astype(F64)) / scale0[:, None, None].astype(F64) | |
| loc_j, scale_j = advancing_block_frames(context, observed, sinh_median) | |
| loc_j = loc_j.astype(F32) | |
| scale_j = scale_j.astype(F32) | |
| loc_j[:, 0] = loc0 | |
| scale_j[:, 0] = scale0 | |
| loc_j = np.repeat(loc_j, PATCH, axis=-1) | |
| scale_j = np.repeat(scale_j, PATCH, axis=-1) | |
| sinh_norm = (quantiles - loc0[:, None, None]) / scale0[:, None, None] | |
| return clamp_nonfinite(sinh_norm * scale_j[..., None] + loc_j[..., None]) | |
| # ------------------------------------------------------------------ graph inputs | |
| class GraphInputs: | |
| """Builds patches / dec / bias / cos / sin for a bucket (L, H) from a normalized block.""" | |
| def __init__(self, forecast_query, rope_cos, rope_sin, rope_p_min): | |
| self.fq = np.asarray(forecast_query, F32).reshape(-1) | |
| self.cos = np.asarray(rope_cos, F32) | |
| self.sin = np.asarray(rope_sin, F32) | |
| self.p_min = int(rope_p_min) | |
| def build(self, scaled, observed, bucket_l, bucket_h, neg=None): | |
| neg = NEG if neg is None else neg | |
| r, t = scaled.shape | |
| lc = t // PATCH | |
| assert t % PATCH == 0 and lc <= bucket_l, (t, bucket_l) | |
| pad = bucket_l - lc | |
| s = bucket_l + bucket_h | |
| patches = np.zeros((r, bucket_l, 2 * PATCH), F32) | |
| patches[:, :pad, PATCH:] = 1.0 | |
| patches[:, pad:, :PATCH] = scaled.reshape(r, lc, PATCH) | |
| patches[:, pad:, PATCH:] = (~observed).reshape(r, lc, PATCH).astype(F32) | |
| obs = np.zeros(bucket_l, bool) | |
| obs[pad:] = observed[0].reshape(lc, PATCH).any(-1) | |
| assert all((observed[i] == observed[0]).all() for i in range(r)), "rows must share the mask" | |
| key_valid = obs | np.maximum.accumulate(obs) | |
| first_obs = int(np.argmax(obs)) if obs.any() else 0 | |
| allow = np.zeros((s, s), bool) | |
| ll = bucket_l | |
| allow[:ll, :ll] = (np.tril(np.ones((ll, ll), bool)) & key_valid[None, :]) | np.eye(ll, dtype=bool) | |
| allow[ll:, :ll] = key_valid[None, :] | |
| allow[ll:, ll:] = np.tril(np.ones((bucket_h, bucket_h), bool)) | |
| bias = np.where(allow, F32(0.0), F32(neg)).reshape(1, 1, s, s) | |
| rows = np.arange(s) - first_obs - self.p_min | |
| assert rows.min() >= 0 and rows.max() < self.cos.shape[0], "RoPE table too small" | |
| cos = self.cos[rows].reshape(1, 1, s, -1) | |
| sin = self.sin[rows].reshape(1, 1, s, -1) | |
| dec = np.broadcast_to(self.fq, (r, bucket_h, self.fq.size)).astype(F32) | |
| return [patches, dec, bias, cos, sin] | |
| # ------------------------------------------------------------------ forecast | |
| def nonnegative_context(inputs, observed): | |
| obs = np.isfinite(inputs) & observed | |
| return bool(obs.any() and ((inputs >= 0) | ~obs).all()) | |
| def ffill(inputs, observed): | |
| if observed.all(): | |
| return inputs | |
| idx = np.where(observed, np.arange(inputs.shape[-1]), 0) | |
| idx = np.maximum.accumulate(idx, axis=-1) | |
| out = np.take_along_axis(inputs, idx, axis=-1) | |
| return np.nan_to_num(out, nan=0.0).astype(F32) | |
| def pad_to_patch(x, observed): | |
| t = x.shape[-1] | |
| if t % PATCH: | |
| pad = PATCH - t % PATCH | |
| x = np.concatenate([np.zeros(x.shape[:-1] + (pad,), F32), x], axis=-1) | |
| observed = np.concatenate([np.zeros(observed.shape[:-1] + (pad,), bool), observed], axis=-1) | |
| return x, observed | |
| def cap_quantiles(q, inputs, observed): | |
| with np.errstate(invalid="ignore"): | |
| ctx_max = np.where(observed, inputs, -np.inf).max(-1, keepdims=True) | |
| ctx_min = np.where(observed, inputs, np.inf).min(-1, keepdims=True) | |
| ctx_max = np.where(np.isneginf(ctx_max), 0.0, ctx_max).astype(F32) | |
| ctx_min = np.where(np.isposinf(ctx_min), 0.0, ctx_min).astype(F32) | |
| _, scale = patch_loc_scale(inputs, observed) | |
| anchor = scale[..., -1:] * F32(CAP_K) | |
| lo = (ctx_min - anchor)[..., None] | |
| hi = (ctx_max + anchor)[..., None] | |
| return np.minimum(np.maximum(q, lo), hi).astype(F32) | |
| def block_ar_decode(context, observed, prediction_length, runner): | |
| """Rows of one pass group share `observed`; returns [R, prediction_length, 21] f32.""" | |
| blocks = [] | |
| remaining = math.ceil(prediction_length / PATCH) | |
| while remaining > 0: | |
| h = min(BLOCK, remaining) | |
| loc, scale = patch_loc_scale(context, observed) | |
| scaled = normalize(context, loc, scale, observed) | |
| preds = np.asarray(runner(scaled, observed, h), F32) | |
| preds = np.clip(preds, -F32_MAX, F32_MAX) | |
| q = denormalize(preds, loc[:, -1:, None], scale[:, -1:, None]) | |
| q = refine_block(q, context, observed) | |
| blocks.append(q) | |
| remaining -= h | |
| if remaining <= 0: | |
| break | |
| med = q[..., MEDIAN] | |
| feedback = np.where(np.isfinite(med) & (np.abs(med) < EXTREME), med, F32(0.0)).astype(F32) | |
| context = np.concatenate([context, feedback], axis=-1) | |
| observed = np.concatenate([observed, np.ones_like(feedback, bool)], axis=-1) | |
| if context.shape[-1] > MAX_CONTEXT: | |
| keep = max(PATCH, (MAX_CONTEXT // PATCH) * PATCH) | |
| context, observed = context[:, -keep:], observed[:, -keep:] | |
| return np.concatenate(blocks, axis=1)[:, :prediction_length] | |
| def forecast(series, prediction_length, runner, mirror_in_one_call=True, symmetric=True, clip=True): | |
| """One univariate series (NaN = missing) -> quantiles [prediction_length, 21] f32. | |
| mirror_in_one_call: run the context and its negation as R = 2 rows of one call per | |
| block (the runner must accept R = 2); False runs two R = 1 passes. | |
| """ | |
| x = np.asarray(series, F32).reshape(1, -1) | |
| observed = np.isfinite(x) | |
| inputs = np.where(observed, x, F32(0.0)).astype(F32) | |
| rows = ffill(inputs, observed) | |
| obs = observed | |
| if rows.shape[-1] > MAX_CONTEXT: | |
| rows, obs = rows[:, -MAX_CONTEXT:], obs[:, -MAX_CONTEXT:] | |
| rows, obs = pad_to_patch(rows, obs) | |
| def finish(q, ctx): | |
| """_forecast_raw's tail: the cap uses the rows of that pass (the negated rows for the mirror).""" | |
| q = cap_quantiles(q, ctx, obs) | |
| if np.isnan(q).any(): | |
| raise ValueError(f"forecast produced {int(np.isnan(q).sum())} NaN values") | |
| return q | |
| if symmetric and mirror_in_one_call: | |
| both = np.concatenate([rows, -rows], axis=0) | |
| q2 = block_ar_decode(both, np.concatenate([obs, obs], axis=0), prediction_length, runner) | |
| q = finish(q2[:1], rows) | |
| qm = finish(q2[1:], -rows) | |
| else: | |
| q = finish(block_ar_decode(rows, obs, prediction_length, runner), rows) | |
| qm = finish(block_ar_decode(-rows, obs, prediction_length, runner), -rows) if symmetric else None | |
| if symmetric: | |
| q = (F32(0.5) * (q + (-qm)[..., ::-1])).astype(F32) | |
| if clip and nonnegative_context(inputs, observed): | |
| q = np.maximum(q, F32(0.0)) | |
| return q[0] | |