cascade.py

"""
A CHEAP MODEL THAT KNOWS WHEN TO ASK FOR HELP.

Escalating between rungs of a nested storage ladder failed, and it failed
for a reason that was nothing to do with the policy: every rung of that
ladder holds a different number of VALUES and costs exactly the same to
RUN. Escalating there buys nothing, because the coarse pass is not cheap --
it is the same matrix product as the fine one.

The compute axis supplies rungs that genuinely differ in cost. A
convolution holding 433 values needs 110,592 multiplies per image; an
unconstrained layer of the same shape needs 3,145,728, a factor of
twenty-eight. That is a real cascade: run the cheap one, answer where it is
confident, and pay for the expensive one only where it is not.

    cost = cheap + (escalated fraction) x expensive

so escalating a quarter of the examples costs about 897,000 multiplies
against the dense model's 3,146,000 -- three and a half times less -- and
the question is what accuracy survives that.

This is the compute axis applied PER EXAMPLE rather than per model, and it
is the first thing in this programme where a routing decision has an
obvious payoff even when it is imperfect: a wrong escalation costs
arithmetic, not accuracy, and a wrong non-escalation costs one example.

Three references. The cheap model alone, the expensive model alone, and an
ORACLE cascade that escalates exactly the examples the cheap model gets
wrong and the expensive one gets right. The oracle bounds what any
confidence rule could achieve; the gap between it and the measured rule is
what a better signal would be worth.
"""

import numpy as np
import time
import json
import os

try:
    import cupy as _cp
    _GPU = _cp.cuda.runtime.getDeviceCount() > 0
except Exception:
    _GPU = False
xp = _cp if _GPU else np
DT = np.float32


def to_dev(a, dtype=DT):
    a = np.asarray(a, dtype=dtype)
    return xp.asarray(a) if _GPU else a


def to_host(a):
    return _cp.asnumpy(a) if _GPU and isinstance(a, _cp.ndarray) else np.asarray(a)


def windowed(g, c_in, k, stride, c_out):
    """A convolution as a tying pattern, over colour input. Its cost as an
    ALGORITHM is one multiply per output unit per tap, not the dense
    product the fold is trained through."""
    go = g // stride
    ni, no = c_in*g*g, c_out*go*go
    ii, jj = np.meshgrid(np.arange(ni), np.arange(no), indexing='ij')
    ci, pi = ii // (g*g), ii % (g*g)
    co, po = jj // (go*go), jj % (go*go)
    off = k//2 if stride == 1 else 0
    dr = pi // g - ((po // go)*stride - off)
    dc = pi % g - ((po % go)*stride - off)
    inside = (dr >= 0) & (dr < k) & (dc >= 0) & (dc < k)
    K = c_in*c_out*k*k + 1
    idx = np.where(inside, (ci*c_out + co)*k*k + dr*k + dc, K-1)
    return idx.ravel().astype(np.int32), K, no, go*go*c_out*(c_in*k*k)



# ---------- deterministic accumulation ------------------------------------

DETERMINISTIC = True    # set False to restore the atomic path
# Switched on after seed 0 was trained three times atomically and
# gave 0.6822, 0.6829 and 0.6851 for the dear arm — a range of
# 0.0029, where the parity question turns on 0.0015. Earlier
# checkpoints are atomic and should be deleted before rerunning.

_FIXED = {}


class FixedScatter:
    """The folded gradient sum, in an order fixed once and reused.

    The atomic path is genuinely non-deterministic on a GPU: the same
    additions in a different order differed by 3.4e-04 on 5 of 5 repeats,
    and a 200-step run drifted 5.8e-03. Amplified over a real run that
    becomes about 0.002 of final accuracy — a floor beneath which no effect
    is measurable however many seeds are added.

    The repair: sort the index once, so every group is a contiguous slice,
    and reduce each slice the same way every time.

    A FIRST VERSION laid every group into one rectangle of (K x longest
    group) and cost 149 GB. Convolution indices are wildly non-uniform: all
    the out-of-window positions map to a single padding group, which for a
    middle layer of a three-layer stack holds 96.8% of 16.7 million entries
    while every real group holds 256. So the rectangle is used only for
    groups near the median, and any oversized group is summed over its own
    slice — a tree reduction over a fixed extent, equally deterministic and
    costing nothing.

    THE FOLD IS WHY THIS IS CHEAP AT ALL: the partition never changes
    during training, so the sort is paid once for the whole run rather than
    once a step."""

    def __init__(self, idx, K, cap=8192):
        h = to_host(idx).astype(np.int64).reshape(-1)
        order = np.argsort(h, kind="stable")
        counts = np.bincount(h, minlength=K)
        starts = np.cumsum(counts) - counts
        big = np.where(counts > cap)[0]
        small = np.where(counts <= cap)[0]
        self.K = K
        self.n = len(h)
        self.order = to_dev(order, np.int64) if _GPU else order
        # the oversized groups: contiguous slices of the sorted gradient
        self.big = [(int(b), int(starts[b]), int(starts[b]+counts[b]))
                    for b in big]
        # everything else, in one rectangle
        self.small = to_dev(small, np.int64) if _GPU else small
        self.width = int(counts[small].max()) if len(small) else 0
        if self.width:
            pos = np.concatenate([np.arange(counts[s]) for s in small]) \
                if len(small) else np.zeros(0, np.int64)
            src = np.concatenate([np.arange(starts[s], starts[s]+counts[s])
                                  for s in small]) if len(small) else \
                np.zeros(0, np.int64)
            row = np.repeat(np.arange(len(small)), counts[small])
            self.src = to_dev(src, np.int64) if _GPU else src
            self.slot = to_dev(row*self.width + pos, np.int64) if _GPU \
                else (row*self.width + pos)
            self.buf = xp.zeros(len(small)*self.width, DT)
        self.bytes = (len(small)*self.width*4 if self.width else 0) \
            + self.n*8
        self._keep = idx           # so the cache key cannot be reused

    def __call__(self, g):
        gs = g.reshape(-1)[self.order]
        out = xp.zeros(self.K, DT)
        if self.width:
            self.buf[:] = 0
            self.buf[self.slot] = gs[self.src]
            out[self.small] = self.buf.reshape(-1, self.width).sum(1)
        for b, a, z in self.big:
            out[b] = gs[a:z].sum()
        return out

def _atomic(dW, idx, K):
    g = xp.zeros(K, DT)
    if _GPU:
        import cupyx
        cupyx.scatter_add(g, idx, dW.reshape(-1))
    else:
        np.add.at(g, idx, dW.reshape(-1))
    return g


def scatter(dW, idx, K):
    if not DETERMINISTIC:
        return _atomic(dW, idx, K)
    key = (id(idx), K, int(np.prod(to_host(idx).shape)))
    if key not in _FIXED:
        _FIXED[key] = FixedScatter(idx, K)
    return _FIXED[key](dW)


def augment(Xb, side, pad, rg):
    """Per-image crop and flip on CHANNEL-MAJOR data."""
    n = Xb.shape[0]
    im = Xb.reshape(n, 3, side, side)
    if pad:
        P = xp.zeros((n, 3, side+2*pad, side+2*pad), dtype=Xb.dtype)
        P[:, :, pad:pad+side, pad:pad+side] = im
        oy = rg.integers(0, 2*pad+1, n); ox = rg.integers(0, 2*pad+1, n)
        oy = xp.asarray(oy) if _GPU else oy
        ox = xp.asarray(ox) if _GPU else ox
        ar = xp.arange(side)
        im = P[xp.arange(n)[:, None, None, None],
               xp.arange(3)[None, :, None, None],
               (oy[:, None]+ar[None, :])[:, None, :, None],
               (ox[:, None]+ar[None, :])[:, None, None, :]]
    fl = rg.random(n) < 0.5
    fl = xp.asarray(fl) if _GPU else fl
    im = xp.where(fl[:, None, None, None], im[:, :, :, ::-1], im)
    return im.reshape(n, -1)


def train(Xtr, Ytr, Xte, layers, nc, cfg, seed, side, label="",
          yte=None, ckpt=None, ckpt_file="cascade_probs.npz"):
    """A stack of folded layers, each normalised.

    layers is a list of (index, K, width, taps). A one-element list is the
    cheap arm; three elements is the expensive one. Normalisation matters:
    without it a three-layer stack was WORSE than a single layer (0.4783
    falling to 0.4478); with it, better (0.6229 rising to 0.6613)."""
    D = Xtr.shape[1]; L = len(layers)
    rg = np.random.default_rng(seed)
    idxs = [to_dev(a, np.int32) if _GPU else a for a, _, _, _ in layers]
    Ks = [K for _, K, _, _ in layers]
    outs = [h for _, _, h, _ in layers]
    taps = [t for _, _, _, t in layers]
    ins = [D] + outs[:-1]
    P = []
    for l in range(L):
        v = rg.normal(0, np.sqrt(2.0/taps[l]), Ks[l]).astype(np.float32)
        v[-1] = 0.0
        P.append(to_dev(v))
    VAL = lambda l: l
    GAM = lambda l: L + l
    BET = lambda l: 2*L + l
    HEAD, OB = 3*L, 3*L + 1
    P += [xp.ones(outs[l], DT) for l in range(L)]
    P += [xp.zeros(outs[l], DT) for l in range(L)]
    P += [to_dev(rg.normal(0, np.sqrt(2.0/outs[-1]), (outs[-1], nc))),
          xp.zeros(nc, DT)]
    M = [xp.zeros_like(p) for p in P]; V = [xp.zeros_like(p) for p in P]
    n = Xtr.shape[0]; t = 0
    ag = np.random.default_rng(seed + 991)
    t0 = time.time()

    def fwd(x, keep=False):
        cache = []; h = x
        for l in range(L):
            W = P[VAL(l)][idxs[l]].reshape(ins[l], outs[l])
            z = h @ W
            var = z.var(1, keepdims=True) + 1e-5
            zn = (z - z.mean(1, keepdims=True))/xp.sqrt(var)
            zs = zn*P[GAM(l)] + P[BET(l)]
            a = xp.maximum(zs, 0)
            if keep:
                cache.append((h, W, var, zn, zs))
            h = a
        return h, cache

    hist = []
    for ep in range(cfg["epochs"]):
        perm = ag.permutation(n)
        run_loss, nb = 0.0, 0
        for st in range(0, n, cfg["batch"]):
            b = perm[st:st+cfg["batch"]]
            x = Xtr[b]
            if cfg["augment"]:
                x = augment(x, side, cfg["aug_pad"], ag)
            y = Ytr[b]
            h, cache = fwd(x, keep=True)
            lg = h @ P[HEAD] + P[OB]
            mx = lg.max(1, keepdims=True)
            e = xp.exp(lg - mx)
            se = e.sum(1, keepdims=True)
            # the cross-entropy, which every step already computes on its
            # way to the gradient and which this programme threw away for
            # its whole length. The knee is invisible in accuracy and
            # obvious here.
            run_loss += float(to_host((-(lg - mx - xp.log(se))*y).sum(1).mean()))
            nb += 1
            d = (e/se - y)/len(b)
            G = [xp.zeros_like(p) for p in P]
            G[HEAD] = h.T @ d; G[OB] = d.sum(0)
            dh = d @ P[HEAD].T
            for l in range(L-1, -1, -1):
                hin, W, var, zn, zs = cache[l]
                dzs = dh*(zs > 0)
                G[GAM(l)] = (dzs*zn).sum(0); G[BET(l)] = dzs.sum(0)
                dzn = dzs*P[GAM(l)]
                dz = (dzn - dzn.mean(1, keepdims=True)
                      - zn*(dzn*zn).mean(1, keepdims=True))/xp.sqrt(var)
                G[VAL(l)] = scatter(hin.T @ dz, idxs[l], Ks[l])
                if l > 0:
                    dh = dz @ W.T
            t += 1
            for i, (p_, gr) in enumerate(zip(P, G)):
                M[i] = 0.9*M[i] + 0.1*gr
                V[i] = 0.999*V[i] + 0.001*gr*gr
                P[i] = p_ - cfg["lr"]*(M[i]/(1-0.9**t)) \
                    / (xp.sqrt(V[i]/(1-0.999**t))+1e-8)
        # progress, because these arms take tens of minutes and a run that
        # prints nothing for half an hour is indistinguishable from a hung one
        if label and ((ep+1) % cfg.get("report_every", 10) == 0
                      or ep == cfg["epochs"]-1):
            msg = f"    {label}  epoch {ep+1:3d}/{cfg['epochs']}"
            if yte is not None:
                # the probabilities are computed here anyway for the
                # progress line, so CHECKPOINT them. A run once died ten
                # minutes into a ninety-minute arm and lost all of it,
                # while the finished arm beside it had been saved.
                pr = []
                for s2 in range(0, Xte.shape[0], 4096):
                    h, _ = fwd(Xte[s2:s2+4096])
                    lg = h @ P[HEAD] + P[OB]
                    e = xp.exp(lg - lg.max(1, keepdims=True))
                    pr.append(to_host(e/e.sum(1, keepdims=True)))
                pr = np.concatenate(pr)
                tl = float(-np.log(np.clip(
                    pr[np.arange(len(yte)), yte], 1e-12, 1.0)).mean())
                acc = float((pr.argmax(1) == yte).mean())
                hist.append(dict(epoch=ep+1, train_loss=run_loss/max(nb, 1),
                                 test_loss=tl, test_acc=acc))
                msg += (f"   train {run_loss/max(nb,1):.4f}"
                        f"  test {tl:.4f}  acc {acc:.4f}")
                if ckpt:
                    old = {}
                    if os.path.exists(ckpt_file):
                        old = dict(np.load(ckpt_file))
                    old[ckpt] = pr; old["yte"] = yte
                    old[ckpt + "_epoch"] = np.array(ep+1)
                    # KEEP EVERY SNAPSHOT, not just the latest. Each epoch's
                    # state is a different model, so one training run already
                    # contains a family — and for a folded model the copies
                    # are almost free: a snapshot of this arm is 433 values.
                    # The probabilities are computed here anyway for the
                    # progress line, so keeping them costs a few hundred
                    # kilobytes and throwing them away costs the experiment.
                    if cfg.get("keep_snapshots", True):
                        old[f"{ckpt}_ep{ep+1:04d}"] = pr
                    old[ckpt + "_hist"] = np.array(
                        [[h_["epoch"], h_["train_loss"], h_["test_loss"],
                          h_["test_acc"]] for h_ in hist], dtype=np.float64)
                    np.savez(ckpt_file, **old)
            print(msg + f"   [{time.time()-t0:.0f}s]", flush=True)
    train.last_hist = hist
    out = []
    for s2 in range(0, Xte.shape[0], 4096):
        h, _ = fwd(Xte[s2:s2+4096])
        lg = h @ P[HEAD] + P[OB]
        e = xp.exp(lg - lg.max(1, keepdims=True))
        out.append(to_host(e/e.sum(1, keepdims=True)))
    return np.concatenate(out)



# ---------- CIFAR-10, verified before use --------------------------------

CIFAR_MEAN = np.array([0.4914, 0.4822, 0.4465])
CIFAR_STD = np.array([0.2470, 0.2435, 0.2616])


def _verified_cifar(extra_root=None, verbose=True):
    """Use a local CIFAR-10 copy if one passes verification.

    Checked rather than trusted: the published per-channel means are
    0.4914, 0.4822, 0.4465 and the deviations 0.2470, 0.2435, 0.2616, so a
    normalised copy fails; and neighbouring pixels in a photograph differ
    far less than random pairs, so a reordered one fails too. A copy named
    "preprocessed" was once accepted by shape alone and trained on for five
    hours with every arm at chance together.

    extra_root takes a kagglehub download:
        import kagglehub
        p = kagglehub.dataset_download("pankrzysiu/cifar10-python")
        main(cifar_path=p)
    """
    import os, glob, pickle
    roots = ([extra_root] if extra_root else []) + \
        sorted(glob.glob("/kaggle/input/*")) + \
        ["/kaggle/working", "./cifar", "./data", "."]
    for root in roots:
        if not root or not os.path.isdir(root):
            continue
        cands = []
        for p in glob.glob(os.path.join(root, "**", "*"), recursive=True):
            if not os.path.isfile(p):
                continue
            try:
                if p.endswith(".npy"):
                    cands.append(np.load(p, allow_pickle=True))
                elif (p.endswith((".pickle", ".pkl", ".p"))
                      or os.path.basename(p).startswith(("data_batch",
                                                         "test_batch"))):
                    with open(p, "rb") as fh:
                        d = pickle.load(fh, encoding="bytes")
                    if isinstance(d, dict):
                        for k, v in d.items():
                            k = k.decode() if isinstance(k, bytes) else k
                            a = np.asarray(v)
                            if a.size > 100 and k in ("data", "labels",
                                                      "fine_labels", "x", "y"):
                                cands.append(a)
            except Exception:
                continue
        imgs = [a for a in cands if a.ndim >= 2 and len(a) >= 1000
                and a.size // len(a) == 3072]
        labs = [np.asarray(a).ravel() for a in cands
                if a.ndim <= 2 and np.issubdtype(np.asarray(a).dtype,
                                                 np.integer)
                and 1000 <= a.size <= 100000]
        if not imgs or not labs:
            continue
        X = np.concatenate([a.reshape(len(a), -1) for a in imgs])
        y = np.concatenate(labs)
        if len(y) != len(X):
            continue
        y = y.astype(np.int64)
        Xf = X.astype(np.float64)
        if Xf.max() > 1.5:
            Xf = Xf/255.0
        why = None
        if int(y.max()) + 1 != 10 or y.min() != 0:
            why = f"labels are 0..{y.max()}"
        elif Xf.min() < -0.01 or Xf.max() > 1.01:
            why = f"values run {Xf.min():.2f} to {Xf.max():.2f}, not raw"
        if why is None:
            best = None
            for lay, shp in (("HWC", (-1, 32, 32, 3)),
                             ("CHW", (-1, 3, 32, 32))):
                im = Xf[:2000].reshape(shp)
                mu = im.mean((0, 1, 2)) if lay == "HWC" else im.mean((0, 2, 3))
                sd = im.std((0, 1, 2)) if lay == "HWC" else im.std((0, 2, 3))
                e = float(np.abs(mu-CIFAR_MEAN).max()
                          + np.abs(sd-CIFAR_STD).max())
                if best is None or e < best[0]:
                    best = (e, lay, mu)
            err, layout, mu = best
            if err > 0.05:
                why = f"channel means {np.round(mu,3)}, so not raw"
            else:
                gm = (Xf[:512].reshape(-1, 32, 32, 3).mean(-1)
                      if layout == "HWC"
                      else Xf[:512].reshape(-1, 3, 32, 32).mean(1))
                adj = float(np.abs(np.diff(gm, axis=2)).mean())
                sh = gm.reshape(len(gm), -1).copy()
                rg = np.random.default_rng(0)
                for r in sh:
                    rg.shuffle(r)
                rnd = float(np.abs(np.diff(sh, axis=1)).mean())
                if adj/max(rnd, 1e-9) >= 0.5:
                    why = "neighbours differ as much as random pairs, so reordered"
        if why:
            if verbose:
                print(f"  rejected {root}: {why}", flush=True)
            continue
        if verbose:
            print(f"  verified CIFAR-10 at {root} ({layout}, {len(X):,} "
                  f"images) — no download needed", flush=True)
        Xr = Xf.astype(np.float32)
        Xr = (Xr.reshape(-1, 32, 32, 3) if layout == "HWC"
              else Xr.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1))
        return Xr, y
    if verbose:
        print("  no verified CIFAR-10 found; downloading", flush=True)
    return None


def load(cfg):
    got = _verified_cifar(cfg.get("cifar_path"))
    if got is not None:
        X, y = got
    else:
        from tensorflow import keras
        (a, b), (c, d) = keras.datasets.cifar10.load_data()
        X = np.concatenate([a, c]).astype(np.float32)/255.0
        y = np.concatenate([b, d]).ravel().astype(np.int64)
    side = 32
    if cfg["grid"] != side:
        s = side // cfg["grid"]
        X = X.reshape(-1, cfg["grid"], s, cfg["grid"], s, 3).mean(axis=(2, 4))
        side = cfg["grid"]
    rg = np.random.default_rng(0); p = rg.permutation(len(X))
    tr, te = p[:cfg["n_train"]], p[cfg["n_train"]:cfg["n_train"]+10000]
    mu = X[tr].mean((0, 1, 2), keepdims=True)
    sd = X[tr].std((0, 1, 2), keepdims=True) + 1e-8
    f = lambda Z: ((Z-mu)/sd).transpose(0, 3, 1, 2).reshape(len(Z), -1)
    Y = np.zeros((len(tr), 10), np.float32); Y[np.arange(len(tr)), y[tr]] = 1
    return f(X[tr]), Y, f(X[te]), y[te], side


# 120 epochs, because at 60 NEITHER ARM HAD CONVERGED and the expensive
# one was pulling away: the gap went 0.0199 to 0.0454 over the last thirty
# epochs, with the cheap arm gaining +0.0083 in its last ten against the
# deep arm's +0.0139. So the cascade was measured mid-training, and which
# way that moves the economics is not predictable — a wider gap makes each
# escalation worth more, but the cheap arm contributes less of the final
# accuracy. One seed: this run answers convergence, not variance.
# ---------- where checkpoints live ----------------------------------------
# Colab wipes the working directory when a session ends, so a run that dies
# loses every checkpoint and starts over. Drive survives. This picks Drive
# when it is mounted and falls back to the working directory otherwise.

def _ckpt_dir():
    for d in ("/content/drive/MyDrive/pin_runs",
              "/kaggle/working", "."):
        try:
            if d.startswith("/content") and not os.path.isdir(
                    "/content/drive/MyDrive"):
                continue
            os.makedirs(d, exist_ok=True)
            t = os.path.join(d, ".w")
            open(t, "w").close(); os.remove(t)
            return d
        except Exception:
            continue
    return "."


CKPT_DIR = _ckpt_dir()


CFG = dict(grid=16, c_in=3, hidden=4096, n_train=50000, epochs=120,
           batch=128, lr=1e-3, augment=True, aug_pad=2, seeds=(0, 1, 2),
           report_every=10, cifar_path=None, reuse_cheap=True,
           keep_snapshots=True)
FRACTIONS = [0.0, 0.05, 0.1, 0.2, 0.3, 0.5, 0.75, 1.0]


def stack(g, ci, ch, depth):
    """A stack of 3x3 convolutions. The first maps the colour input; the
    rest map channels to channels, which is why they cost more."""
    ls = []; cin = ci; gi = g
    for l in range(depth):
        a, K, no, m = windowed(gi, cin, 3, 1, ch)
        ls.append((a, K, no, 9, m))
        cin = ch
    return ls


def one_seed(Xtr, Ytr, Xte, yte, side, cfg, seed):
    """The cheap arm is one convolution; the expensive arm is three.

    A first version used an UNCONSTRAINED layer as the expensive arm and
    the premise collapsed: the convolution beat it, 0.6141 to 0.5570, so
    escalating could only lose. Three stacked convolutions are genuinely
    better (0.6613 against 0.6229 measured separately) and genuinely
    dearer, which is what a cascade needs."""
    g, ci, hid = cfg["grid"], cfg["c_in"], cfg["hidden"]
    ch = max(1, hid//(g*g))
    one = stack(g, ci, ch, 1)
    three = stack(g, ci, ch, 3)
    # Checkpoint after EACH ARM, not after both. A first run was cut off
    # ten epochs from the end of the second arm and lost everything,
    # including a finished first arm — the expensive unit is the arm, so
    # that is where the save belongs.
    # ONE CHECKPOINT FILE PER SEED. A single shared file meant a rerun
    # with several seeds reused seed 0's cheap arm for every one of them,
    # pairing one cheap model against three different deep ones.
    cf = os.path.join(CKPT_DIR, f"cascade_probs_s{seed}.npz")
    have = dict(np.load(cf)) if os.path.exists(cf) else {}
    done = lambda k: (k in have
                      and int(have.get(k + "_epoch", 0)) >= cfg["epochs"])
    cheap = dear = None
    if cfg.get("reuse_cheap"):
        if done("cheap"):
            cheap = have["cheap"]
            print(f"    seed {seed}: reusing the saved cheap arm "
                  f"({(cheap.argmax(1) == yte).mean():.4f})", flush=True)
        if done("dear"):
            dear = have["dear"]
            print(f"    seed {seed}: reusing the saved dear arm "
                  f"({(dear.argmax(1) == yte).mean():.4f})", flush=True)
    if cheap is None:
        cheap = train(Xtr, Ytr, Xte, [l[:4] for l in one], 10, cfg, seed,
                      side, label=f"s{seed} cheap (1 conv) ", yte=yte,
                      ckpt="cheap", ckpt_file=cf)
    if dear is None:
        dear = train(Xtr, Ytr, Xte, [l[:4] for l in three], 10, cfg, seed,
                     side, label=f"s{seed} dear  (3 convs)", yte=yte,
                     ckpt="dear", ckpt_file=cf)
    print(f"    seed {seed} checkpointed in {cf}", flush=True)
    return cheap, dear, sum(l[4] for l in one), sum(l[4] for l in three)


def main(**over):
    CFG.update(over)
    t0 = time.time()
    print("=" * 78)
    print("A CHEAP MODEL THAT KNOWS WHEN TO ASK FOR HELP")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:9s} = {v}")
    g, ci, hid = CFG["grid"], CFG["c_in"], CFG["hidden"]
    ch = max(1, hid//(g*g))
    one = stack(g, ci, ch, 1); three = stack(g, ci, ch, 3)
    mc = sum(l[4] for l in one); md = sum(l[4] for l in three)
    kc = sum(l[1] for l in one); kd = sum(l[1] for l in three)
    print(f"\n  cheap: 1 convolution,  {kc:,} stored values, "
          f"{mc:,} multiplies")
    print(f"  dear:  3 convolutions, {kd:,} stored values, "
          f"{md:,} multiplies")
    print(f"  ratio: {md/mc:.1f}x   (the expensive arm must be BETTER as")
    print(f"         well as dearer, which an unconstrained layer was not)")
    # measured: 166s for the cheap arm and about 1,820s for the dear one
    # at 60 epochs on 50,000 images, one seed
    per = (166 + 1820)*(CFG["epochs"]/60)*(CFG["n_train"]/50000)
    print(f"\n  expect roughly {per/60:.0f} minutes a seed at this setting,")
    print(f"  {len(CFG['seeds'])} seed(s) — progress prints every "
          f"{CFG.get('report_every',10)} epochs so a silent run is")
    print(f"  distinguishable from a hung one")
    print("=" * 78, flush=True)

    Xtr, Ytr, Xte, yte, side = load(CFG)
    Xtr = to_dev(Xtr); Ytr = to_dev(Ytr); Xte = to_dev(Xte)
    runs = []
    for s in CFG["seeds"]:
        cheap, dear, mc, md = one_seed(Xtr, Ytr, Xte, yte, side, CFG, s)
        cc = cheap.argmax(1) == yte
        dc = dear.argmax(1) == yte
        conf = cheap.max(1)
        r = dict(cheap=float(cc.mean()), dear=float(dc.mean()))
        for f in FRACTIONS:
            if f == 0:
                acc = cc.mean()
            elif f == 1:
                acc = dc.mean()
            else:
                thr = np.quantile(conf, f)
                esc = conf <= thr
                acc = np.where(esc, dc, cc).mean()
            r[f"esc {f:.2f}"] = float(acc)
        # the oracle: escalate exactly where cheap is wrong and dear is right
        need = ~cc & dc
        r["oracle"] = float(np.where(need, dc, cc).mean())
        r["oracle frac"] = float(need.mean())
        # and how well confidence identifies those
        if need.sum():
            r["conf finds them"] = float(
                (conf[need] < np.median(conf)).mean())
        runs.append(r)
        print(f"  seed {s}: cheap {r['cheap']:.4f}  dear {r['dear']:.4f}"
              f"   [{time.time()-t0:.0f}s]", flush=True)
        json.dump(runs, open(os.path.join(CKPT_DIR, "cascade.json"),
                            "w"), indent=2)
        cf = os.path.join(CKPT_DIR, f"cascade_probs_s{s}.npz")
        keep = dict(np.load(cf)) if os.path.exists(cf) else {}
        keep.update(cheap=cheap, dear=dear, yte=yte,
                    cheap_epoch=np.array(CFG["epochs"]),
                    dear_epoch=np.array(CFG["epochs"]))
        np.savez(cf, **keep)

    m = {k: float(np.mean([r[k] for r in runs])) for k in runs[0]}
    sd = {k: float(np.std([r[k] for r in runs])) for k in runs[0]}
    print("\n" + "=" * 78)
    print("  ACCURACY AGAINST ARITHMETIC")
    print("=" * 78)
    print(f"  {'escalated':>10s} {'multiplies':>12s} {'vs dense':>9s} "
          f"{'accuracy':>9s} {'vs dense':>9s}")
    for f in FRACTIONS:
        cost = mc + f*md
        a_ = m[f"esc {f:.2f}"]
        print(f"  {f*100:9.0f}% {cost:12,.0f} {md/cost:8.2f}x {a_:9.4f} "
              f"{a_-m['dear']:+9.4f}")
    print(f"\n  seed spread (worst) {max(sd.values()):.4f}")
    if len(runs) > 1:
        print(f"\n  PAIRED against the dear arm — each cascade and the dear")
        print(f"  model it is compared with share a seed, so an unpaired")
        print(f"  spread overstates the noise:\n")
        print(f"  {'escalated':>10s} {'difference':>12s} {'paired sd':>11s}"
              f" {'deviations':>11s}")
        for fr in FRACTIONS:
            if not (0 < fr < 1):
                continue
            d_ = np.array([r[f"esc {fr:.2f}"] - r["dear"] for r in runs])
            n_ = abs(d_.mean())/max(d_.std(), 1e-9)
            print(f"  {fr*100:9.0f}% {d_.mean():+12.4f} {d_.std():11.4f}"
                  f" {n_:11.1f}")
    print(f"  an oracle escalating exactly the {m['oracle frac']*100:.1f}% "
          f"that need it")
    print(f"  would reach {m['oracle']:.4f} at "
          f"{(mc + m['oracle frac']*md)/md:.2f}x the dense cost")

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    tol = 2*max(sd.values())
    # the cheapest escalation that stays within noise of the dense model
    ok = [f for f in FRACTIONS
          if m[f"esc {f:.2f}"] >= m["dear"] - tol and f < 1.0]
    print(f"  cheap alone {m['cheap']:.4f} at {md/mc:.1f}x less arithmetic")
    print(f"  dense alone {m['dear']:.4f}")
    if ok:
        f = min(ok); cost = mc + f*md
        print(f"\n  MATCHING THE DENSE MODEL COSTS {md/cost:.2f}x LESS")
        print(f"  arithmetic: escalating the least-confident "
              f"{f*100:.0f}% reaches")
        print(f"  {m[f'esc {f:.2f}']:.4f} against {m['dear']:.4f}, at "
              f"{cost:,.0f} multiplies.")
        print(f"  A confidence rule is enough here because a wrong escalation")
        print(f"  costs arithmetic rather than accuracy — which is why this")
        print(f"  works where every routing attempt before it did not.")
    else:
        print(f"\n  NO ESCALATION FRACTION MATCHES THE DENSE MODEL within")
        print(f"  noise. The cheap model's confidence does not identify the")
        print(f"  examples it needs help with, so the cascade pays full price")
        print(f"  for partial accuracy.")
    print(f"\n  total {time.time()-t0:.0f}s; wrote cascade.json")


if __name__ == "__main__":
    main()