fluidlayer 2.0.0__py3-none-any.whl

fluidlayer/blocks.py

@@ -0,0 +1,275 @@
+"""
+fluidlayer.blocks — the FluidLayer block.
+
+One block, five stages.  Every stage exists to fix a named weakness of the
+naive "replace attention with diffusion" idea:
+
+  ROUTER      top-k content routing  -> semantic edges
+  TOPOLOGY    semantic edges + temporal chain  -> the complex
+  SHEAF       learned signed restriction maps -> content-adaptive geometry
+  FREEZE      multi-scale band filter bank (exp(-t_j L) - exp(-t_{j+1} L))
+  READ        channel mixing + ONE nonlinearity + residual + sharp route head
+
+Stage notes
+-----------
+* The signed restriction maps make the G-balance phase transition learnable:
+  a balanced component freezes onto a non-trivial sign pattern, an unbalanced
+  one decays to zero.  `tanh` keeps |alpha| <= 1 so L stays bounded and
+  gradients cannot explode.
+* The band filter bank (not a single exponential) is what prevents the
+  "frozen state has rank c" information loss.
+* The route head is a hard top-k gather (winner-take-all).  It is the sharp
+  retrieval path the smooth fluid cannot provide -- this is what keeps the
+  model from being a pure smoother.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .autodiff import Tensor, concat, softmax_rows, softplus
+from .diffusion import (
+    band_times,
+    cheb_coeffs_exp,
+    cheb_coeffs_exp_autodiff,
+    power_iteration_lmax,
+    spectral_bound_autodiff,
+)
+from .topology import (
+    chain_edges,
+    connected_components,
+    edges_from_topk,
+    topo_descriptor,
+)
+
+__all__ = ["FluidLayerBlock", "RMSNorm", "softplus"]
+
+
+class RMSNorm:
+    def __init__(self, d: int, seed: int = 0):
+        self.g = Tensor(np.ones(d))
+        self.eps = 1e-6
+
+    def __call__(self, x: Tensor) -> Tensor:
+        ms = (x * x).mean()
+        inv = (ms + self.eps).pow(-0.5)
+        n = x * inv
+        return n * self.g
+
+
+class FluidLayerBlock:
+    def __init__(self, d: int, k_neighbors: int = 8, bands: int = 6, cheb_m: int = 12,
+                 t_min: float = 1e-3, t_max: float = 1e3, edge_hidden: int = 16,
+                 use_route_head: bool = True, seed: int = 0,
+                 alpha_min: float = 0.05, alpha_max: float | None = None):
+        rng = np.random.default_rng(seed)
+        self.d = d
+        self.k = int(k_neighbors)
+        self.B = int(bands)
+        self.m = int(cheb_m)
+        # Bands live in SPECTRAL time units alpha = t*lam_max/2 so that the
+        # Chebyshev degree m stays sufficient for any graph scale.
+        self.alpha_min = alpha_min
+        self.alpha_max = alpha_max if alpha_max is not None else 0.75 * self.m
+        self.ts = np.geomspace(t_min, t_max, self.B + 1)   # fallback / reference
+        self.use_route_head = use_route_head
+
+        s = d ** -0.5
+        self.Wq = Tensor(rng.normal(0, s, (d, d)))
+        self.Wk = Tensor(rng.normal(0, s, (d, d)))
+        self.Wv = Tensor(rng.normal(0, s, (d, d)))
+        self.Wo = Tensor(rng.normal(0, s, (d, d)))
+        self.Wmix = Tensor(rng.normal(0, s, (d, d)))
+        self.Wr = Tensor(rng.normal(0, s, (d, d))) if use_route_head else None
+        self.edge1 = Tensor(rng.normal(0, (2 * d) ** -0.5, (2 * d, edge_hidden)))
+        self.edge2 = Tensor(rng.normal(0, edge_hidden ** -0.5, (edge_hidden, 1)))
+        self.beta_chain = Tensor(np.array([[0.3]]))       # order / temporal edges
+        # Band gates start NEAR ZERO (softplus(-3) ~ 0.05): the transient is
+        # much larger in magnitude than the harmonic component, so starting
+        # with the bands "open" would swamp the frozen state.
+        self.band_g = Tensor(np.full((self.B, 1), -3.0))    # band gates (softplus)
+        self.harm_g = Tensor(np.array([[1.0]]))            # gate on the HARMONIC branch
+        self.norm = RMSNorm(d)
+        self.diag = {}
+
+    # ------------------------------------------------------------- helpers
+    def _laplacian(self, n: int, edges: np.ndarray, alpha: Tensor) -> Tensor:
+        """Differentiable signed sheaf Laplacian  L = D_abs - A."""
+        m = edges.shape[0]
+        if m == 0:
+            return Tensor(np.zeros((n, n)))
+        idx_uv = (edges[:, 0].astype(np.int64) * n + edges[:, 1].astype(np.int64))
+        idx_vu = (edges[:, 1].astype(np.int64) * n + edges[:, 0].astype(np.int64))
+        Z = Tensor(np.zeros(n * n))
+        # two SEPARATE scatter bases, then added: index_add starts from zeros,
+        # so chaining it would silently discard the first contribution.
+        A = (Z.index_add(idx_uv, alpha) + Z.index_add(idx_vu, alpha)).reshape(n, n)
+        deg = A.abs().sum(axis=1)                          # (n,)
+        diag_idx = np.arange(n) * n + np.arange(n)
+        D = Tensor(np.zeros(n * n)).index_add(diag_idx, deg).reshape(n, n)
+        return D - A
+
+    def _harmonic(self, x: Tensor, edges: np.ndarray):
+        """Exact Hodge projection of the graph Laplacian: the PER-COMPONENT MEAN.
+
+        ker(L) for a graph Laplacian is the space of per-component constants,
+        so the frozen state of the diffusion is the component mean.  Cost
+        O(n + m): union-find for the components, then a scatter-add and a
+        gather.  This is the branch that lets the block represent a GLOBAL
+        aggregate -- the band bank alone only carries the transient.
+        """
+        n, d = x.shape[0], x.shape[1]
+        labels, ncomp = connected_components(n, edges)
+        if ncomp <= 0:
+            return Tensor(np.zeros((n, d)))
+        sums = Tensor(np.zeros((ncomp, d))).index_add(labels, x)
+        ones = Tensor(np.ones(n))
+        counts = Tensor(np.zeros(ncomp)).index_add(labels, ones).reshape(ncomp, 1)
+        comp_mean = sums / counts
+        return comp_mean.gather(labels)
+
+    def _bands(self, L: Tensor, lam_max: float, X: Tensor, lam_max_t=None):
+        """Chebyshev band filter bank; the basis T_k(Ltilde)X is computed ONCE."""
+        n = L.shape[0]
+        inv = (2.0 / lam_max) if lam_max_t is None else (2.0 / lam_max_t)
+        Y = [X]
+        if self.m >= 1:
+            Lt = inv * L - Tensor(np.eye(n))
+            Y.append(Lt @ X)
+            for k in range(2, self.m + 1):
+                Y.append(2.0 * (Lt @ Y[k - 1]) - Y[k - 2])
+        ts = band_times(lam_max, self.m, self.B, self.alpha_min, self.alpha_max)
+        out = []
+        for j in range(self.B):
+            if lam_max_t is not None:
+                a = cheb_coeffs_exp_autodiff(ts[j], lam_max_t, self.m)
+                b = cheb_coeffs_exp_autodiff(ts[j + 1], lam_max_t, self.m)
+                coef = a - b
+                acc = None
+                for k in range(self.m + 1):
+                    if abs(float(coef.data[k])) > 0.0:
+                        term = Y[k] * coef[k]
+                        acc = term if acc is None else acc + term
+            else:
+                a = cheb_coeffs_exp(ts[j], lam_max, self.m)
+                b = cheb_coeffs_exp(ts[j + 1], lam_max, self.m)
+                coef = a - b
+                acc = None
+                for k in range(self.m + 1):
+                    if coef[k] != 0.0:
+                        term = Y[k] * float(coef[k])
+                        acc = term if acc is None else acc + term
+            out.append(acc)
+        self.last_ts = ts
+        return out
+
+    # -------------------------------------------------------------- forward
+    def __call__(self, x: Tensor):
+        n = x.shape[0]
+        d = self.d
+        xn = self.norm(x)
+
+        # ---- ROUTER --------------------------------------------------------
+        q = xn @ self.Wq
+        k = xn @ self.Wk
+        scores = (q @ k.T) * (d ** -0.5)
+        causal = np.triu(np.ones((n, n), dtype=bool), k=1)
+        sc_masked = scores.data - 1e9 * causal
+        idx = np.argsort(-sc_masked, axis=1)[:, : self.k]   # (n, k) neighbours
+        row = np.arange(n)[:, None]
+        sel = scores[row, idx]                                # Tensor (n, k)
+
+        # ---- TOPOLOGY: semantic edges + temporal chain ----------------------
+        sem = edges_from_topk(idx)
+        chain = chain_edges(n)
+
+        # ---- SHEAF: content-adaptive signed restriction maps -----------------
+        if sem.shape[0] > 0:
+            eu = xn[sem[:, 0]]
+            ev = xn[sem[:, 1]]
+            feat = concat([eu, ev], axis=1)
+            h = (feat @ self.edge1).tanh()
+            alpha = (h @ self.edge2).reshape(-1)              # (m,) in (-1, 1)
+        else:
+            alpha = Tensor(np.zeros(0))
+        L_sem = self._laplacian(n, sem, alpha)
+
+        if chain.shape[0] > 0:
+            beta = softplus(self.beta_chain)                   # (1,1) > 0
+            a_ch = beta.reshape(1) * Tensor(np.ones(chain.shape[0]))
+            L_chain = self._laplacian(n, chain, a_ch)
+        else:
+            L_chain = Tensor(np.zeros((n, n)))
+        L = L_sem + L_chain
+
+        # Spectral scale.  We use a GUARANTEED UPPER BOUND (Gershgorin +
+        # smooth max) for the Chebyshev normalisation: an estimate that
+        # undershoots lambda_max makes the recurrence explode.  The bound is
+        # differentiable, so no gradient is lost.
+        lam_bound_t = spectral_bound_autodiff(L)
+        lam_est = power_iteration_lmax(L.data, iters=40)
+        lam_max_t = lam_bound_t
+        lam_max = float(lam_max_t.data)
+        if lam_max < 1e-8:
+            lam_max = 1e-8
+            lam_max_t = Tensor(np.array([[lam_max]]))
+
+        # ---- FREEZE: multi-scale band filter bank --------------------------
+        bands = self._bands(L, lam_max, xn, lam_max_t)
+        gates = softplus(self.band_g).reshape(self.B)
+        z = None
+        for j in range(self.B):
+            term = bands[j] * gates[j]
+            z = term if z is None else z + term
+
+        # ---- HARMONIC BRANCH: the frozen state (exact Hodge projection) -----
+        # For a graph Laplacian, ker(L) = per-component constants, so the
+        # Hodge limit of the diffusion is the PER-COMPONENT MEAN.  The band
+        # bank only carries the transient, so without this branch the block
+        # cannot represent a global aggregate at all (measured: sign acc =
+        # chance on the global-mean task).  Cost: O(n + m) via union-find.
+        all_edges = (
+            np.concatenate([sem, chain], axis=0)
+            if sem.shape[0] or chain.shape[0]
+            else np.zeros((0, 2), dtype=np.int64)
+        )
+        harmonic = self._harmonic(xn, all_edges)
+        _, diag_ncomp = connected_components(n, all_edges)
+        z = z + self.harm_g * harmonic
+
+        # ---- READ: mix channels, one nonlinearity, residual -----------------
+        z = z @ self.Wmix
+        z = z.gelu()
+        out = x + z @ self.Wo
+
+        # ---- sharp retrieval path (winner-take-all over the routed edges) ---
+        if self.use_route_head:
+            v = xn @ self.Wv
+            vg = v.gather_rows(idx)                            # (n, k, d)
+            w = softmax_rows(sel)                              # (n, k)
+            sr = (vg * w[:, :, None]).sum(axis=1)              # (n, d)
+            out = out + sr @ self.Wr
+
+        # ---- diagnostics ---------------------------------------------------
+        try:
+            desc = topo_descriptor(n, sem, np.abs(alpha.data) + 1e-9)
+        except Exception:
+            desc = np.zeros(9)
+        self.diag = {
+            "n_edges": int(sem.shape[0]),
+            "n_components": int(diag_ncomp),
+            "lam_max": float(lam_max),
+            "lam_est": float(lam_est),
+            "alpha_abs_mean": (
+                float(np.abs(alpha.data).mean()) if alpha.data.size else 0.0
+            ),
+            "alpha_neg_frac": (
+                float((alpha.data < 0).mean()) if alpha.data.size else 0.0
+            ),
+            "band_energy": np.array(
+                [float((band.data**2).sum()) for band in bands]
+            ),
+            "topo": desc,
+            "balance_res": float(np.max(np.abs(L.data.sum(axis=1)))),
+        }
+        return out

fluidlayer/bpe.py

@@ -0,0 +1,173 @@
+"""
+fluidlayer/bpe.py — a byte-level BPE tokenizer, in pure Python (no downloads).
+
+WHY THIS EXISTS
+---------------
+Every experiment so far ran at the CHARACTER level (vocab 76).  That is
+convenient but it is not what a real deployment looks like: real LLM
+tokenizers are subword, with a vocabulary of 32k-256k and a very different
+redundancy structure.  The character-level corpus is *maximally* redundant
+(a 12-gram of characters repeats constantly), which flatters a
+content-keyed store.  A subword corpus repeats less.
+
+The point of this file is therefore not "a tokenizer" -- it is the ability
+to re-run every measurement at a realistic token granularity, on CPU,
+without downloading anything.
+
+THE ALGORITHM
+-------------
+1. split the text on a word-ish regex,
+2. count word frequencies,
+3. represent every word as its UTF-8 BYTES (so the tokenizer is lossless
+   for any input, including text the training corpus never saw),
+4. repeatedly merge the most frequent adjacent symbol pair,
+5. emit a vocabulary: the 256 byte tokens + every learned merge.
+
+`encode` is then exact and reversible (`decode(encode(s)) == s` for any
+string, because the byte fallback is always available).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from collections import Counter
+
+__all__ = ["BPETokenizer", "train_bpe"]
+
+_WORD = re.compile(r"\S+|\s+")
+
+
+class BPETokenizer:
+    """A byte-level BPE tokenizer."""
+
+    def __init__(self, merges: dict, id_bytes: dict):
+        # merges: {(a, b): new_id}  -- the merge ORDER matters
+        # id_bytes: {id: bytes}   -- every id, bytes and merged alike
+        self.merges = merges
+        self.ranks = {p: i for i, p in enumerate(merges.keys())}
+        self.id_to_bytes = dict(id_bytes)
+
+    # ------------------------------------------------------------------ size
+    @property
+    def vocab_size(self) -> int:
+        return len(self.id_to_bytes)
+
+    # --------------------------------------------------------------- encode
+    def _word_symbols(self, word: str):
+        b = word.encode("utf-8")
+        return tuple(b)
+
+    def _merge_word(self, symbols):
+        while len(symbols) > 1:
+            best = None
+            best_i = -1
+            for i in range(len(symbols) - 1):
+                r = self.ranks.get((symbols[i], symbols[i + 1]))
+                if r is not None and (best is None or r < best):
+                    best = r
+                    best_i = i
+            if best is None:
+                break
+            a, b = symbols[best_i], symbols[best_i + 1]
+            new = self.merges.get((a, b))
+            if new is None:
+                break
+            symbols = symbols[:best_i] + (new,) + symbols[best_i + 2:]
+        return symbols
+
+    def encode(self, text: str):
+        ids = []
+        for w in _WORD.findall(text):
+            ids.extend(self._merge_word(self._word_symbols(w)))
+        return ids
+
+    def decode(self, ids) -> str:
+        out = bytearray()
+        for i in ids:
+            out += self.id_to_bytes[int(i)]
+        return out.decode("utf-8", errors="replace")
+
+    # ------------------------------------------------------------------ io
+    def save(self, path: str):
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump({
+                "merges": [[list(k), v] for k, v in self.merges.items()],
+                "id_bytes": {str(k): list(v) for k, v in self.id_to_bytes.items()},
+            }, f)
+
+    @staticmethod
+    def load(path: str) -> BPETokenizer:
+        with open(path, encoding="utf-8") as f:
+            d = json.load(f)
+        merges = {tuple(m[0]): m[1] for m in d["merges"]}
+        id_bytes = {int(k): bytes(v) for k, v in d["id_bytes"].items()}
+        return BPETokenizer(merges, id_bytes)
+
+
+def train_bpe(text: str, vocab_size: int = 1024, verbose: bool = True):
+    """Train a byte-level BPE on `text`.  Deterministic, no randomness."""
+    words = Counter(_WORD.findall(text))
+    # each word -> its symbol tuple; count pairs over the whole corpus
+    word_syms = {w: tuple(w.encode("utf-8")) for w in words}
+    merges: dict = {}
+    id_bytes: dict = {i: bytes([i]) for i in range(256)}
+
+    def recount():
+        pairs = Counter()
+        for w, s in word_syms.items():
+            c = words[w]
+            for i in range(len(s) - 1):
+                pairs[(s[i], s[i + 1])] += c
+        return pairs
+
+    n_merges = vocab_size - 256
+    for step in range(n_merges):
+        pairs = recount()
+        if not pairs:
+            break
+        # deterministic tie-break: the highest count, then the lowest ids
+        best = max(pairs.items(), key=lambda kv: (kv[1], -kv[0][0], -kv[0][1]))
+        pair, cnt = best
+        if cnt < 2:
+            break
+        new_id = 256 + len(merges)
+        merges[pair] = new_id
+        id_bytes[new_id] = id_bytes[pair[0]] + id_bytes[pair[1]]
+        # rewrite every word
+        for w in list(word_syms):
+            s = word_syms[w]
+            out = []
+            i = 0
+            while i < len(s):
+                if i + 1 < len(s) and (s[i], s[i + 1]) == pair:
+                    out.append(new_id)
+                    i += 2
+                else:
+                    out.append(s[i])
+                    i += 1
+            word_syms[w] = tuple(out)
+        if verbose and step % 100 == 0:
+            print(f"    merge {step}: {pair} -> {new_id}  (count {cnt})")
+    return BPETokenizer(merges, id_bytes)
+
+
+if __name__ == "__main__":
+    import argparse
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--corpus", default="corpus.txt")
+    ap.add_argument("--vocab", type=int, default=1024)
+    ap.add_argument("--out", default="prior/bpe_1024.json")
+    args = ap.parse_args()
+    with open(args.corpus, encoding="utf-8", errors="ignore") as f:
+        text = f.read()
+    tok = train_bpe(text, args.vocab)
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+    tok.save(args.out)
+    ids = tok.encode("Alice was beginning to get very tired of sitting by "
+                      "her sister on the bank, and of having nothing to do.")
+    print(f"    vocab {tok.vocab_size}  tokens {len(ids)}  "
+          f"-> {ids[:12]} ...")
+    print(f"    round-trip ok: {tok.decode(ids)[:60]!r}")
+    print(f"    wrote {args.out}")

fluidlayer/data.py

@@ -0,0 +1,92 @@
+"""
+fluidlayer.data — tiny corpora + char-level tokenizer.
+
+Prefers a real public-domain text (Project Gutenberg).  If the sandbox has no
+outbound network access it falls back to a deterministic synthetic corpus so
+that the A/B training run is always reproducible.
+"""
+
+from __future__ import annotations
+
+import os
+import urllib.request
+
+import numpy as np
+
+__all__ = ["get_corpus", "CharTokenizer", "batchify", "GUTENBERG_URLS"]
+
+GUTENBERG_URLS = [
+    "https://www.gutenberg.org/files/11/11-0.txt",     # Alice in Wonderland
+    "https://www.gutenberg.org/cache/epub/11/pg11.txt",
+]
+
+
+def get_corpus(max_chars: int = 120_000, cache: str | None = None) -> str:
+    """Return a text corpus (>= max_chars if possible)."""
+    cache = cache or os.path.join(os.path.dirname(__file__), "..", "corpus.txt")
+    if os.path.exists(cache):
+        with open(cache, encoding="utf-8", errors="ignore") as f:
+            t = f.read()
+        if len(t) >= 1000:
+            return t[:max_chars]
+    for url in GUTENBERG_URLS:
+        try:
+            with urllib.request.urlopen(url, timeout=12) as r:
+                raw = r.read().decode("utf-8", errors="ignore")
+            start = raw.find("*** START")
+            if start >= 0:
+                raw = raw[start:]
+            end = raw.find("*** END")
+            if end > 0:
+                raw = raw[:end]
+            raw = raw.replace("\r\n", "\n")
+            if len(raw) > 2000:
+                try:
+                    os.makedirs(os.path.dirname(os.path.abspath(cache)), exist_ok=True)
+                    with open(cache, "w", encoding="utf-8") as f:
+                        f.write(raw)
+                except OSError:
+                    pass
+                return raw[:max_chars]
+        except Exception:
+            continue
+    return _synthetic_corpus(max_chars)
+
+
+def _synthetic_corpus(n: int) -> str:
+    """Deterministic fallback: a Zipfian character process with word structure."""
+    rng = np.random.default_rng(1234)
+    words = ["the", "cat", "sat", "on", "a", "mat", "and", "then", "said", "little",
+             "queen", "of", "hearts", "alice", "rabbit", "hole", "tea", "party",
+             "night", "tree", "house", "river", "cloud", "stone", "light", "voice"]
+    p = np.arange(1, len(words) + 1) ** -1.1
+    p /= p.sum()
+    out = []
+    while sum(len(w) + 1 for w in out) < n:
+        k = rng.choice(len(words), size=rng.integers(4, 14), p=p)
+        out.extend(words[i] for i in k)
+    s = " ".join(out)
+    return s[:n]
+
+
+class CharTokenizer:
+    def __init__(self, text: str):
+        self.chars = sorted(set(text))
+        self.stoi = {c: i for i, c in enumerate(self.chars)}
+        self.itos = self.chars
+
+    def encode(self, text: str) -> np.ndarray:
+        return np.array([self.stoi[c] for c in text], dtype=np.int64)
+
+    def decode(self, ids) -> str:
+        return "".join(self.itos[int(i)] for i in ids)
+
+    @property
+    def vocab(self) -> int:
+        return len(self.chars)
+
+
+def batchify(tokens: np.ndarray, batch_size: int) -> np.ndarray:
+    """(N,) -> (batch_size, N//batch_size) contiguous split."""
+    nb = len(tokens) // batch_size
+    return tokens[: nb * batch_size].reshape(batch_size, nb)