seq/tasks.py

"""
Canonical attention-diagnostic synthetic tasks — the suite the field uses to decide whether a
non-attention architecture can really stand in for a Transformer.

All tasks share one causal next-token interface:
    sample(batch, device) -> (inputs[B,T] long, targets[B,T] long, mask[B,T] {0,1} float)
Loss/accuracy are computed only on masked positions. Token id 0 is a reserved filler/pad.

  * MQAR          Multi-Query Associative Recall (Arora et al. 2023, Zoology/Based). THE test that
                  separates real attention-alternatives from impostors: in-context key->value
                  lookup for many queries, over a gap. Linear models with fixed state struggle as
                  #pairs grows; attention solves it trivially.
  * AssocRecall   single-query AR (Ba et al.; Hyena) — easier sanity version of MQAR.
  * SelectiveCopy Mamba's selective-copy: copy the data tokens in order, ignoring interspersed
                  fillers. Requires INPUT-DEPENDENT gating (content-selective memory).
  * Induction     in-context induction-head probe: [.. A B .. A] -> predict B. The ICL primitive.
"""
from __future__ import annotations

import torch


def _distinct_per_row(B, n, lo, hi, device):
    """n distinct ints in [lo,hi) for each of B rows -> (B,n) long. hi-lo must be >= n."""
    span = hi - lo
    perm = torch.argsort(torch.rand(B, span, device=device), dim=1)[:, :n]
    return perm + lo


class MQAR:
    """Multi-Query Associative Recall (Zoology/Based standard, dense). vocab in [1,V); 0 filler.
    Sequence = [k0 v0 ... k_{D-1} v_{D-1}]  (+optional filler gap)  [q_1 q_2 ... q_M], where each
    query q is one of the bound keys. The TARGET at each query position is that key's bound value
    (NOT a next token — the answer never appears in the input), masked to query positions only.
    Dense queries give rich supervision; capacity is set by the #bindings D (the B1b axis)."""

    def __init__(self, vocab=64, num_pairs=8, num_queries=None, gap=0, key_frac=0.5):
        # DISJOINT token ranges: keys in [1, n_key), values in [n_key, vocab). This removes
        # key/value collisions so recall is unambiguous and the capacity ceiling (B1b) is clean.
        self.n_key = max(num_pairs + 1, int(vocab * key_frac))
        assert self.n_key - 1 >= num_pairs, "need >= num_pairs distinct keys in the key range"
        assert vocab - self.n_key >= 2, "need a value range"
        self.vocab = vocab
        self.D = num_pairs
        self.M = num_queries if num_queries is not None else max(2 * num_pairs, 8)
        self.gap = gap
        self.seq_len = 2 * num_pairs + gap + self.M
        self.name = f"MQAR(V={vocab},keys<{self.n_key},pairs={num_pairs},q={self.M},gap={gap})"

    def sample(self, B, device):
        D, M, V = self.D, self.M, self.vocab
        keys = _distinct_per_row(B, D, 1, self.n_key, device)     # (B,D) distinct keys
        vals = torch.randint(self.n_key, V, (B, D), device=device)  # (B,D) values (disjoint range)
        ctx = torch.empty(B, 2 * D, dtype=torch.long, device=device)
        ctx[:, 0::2] = keys
        ctx[:, 1::2] = vals
        qi = torch.randint(0, D, (B, M), device=device)           # which binding each query hits
        qk = torch.gather(keys, 1, qi)                            # query keys (the input)
        qa = torch.gather(vals, 1, qi)                            # bound values (the target)
        if self.gap > 0:
            filler = torch.zeros(B, self.gap, dtype=torch.long, device=device)
            inp = torch.cat([ctx, filler, qk], dim=1)
        else:
            inp = torch.cat([ctx, qk], dim=1)
        T = inp.shape[1]
        tgt = torch.zeros(B, T, dtype=torch.long, device=device)
        mask = torch.zeros(B, T, device=device)
        qstart = 2 * D + self.gap
        tgt[:, qstart:] = qa                                      # recall target at each query pos
        mask[:, qstart:] = 1.0
        return inp, tgt, mask


class AssocRecall:
    """Single-query associative recall: many KV pairs, one query at the end."""

    def __init__(self, vocab=64, num_pairs=16, gap=0):
        self.inner = MQAR(vocab=vocab, num_pairs=num_pairs, num_queries=1, gap=gap)
        self.vocab = vocab
        self.seq_len = self.inner.seq_len
        self.name = f"AssocRecall(V={vocab},pairs={num_pairs},gap={gap})"

    def sample(self, B, device):
        return self.inner.sample(B, device)


class SelectiveCopy:
    """Copy the K data tokens (in order) out of an L-slot memory region full of fillers.
    vocab: filler=0, data in [1, vocab-1), marker=vocab-1.
    input = [memory(L) with K data tokens scattered, marker, d1, d2, ..., dK]
    predict d1..dK at positions marker..d_{K-1} (next-token), masked there."""

    def __init__(self, vocab=32, mem_len=64, n_data=16, fixed=False):
        assert n_data <= mem_len
        self.vocab = vocab
        self.L = mem_len
        self.K = n_data
        self.fixed = fixed     # True -> data at fixed evenly-spaced positions (control variant)
        self.marker = vocab - 1
        self.seq_len = mem_len + 1 + n_data
        self.name = f"SelectiveCopy(V={vocab},mem={mem_len},k={n_data},{'fixed' if fixed else 'selective'})"

    def sample(self, B, device):
        L, K, V = self.L, self.K, self.vocab
        mem = torch.zeros(B, L, dtype=torch.long, device=device)            # fillers
        if self.fixed:
            base = torch.linspace(0, L - 1, K, device=device).long()        # fixed spacing (control)
            pos = base[None].expand(B, K).clone()
        else:
            pos = torch.argsort(torch.rand(B, L, device=device), dim=1)[:, :K]  # K random slots/row
        pos, _ = torch.sort(pos, dim=1)                                     # left-to-right order
        data = torch.randint(1, V - 1, (B, K), device=device)              # data tokens
        mem.scatter_(1, pos, data)
        marker = torch.full((B, 1), self.marker, dtype=torch.long, device=device)
        inp = torch.cat([mem, marker, data], dim=1)                        # teacher-forced output
        T = inp.shape[1]
        tgt = torch.zeros_like(inp)
        tgt[:, :-1] = inp[:, 1:]
        mask = torch.zeros(B, T, device=device)
        mask[:, L:L + K] = 1.0     # positions: marker (predict d1) .. d_{K-1} (predict dK)
        return inp, tgt, mask


class Induction:
    """In-context induction: prefix contains a unique bigram [A,B]; the final token is A; predict B.
    vocab tokens in [1, vocab); 0 filler. mask=1 only at the final position."""

    def __init__(self, vocab=32, seq_len=64):
        self.vocab = vocab
        self.seq_len = seq_len + 1   # +1 for the trailing query A
        self.name = f"Induction(V={vocab},len={seq_len})"
        self._L = seq_len

    def sample(self, B, device):
        L, V = self._L, self.vocab
        seq = torch.randint(1, V, (B, L), device=device)
        A = torch.randint(1, V, (B,), device=device)
        B_tok = torch.randint(1, V, (B,), device=device)
        # ensure B != A so the answer is informative
        clash = B_tok == A
        B_tok[clash] = (B_tok[clash] % (V - 1)) + 1
        B_tok[B_tok == A] = (A[B_tok == A] % (V - 1)) + 1
        # remove any existing A from the prefix (so A is unique once we place the bigram)
        seq[seq == A[:, None]] = 0
        # also keep the bigram-following slot clean: place [A,B] at a random early position
        ppos = torch.randint(0, L - 2, (B,), device=device)
        ar = torch.arange(B, device=device)
        seq[ar, ppos] = A
        seq[ar, ppos + 1] = B_tok
        # any filler 0 left from removal -> replace with a token guaranteed != A
        filler = (A % (V - 1)) + 1
        zero = seq == 0
        seq = torch.where(zero, filler[:, None].expand_as(seq), seq)
        # but that replacement might have re-introduced A if filler==A (filler!=A by construction)
        # re-place bigram in case a filler overwrote it (positions fixed, so re-assert)
        seq[ar, ppos] = A
        seq[ar, ppos + 1] = B_tok
        query = A[:, None]
        inp = torch.cat([seq, query], dim=1)            # (B, L+1)
        T = inp.shape[1]
        tgt = torch.zeros(B, T, dtype=torch.long, device=device)
        tgt[:, -1] = B_tok
        mask = torch.zeros(B, T, device=device)
        mask[:, -1] = 1.0
        return inp, tgt, mask


class MixedMQAR:
    """MQAR trained over a SPECTRUM of difficulties: each training batch samples the number of KV
    pairs d ~ U[min_pairs, max_pairs] (the standard Zoology training distribution). The easy
    instances supply the gradient that BOOTSTRAPS the recall circuit, so high-D recall becomes
    learnable in feasible compute — where FIXED-high-D training stalls at chance (no foothold for the
    sharp phase transition). EVALUATION is fixed at the target (max_pairs), so the reported number is
    target-D recall. Identical for both models -> fair. n_key is stable across d (= vocab*key_frac
    while d < that), so key/value ranges don't shift with difficulty."""

    def __init__(self, vocab=256, max_pairs=64, num_queries=128, gap=0, min_pairs=1):
        self.vocab = vocab
        self.max_pairs = max_pairs
        self.M = num_queries
        self.gap = gap
        self.min_pairs = min_pairs
        self.seq_len = 2 * max_pairs + gap + num_queries          # max layout (for max_len sizing)
        self.name = f"MixedMQAR(V={vocab},pairs={min_pairs}-{max_pairs},q={num_queries},gap={gap})"

    def _mqar(self, d):
        return MQAR(vocab=self.vocab, num_pairs=d, num_queries=self.M, gap=self.gap)

    def sample(self, B, device):            # TRAINING: a random difficulty per batch
        # Sample the scalar difficulty on CPU: distribution-identical (U[min,max]) but avoids a
        # per-step GPU->CPU sync that serializes the (tiny) model on CUDA. Batch tensors below
        # are still generated on `device`.
        d = int(torch.randint(self.min_pairs, self.max_pairs + 1, (1,)).item())
        return self._mqar(d).sample(B, device)

    def eval_sample(self, B, device):       # EVAL: fixed at the TARGET difficulty (max_pairs)
        return self._mqar(self.max_pairs).sample(B, device)


TASK_REGISTRY = {
    "mqar": MQAR,
    "mixed_mqar": MixedMQAR,
    "assoc": AssocRecall,
    "selcopy": SelectiveCopy,
    "induction": Induction,
}


if __name__ == "__main__":
    dev = torch.device("cpu")
    for name, cls in TASK_REGISTRY.items():
        t = cls()
        x, y, m = t.sample(4, dev)
        print(f"{t.name:<40} x{tuple(x.shape)} masked={int(m.sum().item())}")
        # sanity: an oracle that knows the mapping would score 1.0; random ~1/vocab