statelock 0.1.0__py3-none-any.whl

statelock/__init__.py

@@ -0,0 +1,18 @@
+"""statelock -- verified finite machines inside neural networks.
+
+Three things, in the order you would use them:
+  induce(...)      work the machine out from labelled strings, by search
+  StateDiscoverer  or let a network find it by gradient descent
+  equivalent(...)  check the result against a reference by exhaustive BFS
+  StateLock        attach the verified machine to your model
+
+Everything here is decided by exact equivalence, not by held-out accuracy.
+"""
+from .verify import equivalent, reachable, minimal_size
+from .induce import induce
+from .discover import StateDiscoverer, fit
+from .layer import StateLock, attach
+
+__version__ = "0.1.0"
+__all__ = ["induce", "equivalent", "reachable", "minimal_size",
+           "StateDiscoverer", "fit", "StateLock", "attach"]

statelock/discover.py

@@ -0,0 +1,114 @@
+"""Discover the automaton by gradient descent, with no oracle-driven search.
+
+The state is a learned transition table run as a hard argmax scan with a
+straight-through gradient. What makes it work, and what four earlier attempts
+lacked, is the target: the state is trained to predict its MYHILL-NERODE ROW --
+the labels the sequence would receive under a bundle of probe continuations --
+rather than only the current label. Predicting the current label gives the
+optimiser no reason to separate two states that agree now and differ later, and
+a partition that is coarser than the target's cannot be repaired afterwards,
+because merging only coarsens.
+
+Three settings matter and were each fixed by measurement, not taste:
+  * train on strings that COVER the alphabet, not only in-distribution ones.
+    With grammatical data alone, loss reached 0.0 and every held-out label was
+    right, and exact equivalence still failed, because the table was
+    unconstrained on strings the generator never emits.
+  * interleave sequence lengths rather than staging them, or the short gradient
+    path is forgotten.
+  * hold the straight-through temperature constant. The forward pass is a hard
+    argmax at any temperature, so annealing only sharpens the surrogate until
+    the gradient vanishes.
+
+Exactness tracks the training loss: runs reaching zero verified exact, runs that
+stalled did not. So `fit` reports the loss and you restart on that, without
+needing an answer key.
+"""
+import numpy as np
+import torch
+import torch.nn as nn
+
+__all__ = ["StateDiscoverer", "fit"]
+
+
+class _Scan(nn.Module):
+    def __init__(self, V, K, d, seed=0):
+        super().__init__()
+        g = torch.Generator().manual_seed(seed)
+        self.W = nn.Parameter(torch.randn(V, K, K, generator=g) * 0.5)
+        self.s0 = nn.Parameter(torch.randn(K, generator=g) * 0.5)
+        self.emb = nn.Embedding(K, d)
+        with torch.no_grad():
+            self.emb.weight.normal_(0, 1.0, generator=g)
+        self.V, self.K = V, K
+
+    def forward(self, x, tau=1.0):
+        B, L = x.shape
+        p = torch.softmax(self.s0 / tau, -1).expand(B, self.K)
+        hard = torch.zeros_like(p).scatter_(1, p.argmax(-1, keepdim=True), 1.0)
+        st = hard + p - p.detach()
+        outs = []
+        for t in range(L):
+            logits = torch.einsum('bk,bkj->bj', st, self.W[x[:, t]])
+            p = torch.softmax(logits / tau, -1)
+            hard = torch.zeros_like(p).scatter_(1, p.argmax(-1, keepdim=True), 1.0)
+            st = hard + p - p.detach()
+            outs.append(st)
+        return torch.stack(outs, 1) @ self.emb.weight
+
+    def table(self):
+        return (self.W.argmax(-1).detach().cpu().numpy().astype(np.int64),
+                int(self.s0.argmax().item()))
+
+
+class StateDiscoverer(nn.Module):
+    """A discrete state layer that learns its own transition table."""
+
+    def __init__(self, alphabet_size, n_slots, n_probes, d=64, seed=0):
+        super().__init__()
+        self.scan = _Scan(alphabet_size, n_slots, d, seed)
+        self.row = nn.Linear(d, n_probes)
+
+    def forward(self, x, tau=1.0):
+        """Returns per-position row logits."""
+        return self.row(self.scan(x, tau))
+
+    def states(self, x):
+        """The discovered state index at each position (no gradient)."""
+        delta, s0 = self.scan.table()
+        x = x.cpu().numpy()
+        s = np.full(x.shape[0], s0, np.int64)
+        out = np.empty(x.shape, np.int64)
+        for t in range(x.shape[1]):
+            s = delta[x[:, t], s]; out[:, t] = s
+        return out
+
+    def table(self):
+        return self.scan.table()
+
+
+def fit(model, batches, steps=40000, lr=2e-2, tau=1.0, log_every=0):
+    """Train on an iterable-returning callable `batches(step) -> (x, rows)`.
+
+    x:    LongTensor (B, L) of symbols
+    rows: FloatTensor (B, L, n_probes) of 0/1 probe labels
+    Returns the smoothed final loss. Exactness tracks this: restart if it does
+    not approach zero.
+    """
+    head = [p for n, p in model.named_parameters() if not n.startswith("scan.W")
+            and not n.startswith("scan.s0")]
+    opt = torch.optim.AdamW([{"params": [model.scan.W, model.scan.s0], "lr": lr},
+                             {"params": head, "lr": lr / 3}], weight_decay=0.0)
+    sch = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=[lr, lr / 3], total_steps=steps)
+    lf = nn.BCEWithLogitsLoss(); ema = None
+    for it in range(steps):
+        x, rows = batches(it)
+        loss = lf(model(x, tau), rows)
+        opt.zero_grad(); loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step(); sch.step()
+        v = float(loss.detach())
+        ema = v if ema is None else 0.98 * ema + 0.02 * v
+        if log_every and it % log_every == 0:
+            print(f"step {it} loss {ema:.5f}", flush=True)
+    return ema

statelock/induce.py

@@ -0,0 +1,90 @@
+"""Induce an exact automaton from (string, label) pairs, by search.
+
+The learner never sees a state. It needs an oracle it can query for the label of
+a string it chooses -- `label_fn(word) -> int` -- which is the same access a
+person has when they can run the system they are modelling on inputs of their
+choosing.
+
+Construction: a cover set holding one representative prefix per discovered
+state, closed under the alphabet; a suffix set grown by consistency repair when
+two prefixes share a row but their extensions do not. Both corrections were
+forced by verification failures, not chosen for elegance:
+  * sampling prefixes does not close the table (64k sampled prefixes still left
+    transitions undefined, and the verifier refused);
+  * a fixed suffix set under-separates (K = 149 against a minimal 152, with no
+    unresolved transitions -- a wrong merge that looks nearly right).
+"""
+import numpy as np
+
+__all__ = ["induce"]
+
+
+def induce(alphabet_size, label_fn, suffixes=None, max_rounds=200, max_states=100000):
+    """Return (delta, out, start, info).
+
+    alphabet_size: number of input symbols, which are 0..alphabet_size-1.
+    label_fn: callable taking a tuple of symbols, returning an int label.
+    suffixes: optional starting suffix set (list of tuples). Defaults to the
+        empty string plus every single symbol; more are added as needed.
+    """
+    V = alphabet_size
+    S = list(suffixes) if suffixes else [()] + [(a,) for a in range(V)]
+    P = [()]
+    cache = {}
+
+    def lab(w):
+        if w not in cache: cache[w] = int(label_fn(w))
+        return cache[w]
+
+    def row(p): return tuple(lab(p + s) for s in S)
+
+    rounds = added = 0
+    reps = {}; sid = {}; cover = {}
+    while rounds < max_rounds:
+        rounds += 1
+        e1 = {p + (a,) for p in P for a in range(V)}
+        e2 = {q + (b,) for q in e1 for b in range(V)}
+        cand = sorted(set(P) | e1 | e2, key=lambda w: (len(w), w))
+        rows = {p: row(p) for p in cand}
+        reps = {}; sid = {}; cover = {}
+        for p in cand:
+            r = rows[p]
+            if r not in reps:
+                reps[r] = len(reps); cover[reps[r]] = p
+            sid[p] = reps[r]
+            if len(reps) > max_states:
+                raise RuntimeError("state count exceeded max_states; the target "
+                                   "may not be finite-state at this resolution")
+        bad = None
+        bystate = {}
+        for p in cand: bystate.setdefault(sid[p], []).append(p)
+        for _, ps in bystate.items():
+            if len(ps) < 2: continue
+            p0 = ps[0]
+            for q in ps[1:]:
+                for a in range(V):
+                    u, w = p0 + (a,), q + (a,)
+                    if u in rows and w in rows and rows[u] != rows[w]:
+                        j = next(i for i in range(len(S)) if rows[u][i] != rows[w][i])
+                        bad = (a,) + S[j]; break
+                if bad: break
+            if bad: break
+        if bad is not None and bad not in S:
+            S.append(bad); added += 1; continue
+        newP = sorted(set(cover.values()), key=lambda w: (len(w), w))
+        if set(newP) == set(P): break
+        P = newP
+
+    K = len(reps)
+    delta = np.zeros((V, K), np.int64); unresolved = 0
+    for i in range(K):
+        p = cover[i]
+        for a in range(V):
+            q = p + (a,)
+            if q in sid: delta[a, i] = sid[q]
+            else: unresolved += 1
+    out = np.array([lab(cover[i]) for i in range(K)], np.int64)
+    info = {"states": K, "suffixes": len(S), "suffixes_added": added,
+            "rounds": rounds, "unresolved_transitions": unresolved,
+            "queries": len(cache)}
+    return delta, out, sid[()], info

statelock/layer.py

@@ -0,0 +1,70 @@
+"""Attach a verified machine to a model.
+
+`StateLock` turns a transition table into an embedding lookup over machine
+states and adds it to a hidden-state tensor. The table is fixed: it is never
+trained, so whatever was verified about it stays true inside the network.
+
+Two ways to use it:
+  * directly, adding the state embedding to your own hidden states;
+  * `attach(hf_model, lock, input_ids_fn)` installs a forward hook on a Hugging
+    Face model's embedding layer, so the state is added to the token embeddings
+    with no change to the model's code.
+"""
+import numpy as np
+import torch
+import torch.nn as nn
+
+__all__ = ["StateLock", "attach"]
+
+
+class StateLock(nn.Module):
+    def __init__(self, delta, start=0, d_model=768, out=None):
+        super().__init__()
+        delta = np.asarray(delta, dtype=np.int64)
+        self.register_buffer("delta", torch.from_numpy(delta))
+        self.start = int(start)
+        self.n_states = delta.shape[1]
+        self.emb = nn.Embedding(self.n_states, d_model)
+        nn.init.normal_(self.emb.weight, std=0.02)
+        self.register_buffer("out", torch.from_numpy(np.asarray(out, dtype=np.int64))
+                             if out is not None else torch.zeros(0, dtype=torch.long))
+
+    @torch.no_grad()
+    def run(self, x):
+        """State index at every position. Exact, O(length), never trained."""
+        d = self.delta.cpu().numpy(); xi = x.cpu().numpy()
+        s = np.full(xi.shape[0], self.start, np.int64)
+        st = np.empty(xi.shape, np.int64)
+        for t in range(xi.shape[1]):
+            s = d[xi[:, t], s]; st[:, t] = s
+        return torch.from_numpy(st).to(x.device)
+
+    def forward(self, hidden, x):
+        """hidden: (B, L, d). x: (B, L) symbols. Returns hidden + state embedding."""
+        return hidden + self.emb(self.run(x))
+
+
+def attach(hf_model, lock, symbol_fn=None):
+    """Add the machine's state to a Hugging Face model's token embeddings.
+
+    symbol_fn maps input_ids to machine symbols; omit it when the model's token
+    ids already are the machine's alphabet. Returns a handle; call
+    handle.remove() to detach.
+    """
+    emb = hf_model.get_input_embeddings()
+    state = {"ids": None}
+
+    def pre_hook(_module, args, kwargs):
+        ids = args[0] if args else kwargs.get("input_ids")
+        state["ids"] = ids
+        return None
+
+    def hook(_module, args, output):
+        ids = args[0] if args else state["ids"]
+        if ids is None:
+            return output
+        sym = ids if symbol_fn is None else symbol_fn(ids)
+        return output + lock.emb(lock.run(sym)).to(output.dtype)
+
+    hf_model.register_forward_pre_hook(pre_hook, with_kwargs=True)
+    return emb.register_forward_hook(hook)

statelock/verify.py

@@ -0,0 +1,65 @@
+"""Exact equivalence checking. This is the load-bearing function of the whole
+library: it decides whether two automata agree on EVERY string, by breadth-first
+search over reachable pairs of states, with no length cutoff.
+
+Accuracy on sampled strings cannot substitute. In the work this library comes
+from, a bounded enumeration to length 11 reported 6/6 correct and disagreed with
+true equivalence on 22 of 48 deliberately corrupted targets.
+"""
+from collections import deque
+import numpy as np
+
+__all__ = ["equivalent", "reachable", "minimal_size"]
+
+
+def equivalent(delta_a, out_a, delta_b, out_b, start_a=0, start_b=0):
+    """True if the two DFAs produce the same output on every input string.
+
+    delta_*: int array (alphabet_size, n_states); delta[x, s] is the next state.
+    out_*:   int array (n_states,); the output symbol emitted in that state.
+    """
+    delta_a, out_a = np.asarray(delta_a), np.asarray(out_a)
+    delta_b, out_b = np.asarray(delta_b), np.asarray(out_b)
+    A = delta_a.shape[0]
+    if delta_b.shape[0] != A:
+        raise ValueError("alphabet sizes differ")
+    seen = {(int(start_a), int(start_b))}
+    q = deque(seen)
+    while q:
+        a, b = q.popleft()
+        for x in range(A):
+            a2, b2 = int(delta_a[x, a]), int(delta_b[x, b])
+            if int(out_a[a2]) != int(out_b[b2]):
+                return False
+            if (a2, b2) not in seen:
+                seen.add((a2, b2)); q.append((a2, b2))
+    return True
+
+
+def reachable(delta, start=0):
+    """Indices of states reachable from `start`."""
+    delta = np.asarray(delta); A = delta.shape[0]
+    seen = {int(start)}; q = deque(seen)
+    while q:
+        s = q.popleft()
+        for x in range(A):
+            t = int(delta[x, s])
+            if t not in seen:
+                seen.add(t); q.append(t)
+    return sorted(seen)
+
+
+def minimal_size(delta, out):
+    """Number of states in the minimal equivalent DFA (Moore refinement)."""
+    delta, out = np.asarray(delta), np.asarray(out)
+    K = delta.shape[1]; part = out.copy()
+    while True:
+        sig = [tuple([int(part[i])] + [int(part[delta[x, i]]) for x in range(delta.shape[0])])
+               for i in range(K)]
+        m = {}; new = np.empty(K, np.int64)
+        for i, s in enumerate(sig):
+            if s not in m: m[s] = len(m)
+            new[i] = m[s]
+        if len(m) == len(set(part.tolist())):
+            return len(set(part.tolist()))
+        part = new

statelock-0.1.0.dist-info/METADATA

@@ -0,0 +1,128 @@
+Metadata-Version: 2.4
+Name: statelock
+Version: 0.1.0
+Summary: Verified finite machines inside neural networks: induce one from labelled strings or discover it by gradient descent, check it by exhaustive equivalence, and attach it to your model.
+License: MIT
+Project-URL: Source, https://github.com/RaheemNazir/statelock
+Keywords: state tracking,automata,length generalization,transformers,neural algorithmic reasoning,verification
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.22
+Requires-Dist: torch>=2.0
+Dynamic: license-file
+
+# statelock
+
+Transformers lose track of state as inputs get longer. On a prose connectivity
+task, a model goes from 0.92 accuracy at its training length to 0.50 at 16x
+length, and growing it from 1.3M to 88M parameters does not help: the largest
+model is the best in distribution and the worst at length.
+
+`statelock` gives the model a finite machine to consult, and the machine is
+**checked exhaustively** rather than measured on a test set. With one attached,
+the same model holds 1.0000 at 16x length.
+
+```python
+from statelock import induce, equivalent, StateLock
+
+delta, out, start, info = induce(alphabet_size=8, label_fn=my_system)   # find it
+assert equivalent(delta, out, ref_delta, ref_out, start, ref_start)      # check it
+lock = StateLock(delta, start, d_model=768)                              # attach it
+hidden = lock(hidden, symbols)            # inside your model's forward
+```
+
+That is the whole interface.
+
+## Install
+
+```
+pip install statelock
+```
+
+## The three things it does
+
+**Induce.** Work the machine out from labelled strings. No state supervision, no
+hints. You supply `label_fn(word) -> int`, the label your system gives a string
+you choose. On a 152-state target it recovers exactly 152 states, equal to the
+minimal machine, and passes exact equivalence.
+
+**Discover.** Or let a network find the machine during ordinary training, with
+no search at all:
+
+```python
+from statelock import StateDiscoverer, fit
+model = StateDiscoverer(alphabet_size=6, n_slots=32, n_probes=64)
+loss = fit(model, batches, steps=3000)     # exactness tracks this loss
+delta, start = model.table()
+```
+
+The state is trained to predict its Myhill-Nerode row, the labels a bundle of
+probe continuations would produce, not just the current label. Predicting the
+current label alone gives the optimiser no reason to separate two states that
+agree now and differ later, and four earlier attempts failed for exactly that
+reason. Exactness tracks the training loss: a run that reaches zero verifies,
+one that stalls does not, so you restart on the loss without needing an answer
+key.
+
+**Verify.** `equivalent(...)` decides agreement on *every* string by BFS over
+reachable state pairs, with no length cutoff. This is not optional rigour. A
+bounded enumeration to length 11 once reported 6/6 correct here and disagreed
+with true equivalence on 22 of 48 corrupted targets.
+
+## With Hugging Face
+
+```python
+from statelock import StateLock, attach
+lock = StateLock(delta, start, d_model=model.config.hidden_size)
+handle = attach(model, lock, symbol_fn=my_tokens_to_symbols)   # forward hook
+```
+
+No change to the model's code. `handle.remove()` detaches.
+
+## What it is for, and what it is not for
+
+Use it where the task has a discrete core with a state count in the hundreds:
+entity and relation tracking across a document, protocol and permission state,
+bracket and scope structure, small formal languages, board and game state.
+
+It will not help where there is no such core, and it does not make a model
+better at anything except keeping that state. The machine is a component you
+attach, and attaching it is a design decision about your task.
+
+## Measured
+
+Prose connectivity, dense per-position labels, trained at length 128:
+
+| model | arm | L=128 | L=2048 (16x) |
+|---|---|---|---|
+| 1.3M | none | 0.9155 | 0.5848 |
+| 1.3M | shuffled machine | 0.8984 | 0.6033 |
+| 1.3M | statelock | 1.0000 | 0.9998 |
+| 27M | none | 0.9099 | 0.6191 |
+| 27M | statelock | 1.0000 | 1.0000 |
+| 88M | none | 0.9167 | 0.5028 |
+| 88M | statelock | 1.0000 | 1.0000 |
+
+Constant-predictor baseline 0.5204. The shuffled-machine row is the control that
+matters: same parameters, same embedding, random transitions, no gain.
+
+Discovery (the network finds the machine itself, then it is verified):
+
+| target minimal states | slots offered | exact | loss on exact runs |
+|---|---|---|---|
+| 11 | 32 | 2/3 seeds | 0.0 |
+| 38 | 128 | 1/3 seeds | 0.0 |
+
+## Reproduce
+
+```
+python tests/test_all.py
+```
+
+Induction to 152 states, 40 corruptions all flagged by the verifier, exact
+running at length 4096, and a discovery run that verifies. Runs on CPU.
+
+## Licence
+
+MIT.

statelock-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+statelock/__init__.py,sha256=bYYRj2IVMb-l0qv2k4RcyYyWtcbxWJ4BAlO2-7QoDvI,775
+statelock/discover.py,sha256=7GDOIzW7KACAK5QJUiS5PgGmv2_iy6j1J7mvZKJ212Q,4817
+statelock/induce.py,sha256=6n7VyKTvYutcbLjvLPyW2GgK-oqGUDNs-ioEsHeESEA,3689
+statelock/layer.py,sha256=V39porJMt2sWKy-9gjIswUOrQ5CUvVD8nvIAtd27Pck,2705
+statelock/verify.py,sha256=sFhV01_iggZE1dnQMu4doWX2EngsysRPZBr9S7KzYtE,2435
+statelock-0.1.0.dist-info/licenses/LICENSE,sha256=JK2p0N_CKg82aKhciBBf1_pONnCSyhuOAAn3Gw0P7WU,1075
+statelock-0.1.0.dist-info/METADATA,sha256=_--5kw7Tau553QUWlgDDRDuL7MtmQ2JxPjKpqq1L328,4728
+statelock-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+statelock-0.1.0.dist-info/top_level.txt,sha256=OHWmWuALtqlMVydbdI0I3LPMLEFJe7XJj5HaEAF__PI,10
+statelock-0.1.0.dist-info/RECORD,,

statelock-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

statelock-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Abdul Raheem Nazir
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

statelock-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+statelock