continuousintelligence 0.1.0__py3-none-any.whl

continuousintelligence/ci_proto/_groups.py

@@ -0,0 +1,198 @@
+"""Group context: what the other series of an evaluation window add to each series' forecast.
+
+All three functions take the univariate series of one window, each as its observed history
+(``histories``) and that history's epoch-ns timestamps (``stamps``), in the window's order. Every
+value they return comes from those histories, and a series only pools with series whose history
+ends at the same timestamp, its forecast origin, so no value after the origin is used.
+
+- :func:`static_group_rows` gives each series past covariate rows: the mean of its static
+  group's other z-scored histories, as Chronos-2's groups pool series by their metadata.
+- :func:`joint_neighbour_rows` gives each series extra target rows: the series whose z-scored
+  histories correlate best with its own, forecast jointly with it.
+- :func:`chunk_groups` splits the window into chunks that are forecast jointly, every series a
+  target, as t0-gamma's ``cross_inference`` and Chronos-2's ``cross_learning`` do on fev-bench.
+
+The z-scores are taken on a grid of the window's last ``span`` distinct timestamps. Each row
+these functions return covers its series' own history timestamps and is missing where no value
+is available.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Final
+
+import numpy as np
+from numpy.typing import NDArray
+
+# The timestamps that z-scores and correlations are taken over.
+SPAN: Final = 2048
+
+# Two histories correlate only over at least this many common steps.
+_MIN_OVERLAP: Final = 32
+
+# A window whose largest origin group holds more series is left without joint neighbours, whose
+# correlation matrix would grow with the square of the group.
+_MAX_JOINT_GROUP: Final = 8000
+
+
+def _codes(values: Sequence[object]) -> NDArray[np.int64]:
+    """Number each distinct value; a group's members share a number."""
+    return np.unique(np.asarray(values), return_inverse=True)[1].astype(np.int64)
+
+
+def _origins(stamps: Sequence[NDArray[np.int64]]) -> NDArray[np.int64]:
+    """Number each series by the last timestamp of its history, its forecast origin."""
+    return _codes([int(s[-1]) if s.size else -1 for s in stamps])
+
+
+def _z_scores(
+    histories: Sequence[NDArray[np.float64]], stamps: Sequence[NDArray[np.int64]], span: int
+) -> tuple[NDArray[np.int64], NDArray[np.float32]]:
+    """Return the timestamp grid and each history z-scored onto it.
+
+    A history without spread stays missing.
+    """
+    grid = np.unique(np.concatenate([s[-span:] for s in stamps]))[-span:]
+    z = np.full((len(histories), grid.size), np.nan, dtype=np.float32)
+    for i, (row, s) in enumerate(zip(histories, stamps, strict=True)):
+        with np.errstate(all="ignore"):
+            scale = np.nanstd(row)
+        if not np.isfinite(scale) or scale <= 0:
+            continue
+        at = np.searchsorted(grid, s)
+        ok = (at < grid.size) & (grid[np.minimum(at, grid.size - 1)] == s)
+        z[i, at[ok]] = ((row[ok] - np.nanmean(row)) / scale).astype(np.float32)
+    return grid, z
+
+
+def _on_own_stamps(
+    values: NDArray[np.float64], grid: NDArray[np.int64], stamps: NDArray[np.int64]
+) -> NDArray[np.float64]:
+    """Read a grid row at a series' own history timestamps; missing off the grid."""
+    if (
+        values.shape == stamps.shape
+        and np.array_equal(grid, stamps)
+        and np.all(grid[1:] > grid[:-1])
+    ):
+        # Duplicate timestamps must retain searchsorted's first-occurrence semantics.
+        return values.copy()
+    at = np.searchsorted(grid, stamps)
+    ok = (at < grid.size) & (grid[np.minimum(at, grid.size - 1)] == stamps)
+    row = np.full(stamps.size, np.nan)
+    row[ok] = values[at[ok]]
+    return row
+
+
+def static_group_rows(
+    histories: Sequence[NDArray[np.float64]],
+    stamps: Sequence[NDArray[np.int64]],
+    statics: Sequence[Sequence[object]],
+    columns: int,
+    span: int = SPAN,
+) -> list[NDArray[np.float64]] | None:
+    """Return each series' static-group rows, one per grouping column, or ``None`` if none groups.
+
+    ``statics`` holds each static column's value for every series. A column groups the window
+    when it has at least two distinct values, fewer than the window's series, and a median group
+    size of at least three. The ``columns`` of them with the most groups are kept, ties in column
+    order. Each gives every series one row: the mean of its same-origin group's other z-scored
+    histories. A missing value is a group of its own. A window of fewer than three series gets
+    no rows.
+    """
+    if columns < 1:
+        raise ValueError(f"static groups need at least one column, got {columns}")
+    count = len(histories)
+    if count < 3:
+        return None
+    groupings = []
+    for values in statics:
+        codes = _codes([repr(value) for value in values])
+        sizes = np.bincount(codes)
+        if 2 <= sizes.size < count and np.median(sizes[codes]) >= 3:
+            groupings.append((sizes.size, codes))
+    chosen = sorted(groupings, key=lambda grouping: -grouping[0])[:columns]
+    if not chosen:
+        return None
+    grid, z = _z_scores(histories, stamps, span)
+    seen = np.isfinite(z)
+    filled = np.where(seen, z, np.float32(0.0))
+    origin = _origins(stamps)
+    rows: list[list[NDArray[np.float64]]] = [[] for _ in histories]
+    for _, codes in chosen:
+        group = _codes((codes * (origin.max() + 1) + origin).tolist())
+        sums_by_group = np.zeros((group.max() + 1, grid.size), dtype=np.float32)
+        counts_by_group = np.zeros_like(sums_by_group)
+        np.add.at(sums_by_group, group, filled)
+        np.add.at(counts_by_group, group, seen.astype(np.float32))
+        sums = sums_by_group[group].astype(np.float64)
+        counts = counts_by_group[group].astype(np.float64) - seen
+        with np.errstate(all="ignore"):
+            means = (sums - filled) / counts
+        means[counts < 1] = np.nan
+        for i, s in enumerate(stamps):
+            rows[i].append(_on_own_stamps(means[i], grid, s))
+    return [np.stack(own) for own in rows]
+
+
+def joint_neighbour_rows(
+    histories: Sequence[NDArray[np.float64]],
+    stamps: Sequence[NDArray[np.int64]],
+    size: int,
+    span: int = SPAN,
+) -> list[NDArray[np.float64]] | None:
+    """Return each series' history followed by its ``size - 1`` neighbours', or ``None``.
+
+    A series' neighbours are the same-origin series whose z-scored histories correlate best with
+    its own over at least 32 common grid steps, each placed on the series' own history
+    timestamps. A series with fewer neighbours repeats its own history in the remaining rows. A
+    window of fewer than two series, or whose largest origin group holds more than 8,000, gets
+    none.
+    """
+    if size < 2:
+        raise ValueError(f"joint neighbours need a group of at least two series, got {size}")
+    if len(histories) < 2:
+        return None
+    _, z = _z_scores(histories, stamps, span)
+    seen = np.isfinite(z)
+    filled = np.where(seen, z, np.float32(0.0))
+    origin = _origins(stamps)
+    if np.bincount(origin).max() > _MAX_JOINT_GROUP:
+        return None
+    neighbours: list[list[int]] = [[] for _ in histories]
+    for group in range(origin.max() + 1):
+        members = np.flatnonzero(origin == group)
+        if members.size < 2:
+            continue
+        zf, sf = filled[members], seen[members].astype(np.float32)
+        overlap = sf @ sf.T
+        corr = (zf @ zf.T) / np.maximum(overlap, 1.0)
+        corr[overlap < _MIN_OVERLAP] = -np.inf
+        np.fill_diagonal(corr, -np.inf)
+        for r, member in enumerate(members.tolist()):
+            order = np.argsort(-corr[r], kind="stable")[: size - 1]
+            neighbours[member] = [int(members[o]) for o in order if np.isfinite(corr[r, o])]
+    out = []
+    for i, (own, s) in enumerate(zip(histories, stamps, strict=True)):
+        rows = [own] + [_on_own_stamps(histories[j], stamps[j], s) for j in neighbours[i]]
+        rows += [own] * (size - len(rows))
+        out.append(np.stack(rows))
+    return out
+
+
+def chunk_groups(origins: Sequence[int], size: int) -> list[list[int]]:
+    """Split a window's series into consecutive chunks of up to ``size`` with one origin.
+
+    ``origins`` holds each series' last history timestamp. Series keep the window's order within
+    a chunk, and a chunk may hold one series.
+    """
+    if size < 2:
+        raise ValueError(f"chunks need room for at least two series, got {size}")
+    by_origin: dict[int, list[int]] = {}
+    for index, origin in enumerate(origins):
+        by_origin.setdefault(origin, []).append(index)
+    return [
+        members[start : start + size]
+        for members in by_origin.values()
+        for start in range(0, len(members), size)
+    ]

continuousintelligence/ci_proto/_mamba.py

@@ -0,0 +1,190 @@
+"""PyTorch Mamba-3 recurrence for CPU and CUDA inference.
+
+Adapted by Continuous Intelligence from Mamba's module and
+SISO/MIMO reference calculations:
+https://github.com/state-spaces/mamba/tree/a14b1dff0454a3bc27d9eb31355dc01e4b2490ec
+
+RMS normalisation is adapted from xLSTM Large:
+https://github.com/NX-AI/xlstm/blob/ab22eadbd245f293dd8dec38ed29963d73758a12/xlstm/xlstm_large/components.py
+
+Copyright (c) 2026, Dao AI Lab, Goombalab.
+Copyright 2024 NXAI GmbH and its affiliates 2024
+Upstream portions are licensed under Apache-2.0; see NOTICE.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor, nn
+
+
+class RMSNorm(nn.Module):
+    """RMS norm with optional float32 reductions."""
+
+    def __init__(
+        self,
+        num_features: int,
+        eps: float = 1e-06,
+        force_float32_reductions: bool = True,
+        use_bias: bool = False,
+    ) -> None:
+        super().__init__()
+        self.eps = eps
+        self.force_float32_reductions = force_float32_reductions
+        self.weight = nn.Parameter(torch.ones(num_features))
+        self.bias = nn.Parameter(torch.zeros(num_features)) if use_bias else None
+
+    def forward(self, x: Tensor) -> Tensor:
+        in_dtype = x.dtype
+        if self.force_float32_reductions:
+            x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        x = x.to(in_dtype)
+        x = x * self.weight
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+
+def _rotate_pairs(x: Tensor, cos: Tensor, sin: Tensor, interleaved: bool) -> Tensor:
+    """Rotate pairs of the last axis; angles beyond ``cos`` leave their pairs unrotated."""
+    if interleaved:
+        pairs = x.view(*x.shape[:-1], -1, 2)
+        first, second = (pairs[..., 0], pairs[..., 1])
+    else:
+        half = x.shape[-1] // 2
+        first, second = (x[..., :half], x[..., half:])
+    if cos.shape[-1] < first.shape[-1]:
+        pad = first.shape[-1] - cos.shape[-1]
+        cos = F.pad(cos, (0, pad), value=1.0)
+        sin = F.pad(sin, (0, pad), value=0.0)
+    rotated_first = first * cos - second * sin
+    rotated_second = first * sin + second * cos
+    if interleaved:
+        return torch.stack([rotated_first, rotated_second], dim=-1).view_as(x)
+    return torch.cat([rotated_first, rotated_second], dim=-1)
+
+
+class Mamba3Reference(nn.Module):
+    """Mamba-3 block evaluated step by step in plain PyTorch."""
+
+    def __init__(
+        self,
+        d_model: int,
+        d_state: int = 128,
+        expand: int = 2,
+        headdim: int = 64,
+        is_mimo: bool = False,
+        mimo_rank: int = 4,
+    ) -> None:
+        super().__init__()
+        if type(d_state) is not int or d_state < 4 or d_state % 2:
+            raise ValueError("mamba_state_dim must be an even integer of at least 4")
+        self.d_state = d_state
+        self.is_mimo = is_mimo
+        self.mimo_rank = mimo_rank if is_mimo else 1
+        self.d_inner = expand * d_model
+        if self.d_inner % headdim != 0:
+            raise ValueError(f"expand * d_model {self.d_inner} not divisible by headdim {headdim}")
+        self.headdim = headdim
+        self.nheads = self.d_inner // headdim
+        split = d_state // 2 - d_state // 2 % 2
+        self.num_rope_angles = split // 2
+        bc = self.d_state * self.mimo_rank
+        d_in_proj = 2 * self.d_inner + 2 * bc + 3 * self.nheads + self.num_rope_angles
+        self.in_proj = nn.Linear(d_model, d_in_proj, bias=False)
+        dt = torch.exp(
+            torch.rand(self.nheads) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
+        ).clamp(min=1e-4)
+        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
+        self.B_bias = nn.Parameter(torch.ones(self.nheads, self.mimo_rank, d_state))
+        self.C_bias = nn.Parameter(torch.ones(self.nheads, self.mimo_rank, d_state))
+        self.B_norm = RMSNorm(d_state, eps=1e-05)
+        self.C_norm = RMSNorm(d_state, eps=1e-05)
+        if is_mimo:
+            shape = (self.nheads, self.mimo_rank, headdim)
+            self.mimo_x = nn.Parameter(torch.ones(shape) / self.mimo_rank)
+            self.mimo_z = nn.Parameter(torch.ones(shape))
+            self.mimo_o = nn.Parameter(torch.ones(shape) / self.mimo_rank)
+        self.D = nn.Parameter(torch.ones(self.nheads))
+        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
+
+    def forward(self, u: Tensor) -> Tensor:
+        batch, seq, _ = u.shape
+        bc = self.d_state * self.mimo_rank
+        z, x, B, C, dd_dt, dd_A, trap, angles = torch.split(
+            self.in_proj(u),
+            [
+                self.d_inner,
+                self.d_inner,
+                bc,
+                bc,
+                self.nheads,
+                self.nheads,
+                self.nheads,
+                self.num_rope_angles,
+            ],
+            dim=-1,
+        )
+        z = rearrange(z, "b l (h p) -> b l h p", p=self.headdim).to(torch.float32)
+        x = rearrange(x, "b l (h p) -> b l h p", p=self.headdim).to(torch.float32)
+        B = rearrange(B, "b l (r g n) -> b l r g n", r=self.mimo_rank, g=1)
+        C = rearrange(C, "b l (r g n) -> b l r g n", r=self.mimo_rank, g=1)
+        A = torch.clamp(-F.softplus(dd_A.to(torch.float32)), max=-1e-4)
+        DT = F.softplus(dd_dt.to(torch.float32) + self.dt_bias)
+        alpha = torch.exp(A * DT)
+        gate = torch.sigmoid(trap.to(torch.float32))
+        beta = (1.0 - gate) * DT * alpha
+        gamma = gate * DT
+        k = self.B_norm(B).expand(-1, -1, -1, self.nheads, -1).to(torch.float32)
+        q = self.C_norm(C).expand(-1, -1, -1, self.nheads, -1).to(torch.float32)
+        k = k + rearrange(self.B_bias, "h r n -> r h n")
+        q = q + rearrange(self.C_bias, "h r n -> r h n")
+        turns = torch.tanh(angles.to(torch.float32))[:, :, None, :] * DT[..., None] * math.pi
+        cumulative = torch.cumsum(turns, dim=1)
+        cos, sin = (torch.cos(cumulative)[:, :, None], torch.sin(cumulative)[:, :, None])
+        q = _rotate_pairs(q, cos, sin, interleaved=not self.is_mimo)
+        k = _rotate_pairs(k, cos, sin, interleaved=not self.is_mimo)
+        if self.is_mimo:
+            v = torch.einsum("blhp,hrp->blrhp", x, self.mimo_x.to(torch.float32))
+            gates = torch.einsum("blhp,hrp->blrhp", z, self.mimo_z.to(torch.float32))
+        else:
+            v, gates = (x[:, :, None], z[:, :, None])
+        state = u.new_zeros((batch, self.nheads, self.headdim, self.d_state), dtype=torch.float32)
+        previous = torch.zeros_like(state)
+        outputs = []
+        for t in range(seq):
+            current = torch.einsum("brhn,brhp->bhpn", k[:, t], v[:, t])
+            state = (
+                alpha[:, t, :, None, None] * state
+                + beta[:, t, :, None, None] * previous
+                + gamma[:, t, :, None, None] * current
+            )
+            previous = current
+            step = torch.einsum("bhpn,brhn->brhp", state, q[:, t])
+            step = step + self.D[None, None, :, None] * v[:, t]
+            step = step * F.silu(gates[:, t])
+            if self.is_mimo:
+                step = torch.einsum("brhp,hrp->bhp", step, self.mimo_o.to(torch.float32))
+            else:
+                step = step[:, 0]
+            outputs.append(step)
+        y = rearrange(torch.stack(outputs, dim=1), "b l h p -> b l (h p)")
+        return self.out_proj(y.to(u.dtype))
+
+
+class Mamba3(Mamba3Reference):
+    """Checkpoint-compatible cell with fused CUDA inference when configured by the loader."""
+
+    fused = False
+
+    def forward(self, u: Tensor) -> Tensor:
+        if self.fused and u.is_cuda and not torch.is_grad_enabled():
+            from ._cuda import mamba3_forward
+
+            return mamba3_forward(self, u)
+        return super().forward(u)