#!/usr/bin/env python3
"""Independent checker for the integer operator corpus (extensions/integer-operators.md).

Python standard library only; it shares no code with wai-ops. Written from the
extension's text (and, for the source forms it lowers, integer-payloads), it:

  - recomputes every vector file under ops-conformance/vectors/, stage/ and graphs/:
    each graph vector's `load` (WHG1 parse and load steps 1-8, section 10.5), its
    `profiles` (section 14.1, from the bounds pass and verdict of sections 8.1-8.5),
    and each decode's `expect` (the run of section 4, every graph operator of
    section 5), and each stage vector's `expect` (section 5E); it writes each file
    again canonically (sections 8.5 and 14.1) and requires the bytes of the
    committed file;
  - re-lowers every parameter set certificates/index.json lists, by the canonical
    lowerings of section 13, and requires each lowered graph's bytes, each set's
    digest (SPEC 3.1), and that every lowered graph loads; and requires the graph
    of each graphs/lower_*.jsonl vector to be the role of the set its id names, as
    re-lowered.

The bounds pass decides the verdict only. It does not write the certificate of
section 8.5, which the certifier half of this checker will.

    python3 -I tools/wai_ops_check.py ops-conformance            # check
    python3 -I tools/wai_ops_check.py ops-conformance --write DIR # also write the files

Exit status 1 on any difference or failure.
"""
import hashlib
import json
import math
import os
import struct
import sys

CAP = 1 << 26
I64_MIN, I64_MAX = -(1 << 63), (1 << 63) - 1
I32_MIN, I32_MAX = -(1 << 31), (1 << 31) - 1


class Refuse(Exception):
    """A refusal with a code of section 4 or 5E, and, for a stage operator, `at`."""

    def __init__(self, code, at=None):
        super().__init__(code)
        self.code = code
        self.at = at


# --------------------------------------------------------------- section 2


def R(p, s, mode):
    """R_floor, R_half_up, R_half_even (section 2)."""
    if s == 0:
        return p
    if mode == 0:
        return p >> s
    if mode == 1:
        return (p + (1 << (s - 1))) >> s
    q = p >> s
    r = p - (q << s)
    h = 1 << (s - 1)
    if r > h or (r == h and q % 2 == 1):
        return q + 1
    return q


def exp_neg_q16(x):
    """exp_neg_q16(x), x <= 0, in the order section 2 writes it."""
    u = -x
    n = u // 45426
    g = 45426 * n
    r = x + g
    t = r + 88670
    t2 = t * t
    q = t2 >> 16
    c = 23495 * q
    y = (c >> 16) + 22544
    return 0 if n >= 63 else y >> n


def rsqrt_q16(v):
    r = math.isqrt(v)
    return 0 if r == 0 else (1 << 32) // r


def trunc_div(a, b):
    q = abs(a) // abs(b)
    return q if (a >= 0) == (b > 0) else -q


# --------------------------------------------------------------- WHG1 (section 10)

NAMES = {
    0x01: "conv", 0x02: "rescale", 0x03: "chscale", 0x04: "lut16", 0x05: "add",
    0x06: "concat", 0x07: "chunk_add", 0x08: "split", 0x09: "pixel_shuffle",
    0x0A: "pixel_unshuffle", 0x0B: "clamp", 0x0C: "relu", 0x0D: "leaky",
    0x0E: "pad_rep", 0x0F: "crop_to", 0x10: "mul", 0x11: "conv_transpose",
    0x12: "igdn", 0x13: "abs", 0x14: "crop", 0x15: "subpos_split",
    0x16: "subpos_merge", 0x17: "fill", 0x18: "layernorm_c", 0x19: "l2norm",
    0x1A: "softmax", 0x1B: "sigmoid", 0x1C: "elu", 0x1D: "gram", 0x1E: "apply",
}
CODE = {v: k for k, v in NAMES.items()}
TWO_IN = {"add", "crop_to", "mul", "gram", "apply"}


def arity(name):
    if name in TWO_IN:
        return (2, 1)
    if name == "concat":
        return (None, 1)
    if name == "subpos_split":
        return (1, 4)
    if name == "subpos_merge":
        return (4, 1)
    return (1, 1)


class Reader:
    def __init__(self, b, pos=0):
        self.b, self.p = b, pos

    def take(self, n):
        if n < 0 or self.p + n > len(self.b):
            raise Refuse("model_mismatch")
        s = self.b[self.p:self.p + n]
        self.p += n
        return s

    def u(self, n):
        return int.from_bytes(self.take(n), "little")

    def i(self, n):
        return int.from_bytes(self.take(n), "little", signed=True)

    def ints(self, count, width):
        raw = self.take(count * width)
        return [int.from_bytes(raw[k:k + width], "little", signed=True) for k in range(0, len(raw), width)]

    def zeros(self, n):
        if any(self.take(n)):
            raise Refuse("model_mismatch")


def bad():
    raise Refuse("model_mismatch")


def read_rq(rd, C):
    """The RQ block (section 10.4) over C channels."""
    rq = {}
    rq["round"], rq["bias"], rq["m_kind"], rq["s_kind"] = rd.u(1), rd.u(1), rd.u(1), rd.u(1)
    rq["m_bytes"], rq["b_bytes"] = rd.u(1), rd.u(1)
    if rd.u(2) != 0:
        bad()
    if rq["round"] > 2 or rq["bias"] > 2 or rq["m_kind"] > 2 or rq["s_kind"] not in (1, 2):
        bad()
    if rq["m_kind"] == 0:
        if rq["m_bytes"] != 0:
            bad()
    elif rq["m_bytes"] not in (4, 8):
        bad()
    if rq["bias"] == 0:
        if rq["b_bytes"] != 0:
            bad()
    elif rq["b_bytes"] not in (4, 8):
        bad()
    nm = {0: 0, 1: 1, 2: C}[rq["m_kind"]]
    rq["mult"] = rd.ints(nm, rq["m_bytes"]) if nm else []
    ns = 1 if rq["s_kind"] == 1 else C
    rq["shift"] = list(rd.take(ns))
    if any(s > 127 for s in rq["shift"]):
        bad()
    rq["biasv"] = rd.ints(C, rq["b_bytes"]) if rq["bias"] else []
    return rq


def rq_m(rq, o):
    return None if rq["m_kind"] == 0 else rq["mult"][0 if rq["m_kind"] == 1 else o]


def rq_s(rq, o):
    return rq["shift"][0 if rq["s_kind"] == 1 else o]


def rq_apply(rq, o, v):
    """rq_o(v), section 5.1."""
    t = v + (rq["biasv"][o] if rq["bias"] == 1 else 0)
    m = rq_m(rq, o)
    p = t * (1 if m is None else m)
    q = R(p, rq_s(rq, o), rq["round"])
    return q + (rq["biasv"][o] if rq["bias"] == 2 else 0)


def read_operand(rd, n_scalars):
    kind = rd.u(1)
    rd.zeros(3)
    v = rd.u(4)
    if kind == 0:
        return ("const", v)
    if kind == 1 and v < n_scalars:
        return ("scalar", v)
    bad()


def read_params(name, pb, cin, n_scalars):
    rd = Reader(pb)
    P = {}
    if name == "conv":
        for k in ("cout", "kh", "kw", "sh", "sw", "ph", "pw", "groups"):
            P[k] = rd.u(4)
        P["w_bytes"] = rd.u(1)
        rd.zeros(3)
        if min(P["cout"], P["kh"], P["kw"], P["sh"], P["sw"], P["groups"]) == 0 or P["w_bytes"] not in (1, 2, 4, 8):
            bad()
        if cin % P["groups"] or P["cout"] % P["groups"]:
            bad()
        P["rq"] = read_rq(rd, P["cout"])
        P["w"] = rd.ints(P["cout"] * (cin // P["groups"]) * P["kh"] * P["kw"], P["w_bytes"])
    elif name == "conv_transpose":
        for k in ("cout", "kh", "kw", "sh", "sw", "ph", "pw", "oph", "opw"):
            P[k] = rd.u(4)
        P["w_bytes"] = rd.u(1)
        rd.zeros(3)
        if min(P["cout"], P["kh"], P["kw"], P["sh"], P["sw"]) == 0 or P["w_bytes"] not in (1, 2, 4, 8):
            bad()
        P["rq"] = read_rq(rd, P["cout"])
        P["w"] = rd.ints(cin * P["cout"] * P["kh"] * P["kw"], P["w_bytes"])
    elif name == "rescale":
        P["rq"] = read_rq(rd, cin)
    elif name == "chscale":
        P["scalar"], P["n_rows"] = rd.u(4), rd.u(4)
        P["round"], P["m_bytes"], P["shift"] = rd.u(1), rd.u(1), rd.u(1)
        rd.zeros(1)
        if P["scalar"] >= n_scalars or P["n_rows"] == 0 or P["m_bytes"] not in (4, 8) or P["shift"] > 127 or P["round"] > 2:
            bad()
        P["M"] = rd.ints(P["n_rows"] * cin, P["m_bytes"])
    elif name == "mul":
        P["rq"] = read_rq(rd, cin)
        if P["rq"]["bias"] != 0:
            bad()
    elif name == "lut16":
        P["e_in"], P["e_out"] = rd.i(1), rd.i(1)
        rd.zeros(2)
        P["T"] = rd.ints(65536, 2)
        if P["T"][0] != 0:
            bad()
    elif name in ("add", "concat", "chunk_add", "relu", "abs", "crop_to", "softmax", "sigmoid"):
        pass
    elif name == "split":
        P["start"], P["count"] = rd.u(4), rd.u(4)
        if P["count"] == 0:
            bad()
    elif name in ("pixel_shuffle", "pixel_unshuffle"):
        P["r"] = rd.u(4)
        if P["r"] == 0:
            bad()
    elif name == "clamp":
        P["lo"], P["hi"] = rd.i(8), rd.i(8)
        if P["lo"] > P["hi"]:
            bad()
    elif name == "leaky":
        P["a"], P["f"] = rd.i(8), rd.u(4)
        if P["f"] > 63:
            bad()
    elif name == "pad_rep":
        P["m"] = rd.u(4)
        if P["m"] == 0:
            bad()
    elif name == "crop":
        P["ops"] = [read_operand(rd, n_scalars) for _ in range(4)]
    elif name in ("subpos_split", "subpos_merge"):
        P["order"] = list(rd.take(4))
        if sorted(P["order"]) != [0, 1, 2, 3]:
            bad()
    elif name == "fill":
        P["kind"] = rd.u(1)
        rd.zeros(7)
        P["value"] = rd.i(8)
        if not (P["kind"] == 0 or (P["kind"] == 1 and 0 <= P["value"] < n_scalars)):
            bad()
    elif name == "igdn":
        P["gf"], P["frac"] = rd.u(4), rd.u(4)
        P["g_bytes"] = rd.u(1)
        rd.zeros(7)
        if P["gf"] > 127 or P["frac"] > 127 or P["g_bytes"] not in (4, 8):
            bad()
        P["gamma"] = rd.ints(cin * cin, P["g_bytes"])
        P["beta"] = rd.ints(cin, 8)
    elif name == "layernorm_c":
        P["eps"] = rd.u(8)
        P["g_bytes"] = rd.u(1)
        rd.zeros(7)
        if P["eps"] == 0 or P["g_bytes"] not in (4, 8):
            bad()
        P["gamma"] = rd.ints(cin, P["g_bytes"])
        P["beta"] = rd.ints(cin, 8)
    elif name == "l2norm":
        P["eps"] = rd.u(8)
        if P["eps"] == 0:
            bad()
    elif name == "elu":
        P["alpha"] = rd.i(8)
    elif name in ("gram", "apply"):
        P["groups"] = rd.u(4)
        P["rq"] = read_rq(rd, cin)
        if P["groups"] == 0 or P["rq"]["m_kind"] != 0 or P["rq"]["bias"] != 0 or P["rq"]["s_kind"] != 1:
            bad()
    if rd.p != len(pb):
        bad()
    return P


def parse_whg1(b):
    """Load steps 1-6 of section 10.5. Every refusal is model_mismatch."""
    if len(b) < 32 or b[0:4] != b"WHG1" or b[4] != 1 or b[5] != 1:
        bad()
    rd = Reader(b, 6)
    r0 = rd.u(2)
    n_t, n_o = rd.u(4), rd.u(4)
    n_i, n_out, n_s = rd.u(2), rd.u(2), rd.u(2)
    r1 = rd.u(2)
    max_elems = rd.u(4)
    r2 = rd.u(4)
    if r0 or r1 or r2:
        bad()
    if not (n_t and n_o and n_i and n_out) or not (1 <= max_elems <= CAP):
        bad()
    if 32 + 16 * n_s + 32 * n_i + 4 * n_out + 24 * n_t + 12 * n_o > len(b):
        bad()
    g = {"max_elems": max_elems, "scalars": [], "inputs": [], "outputs": [], "tensors": [], "ops": []}
    for _ in range(n_s):
        lo, hi = rd.i(8), rd.i(8)
        if lo > hi:
            bad()
        g["scalars"].append((lo, hi))
    for _ in range(n_i):
        t, mh, mw = rd.u(4), rd.u(4), rd.u(4)
        kind = rd.u(1)
        rd.zeros(3)
        src = rd.u(4)
        rd.zeros(4)
        init = rd.i(8)
        if (mh == 0) != (mw == 0) or kind > 1 or t >= n_t:
            bad()
        if kind == 0 and (src or init):
            bad()
        g["inputs"].append({"tensor": t, "max_h": mh, "max_w": mw, "kind": kind, "src": src, "init": init})
    for _ in range(n_out):
        t = rd.u(4)
        if t >= n_t:
            bad()
        g["outputs"].append(t)
    for _ in range(n_t):
        ch = rd.u(4)
        ex = rd.i(1)
        pol = rd.u(1)
        rd.zeros(2)
        lo, hi = rd.i(8), rd.i(8)
        if ch == 0 or pol > 2:
            bad()
        if pol == 0 and (lo or hi):
            bad()
        if pol != 0 and lo > hi:
            bad()
        g["tensors"].append({"c": ch, "e": ex, "policy": pol, "lo": lo, "hi": hi})
    for _ in range(n_o):
        code, ni, no = rd.u(1), rd.u(1), rd.u(1)
        rd.zeros(1)
        name = NAMES.get(code)
        if name is None:
            bad()
        ai, ao = arity(name)
        if no != ao or (ai is None and ni < 2) or (ai is not None and ni != ai):
            bad()
        ins = [rd.u(4) for _ in range(ni)]
        outs = [rd.u(4) for _ in range(no)]
        if any(t >= n_t for t in ins + outs):
            bad()
        plen = rd.u(4)
        pb = rd.take(plen)
        cin = g["tensors"][ins[0]]["c"]
        g["ops"].append({"name": name, "ins": ins, "outs": outs, "p": read_params(name, pb, cin, n_s)})
    if rd.p != len(b):
        bad()
    rules(g)
    return g


def rules(g):
    """Load step 6: the graph rules of section 10.3, state, exponents (3, 5.6)."""
    T = g["tensors"]
    defined = [False] * len(T)
    for i in g["inputs"]:
        if defined[i["tensor"]] or T[i["tensor"]]["policy"] == 0:
            bad()
        defined[i["tensor"]] = True
    if g["inputs"][0]["kind"] != 0:
        bad()
    if len({i["max_h"] != 0 for i in g["inputs"]}) != 1:
        bad()
    for op in g["ops"]:
        if not all(defined[t] for t in op["ins"]):
            bad()
        for t in op["outs"]:
            if defined[t]:
                bad()
            defined[t] = True
    if not all(defined):
        bad()
    for i in g["inputs"]:
        if i["kind"] == 1:
            if i["src"] >= len(g["outputs"]):
                bad()
            rec, out = T[i["tensor"]], T[g["outputs"][i["src"]]]
            if out["c"] != rec["c"] or not (rec["lo"] <= i["init"] <= rec["hi"]):
                bad()
    for op in g["ops"]:
        n, P = op["name"], op["p"]
        C = T[op["ins"][0]]["c"]
        cs = [T[t]["c"] for t in op["ins"]]
        want = C
        if n in ("conv", "conv_transpose"):
            want = P["cout"]
        elif n == "add":
            want = C if cs[1] == C else None
        elif n == "mul":
            want = C if cs[1] in (C, 1) else None
        elif n == "concat":
            want = sum(cs)
        elif n == "chunk_add":
            want = C // 2 if C % 2 == 0 else None
        elif n == "split":
            want = P["count"] if P["start"] + P["count"] <= C else None
        elif n == "pixel_shuffle":
            want = C // (P["r"] ** 2) if C % (P["r"] ** 2) == 0 else None
        elif n == "pixel_unshuffle":
            want = C * P["r"] ** 2
        elif n == "subpos_merge":
            want = C if all(c == C for c in cs) else None
        elif n == "fill":
            want = T[op["outs"][0]]["c"]
        elif n in ("gram", "apply"):
            want = C if cs[1] == C and C % P["groups"] == 0 else None
        elif n == "chscale":
            lo, hi = g["scalars"][P["scalar"]]
            want = C if lo >= 0 and hi <= P["n_rows"] - 1 else None
        if want is None or any(T[t]["c"] != want for t in op["outs"]):
            bad()
        if n in ("add", "concat", "chunk_add"):
            es = [T[t]["e"] for t in op["ins"] + op["outs"]]
            if not (all(e == -128 for e in es) or (all(e != -128 for e in es) and len(set(es)) == 1)):
                bad()
        if n == "lut16":
            if P["e_in"] != T[op["ins"][0]]["e"] or P["e_out"] != T[op["outs"][0]]["e"]:
                bad()


# --------------------------------------------------------------- plans (section 4, 10.5 step 7)


class Big(Exception):
    pass


class Shape(Exception):
    pass


def capped(v):
    if v > CAP:
        raise Big()
    return v


def op_sizes(g, op, hw, scal):
    """The output sizes of one operator, from its inputs' sizes. `scal` is None at load
    (scalars by their ranges) or the scalar values."""
    n, P = op["name"], op["p"]
    h, w = hw[op["ins"][0]]
    same = lambda o: hw[o] == (h, w)
    if n == "conv":
        hp, wp = capped(h + 2 * P["ph"]), capped(w + 2 * P["pw"])
        if hp < P["kh"] or wp < P["kw"]:
            raise Shape()
        return [((hp - P["kh"]) // P["sh"] + 1, (wp - P["kw"]) // P["sw"] + 1)]
    if n == "conv_transpose":
        uh = capped((h - 1) * P["sh"] + P["kh"] + P["oph"])
        uw = capped((w - 1) * P["sw"] + P["kw"] + P["opw"])
        capped(P["cout"] * uh * uw)
        if uh <= 2 * P["ph"] or uw <= 2 * P["pw"]:
            raise Shape()
        return [(uh - 2 * P["ph"], uw - 2 * P["pw"])]
    if n in ("add", "mul"):
        if not same(op["ins"][1]):
            raise Shape()
        return [(h, w)]
    if n in ("concat", "subpos_merge"):
        # subpos_merge's output sides are sizes it forms: capped before the shape rule.
        out = (capped(2 * h), capped(2 * w)) if n == "subpos_merge" else (h, w)
        if not all(same(t) for t in op["ins"]):
            raise Shape()
        return [out]
    if n == "pixel_shuffle":
        return [(capped(h * P["r"]), capped(w * P["r"]))]
    if n == "pixel_unshuffle":
        if h % P["r"] or w % P["r"]:
            raise Shape()
        return [(h // P["r"], w // P["r"])]
    if n == "pad_rep":
        m = P["m"]
        return [(capped(-(-h // m) * m), capped(-(-w // m) * m))]
    if n == "crop_to":
        rh, rw = hw[op["ins"][1]]
        if rh > h or rw > w:
            raise Shape()
        return [(rh, rw)]
    if n == "crop":
        y0, x0, ch, cw = P["ops"]

        def side(start, length, size):
            if scal is not None:
                val = lambda o: o[1] if o[0] == "const" else scal[o[1]]
                s0, ln = val(start), val(length)
                if s0 < 0 or ln < 1 or s0 + ln > size:
                    raise Shape()
                return ln
            s0 = start[1] if start[0] == "const" else 0
            ln = length[1] if length[0] == "const" else min(g["scalars"][length[1]][1], size - s0)
            if ln < 1 or s0 + ln > size:
                raise Shape()
            return ln

        return [(side(y0, ch, h), side(x0, cw, w))]
    if n == "subpos_split":
        if h % 2 or w % 2:
            raise Shape()
        return [(h // 2, w // 2)] * 4
    if n == "gram":
        if not same(op["ins"][1]):
            raise Shape()
        return [(1, g["tensors"][op["ins"][0]]["c"] // P["groups"])]
    if n == "apply":
        d = g["tensors"][op["ins"][0]]["c"] // P["groups"]
        if (h, w) != (1, d):
            raise Shape()
        return [hw[op["ins"][1]]]
    return [(h, w)]


def walk(g, hw, scal, load):
    for op in g["ops"]:
        try:
            outs = op_sizes(g, op, hw, scal)
        except Big:
            raise Refuse("too_large")
        except Shape:
            raise Refuse("model_mismatch" if load else "geometry_mismatch")
        for t, (oh, ow) in zip(op["outs"], outs):
            if g["tensors"][t]["c"] * oh * ow > g["max_elems"]:
                raise Refuse("too_large")
            hw[t] = (oh, ow)


def load(b):
    """Load steps 1-8 of section 10.5. Step 8 is the bounds pass and the verdict
    (section 8): a graph whose verdict meets neither profile is refused, so a graph
    that loads meets acc64 at least, and g["verdict"] says whether it meets acc32."""
    g = parse_whg1(b)
    if g["inputs"][0]["max_h"] == 0:
        if any(op["name"] in ("gram", "l2norm", "softmax") for op in g["ops"]):
            bad()
        g["sizes"] = None
    else:
        hw = {}
        for i in g["inputs"]:
            c = g["tensors"][i["tensor"]]["c"]
            if c * i["max_h"] * i["max_w"] > g["max_elems"]:
                raise Refuse("too_large")
            hw[i["tensor"]] = (i["max_h"], i["max_w"])
        walk(g, hw, None, True)
        g["sizes"] = hw
    g["verdict"] = bounds(g)
    if not g["verdict"]["acc64"]:
        bad()
    return g


# --------------------------------------------------------------- bounds (section 8)

SAT = 1 << 255
# Section 7.1: the width of each class under each profile, in bits.
WIDTH = {
    "acc64": {"tensor": 64, "acc": 64, "wide": 128, "product": 128},
    "acc32": {"tensor": 32, "acc": 32, "wide": 32, "product": 64},
}
# The class of each named intermediate (sections 5 and 2).
CLASS = {
    "conv.term": "acc", "conv.acc": "acc", "convT.term": "acc", "convT.acc": "acc",
    "rq.t": "acc", "rq.p": "product", "rq.q": "acc", "rq.y": "acc",
    "mul.p": "acc", "leaky.p": "acc",
    "gram.term": "acc", "gram.acc": "acc", "apply.term": "acc", "apply.acc": "acc",
    "igdn.x2": "wide", "igdn.gx2": "product", "igdn.term": "wide", "igdn.N": "wide",
    "igdn.r": "acc", "igdn.xr": "wide",
    "ln.S": "wide", "ln.T": "wide", "ln.m": "wide", "ln.v": "wide", "ln.dk": "wide", "ln.ng": "wide",
    "ln.d": "acc", "rsq.num": "acc", "ln.k": "acc", "ln.n": "acc",
    "l2.T": "wide", "l2.vk": "wide", "l2.k": "acc",
    "sm.diff": "acc", "sm.sum": "acc", "sm.num": "acc", "sig.num": "acc", "elu.p": "wide",
    "exp.negx": "acc", "exp.nl": "acc", "exp.t2": "acc", "exp.c": "acc",
}


def fits_bits(iv, n):
    return -(1 << (n - 1)) <= iv[0] and iv[1] <= (1 << (n - 1)) - 1


def hull(ivs):
    return (min(v[0] for v in ivs), max(v[1] for v in ivs))


def isqrt_ceil(n):
    r = math.isqrt(n)
    return r if r * r == n else r + 1


class Node:
    """One node of the bounds pass. Every +, -, x and floor division (a shift
    included) goes through sat (section 8.1), and the node is saturated when sat
    changes any result. min, max, |.|, negation and isqrt cannot leave
    [-2^255, 2^255] when their operands lie in it. `rec` holds each named
    intermediate's interval, hulled over channels."""

    def __init__(self):
        self.saturated = False
        self.rec = {}

    def sat(self, v):
        if v > SAT:
            self.saturated = True
            return SAT
        if v < -SAT:
            self.saturated = True
            return -SAT
        return v

    def add(self, a, b):
        return self.sat(a + b)

    def sub(self, a, b):
        return self.sat(a - b)

    def mul(self, a, b):
        return self.sat(a * b)

    def div(self, a, b):
        return self.sat(a // b)

    def shr(self, a, s):
        return self.sat(a >> s)

    def note(self, key, iv):
        old = self.rec.get(key)
        self.rec[key] = iv if old is None else (min(old[0], iv[0]), max(old[1], iv[1]))
        return iv

    def term(self, w, x):
        """The term interval of a weight w over x (section 8.2)."""
        a, b = self.mul(w, x[0]), self.mul(w, x[1])
        return (min(a, b), max(a, b))

    def corner(self, x, y):
        c = [self.mul(x[0], y[0]), self.mul(x[0], y[1]), self.mul(x[1], y[0]), self.mul(x[1], y[1])]
        return (min(c), max(c))

    def A(self, terms, b):
        """The partial-sum interval A(K, b) of section 8.2."""
        lo, hi = min(0, b), max(0, b)
        for t in terms:
            lo = self.add(lo, min(0, t[0]))
            hi = self.add(hi, max(0, t[1]))
        return (lo, hi)

    def R(self, p, s, mode):
        """R_floor, R_half_up, R_half_even (section 2), as written there."""
        if s == 0:
            return p
        if mode == 0:
            return self.shr(p, s)
        if mode == 1:
            return self.shr(self.add(p, 1 << (s - 1)), s)
        q = self.shr(p, s)
        r = self.sub(p, self.mul(q, 1 << s))
        h = 1 << (s - 1)
        return self.add(q, 1) if r > h or (r == h and q % 2 == 1) else q

    def rq(self, t, m, s, mode, post):
        """The rq.p, rq.q and rq.y rules of section 8.3 for one channel."""
        self.note("rq.t", t)
        P = self.term(m, t)
        a = (1 << (s - 1)) if mode != 0 and s >= 1 else 0
        self.note("rq.p", (P[0], self.add(P[1], a)))
        q = (self.R(P[0], s, mode), self.R(P[1], s, mode))
        self.note("rq.q", q)
        return self.note("rq.y", (self.add(q[0], post), self.add(q[1], post)))

    def exp(self, u):
        """The exp_neg_q16 rule of section 8.3, given the interval of u = -x."""
        self.note("exp.negx", u)
        self.note("exp.nl", (self.mul(45426, self.div(u[0], 45426)), self.mul(45426, self.div(u[1], 45426))))
        self.note("exp.t2", (self.mul(43245, 43245), self.mul(88670, 88670)))
        self.note("exp.c", (0, self.mul(23495, self.div(self.mul(88670, 88670), 1 << 16))))


def bpre(rq, o):
    return rq["biasv"][o] if rq["bias"] == 1 else 0


def bpost(rq, o):
    return rq["biasv"][o] if rq["bias"] == 2 else 0


def rq_mult(rq, o):
    m = rq_m(rq, o)
    return 1 if m is None else m


def node_bounds(nb, g, op, iv):
    """The producer results of one operator (section 8.3), recording its named
    intermediates in nb. Refuses (model_mismatch) when lut16's input range is not
    within [-32767, 32767] (section 5.6)."""
    T = g["tensors"]
    n, P = op["name"], op["p"]
    x = iv[op["ins"][0]]
    l, h = x
    C = T[op["ins"][0]]["c"]
    if n == "conv":
        rq, cg, kk = P["rq"], C // P["groups"], P["kh"] * P["kw"]
        outs = []
        for o in range(P["cout"]):
            taps = [nb.term(P["w"][(o * cg + ci) * kk + k], x) for ci in range(cg) for k in range(kk)]
            nb.note("conv.term", hull(taps))
            acc = nb.note("conv.acc", nb.A(taps, bpre(rq, o)))
            outs.append(nb.rq(acc, rq_mult(rq, o), rq_s(rq, o), rq["round"], bpost(rq, o)))
        return [hull(outs)]
    if n == "conv_transpose":
        rq, cout, kh, kw, sh, sw = P["rq"], P["cout"], P["kh"], P["kw"], P["sh"], P["sw"]
        outs = []
        for o in range(cout):
            every, by_res = [], {}
            for ci in range(C):
                for i in range(kh):
                    for j in range(kw):
                        t = nb.term(P["w"][((ci * cout + o) * kh + i) * kw + j], x)
                        every.append(t)
                        by_res.setdefault((i % sh, j % sw), []).append(t)
            nb.note("convT.term", hull(every))
            b = bpre(rq, o)
            accs = [nb.A(by_res[r], b) for r in sorted(by_res)]
            if len(by_res) < sh * sw:
                # A residue that no tap has: A over no terms.
                accs.append(nb.A([], b))
            acc = nb.note("convT.acc", hull(accs))
            outs.append(nb.rq(acc, rq_mult(rq, o), rq_s(rq, o), rq["round"], bpost(rq, o)))
        return [hull(outs)]
    if n == "rescale":
        rq = P["rq"]
        t = lambda c: (nb.add(l, bpre(rq, c)), nb.add(h, bpre(rq, c)))
        return [hull([nb.rq(t(c), rq_mult(rq, c), rq_s(rq, c), rq["round"], bpost(rq, c)) for c in range(C)])]
    if n == "chscale":
        lo_r, hi_r = g["scalars"][P["scalar"]]
        s, mode = P["shift"], P["round"]
        a = (1 << (s - 1)) if mode != 0 and s >= 1 else 0
        outs = []
        for c in range(C):
            ps, qs = [], []
            for r in range(lo_r, hi_r + 1):
                Pm = nb.term(P["M"][r * C + c], x)
                ps.append((Pm[0], nb.add(Pm[1], a)))
                qs.append((nb.R(Pm[0], s, mode), nb.R(Pm[1], s, mode)))
            nb.note("rq.t", x)
            nb.note("rq.p", hull(ps))
            q = nb.note("rq.q", hull(qs))
            outs.append(nb.note("rq.y", q))
        return [hull(outs)]
    if n == "mul":
        rq = P["rq"]
        mp = nb.note("mul.p", nb.corner(x, iv[op["ins"][1]]))
        return [hull([nb.rq(mp, rq_mult(rq, c), rq_s(rq, c), rq["round"], 0) for c in range(C)])]
    if n in ("gram", "apply"):
        rq = P["rq"]
        t = nb.note(n + ".term", nb.corner(x, iv[op["ins"][1]]))
        if n == "gram":
            L = max(g["sizes"][u][0] * g["sizes"][u][1] for u in op["ins"])
        else:
            L = C // P["groups"]
        acc = nb.note(n + ".acc", (nb.mul(L, min(0, t[0])), nb.mul(L, max(0, t[1]))))
        return [nb.rq(acc, 1, rq_s(rq, 0), rq["round"], 0)]
    if n == "add":
        b = iv[op["ins"][1]]
        return [(nb.add(l, b[0]), nb.add(h, b[1]))]
    if n == "chunk_add":
        return [(nb.mul(2, l), nb.mul(2, h))]
    if n == "clamp":
        return [(min(max(l, P["lo"]), P["hi"]), min(max(h, P["lo"]), P["hi"]))]
    if n == "relu":
        return [(max(l, 0), max(h, 0))]
    if n == "abs":
        return [(l, h) if l >= 0 else ((-h, -l) if h <= 0 else (0, max(-l, h)))]
    if n == "leaky":
        parts = []
        if l < 0:
            p = nb.note("leaky.p", nb.corner((l, min(h, -1)), (P["a"], P["a"])))
            parts.append((nb.shr(p[0], P["f"]), nb.shr(p[1], P["f"])))
        else:
            nb.note("leaky.p", (0, 0))
        if h >= 0:
            parts.append((max(l, 0), h))
        return [hull(parts)]
    if n == "lut16":
        if l < -32767 or h > 32767:
            bad()
        vals = P["T"][l + 32768:h + 32769]
        return [(min(vals), max(vals))]
    if n == "igdn":
        l2, h2 = nb.mul(l, l), nb.mul(h, h)
        x2 = nb.note("igdn.x2", (0, max(l2, h2)) if l <= 0 <= h else (min(l2, h2), max(l2, h2)))
        outs = []
        for i in range(C):
            terms = []
            for j in range(C):
                gm = P["gamma"][i * C + j]
                gx = nb.note("igdn.gx2", nb.corner((gm, gm), x2))
                terms.append(nb.note("igdn.term", (nb.shr(gx[0], P["gf"]), nb.shr(gx[1], P["gf"]))))
            nb.note("igdn.N", nb.A(terms, P["beta"][i]))
            # The final norm, a left fold over j from beta_i; not recorded.
            f_lo = f_hi = P["beta"][i]
            for t in terms:
                f_lo, f_hi = nb.add(f_lo, t[0]), nb.add(f_hi, t[1])
            r = nb.note("igdn.r", (math.isqrt(max(0, f_lo)), math.isqrt(max(0, f_hi))))
            xr = nb.note("igdn.xr", nb.corner(x, r))
            outs.append((nb.shr(xr[0], P["frac"]), nb.shr(xr[1], P["frac"])))
        return [hull(outs)]
    if n == "layernorm_c":
        e = P["eps"]
        r0 = math.isqrt(e)
        nb.note("ln.S", (nb.mul(C, min(0, l)), nb.mul(C, max(0, h))))
        half = C // 2
        nb.note("ln.m", (nb.sub(nb.mul(C, l), half), nb.add(nb.mul(C, h), half)))
        d = nb.sub(h, l)
        nb.note("ln.d", (nb.sub(l, h), d))
        dd = nb.mul(d, d)
        nb.note("ln.T", (0, nb.mul(C, dd)))
        nb.note("ln.v", (e, nb.add(dd, e)))
        nb.note("rsq.num", (1 << 32, 1 << 32))
        k = nb.div(1 << 32, r0)
        nb.note("ln.k", (0, k))
        ic = isqrt_ceil(C)
        big = nb.mul(nb.mul(ic, 1 << 32), nb.add(r0, 1))
        D = min(nb.mul(d, k), -nb.div(-big, r0))
        nb.note("ln.dk", (-D, D))
        nmax = nb.add(nb.div(nb.mul(nb.mul(ic, 1 << 16), nb.add(r0, 1)), r0), 1)
        N = min(nb.add(nb.shr(nb.mul(d, k), 16), 1), nmax)
        nb.note("ln.n", (-N, N))
        outs = []
        for c in range(C):
            gm = P["gamma"][c]
            ng = nb.note("ln.ng", nb.corner((-N, N), (gm, gm)))
            outs.append((nb.add(nb.shr(ng[0], 16), P["beta"][c]), nb.add(nb.shr(ng[1], 16), P["beta"][c])))
        return [hull(outs)]
    if n == "l2norm":
        e = P["eps"]
        hh, ww = g["sizes"][op["ins"][0]]
        X = max(abs(l), abs(h))
        nb.note("l2.T", (0, nb.add(nb.mul(nb.mul(hh, ww), nb.mul(X, X)), e)))
        nb.note("rsq.num", (1 << 32, 1 << 32))
        nb.note("l2.k", (0, nb.div(1 << 32, math.isqrt(e))))
        nb.note("l2.vk", (-(1 << 32), 1 << 32))
        return [(-(1 << 16), 1 << 16)]
    if n == "softmax":
        nb.note("sm.diff", (nb.sub(l, h), 0))
        nb.exp((0, nb.sub(h, l)))
        wm = g["sizes"][op["ins"][0]][1]
        nb.note("sm.sum", (0, nb.mul(wm, 65553)))
        nb.note("sm.num", (0, nb.add(nb.mul(65553, 1 << 16), nb.div(nb.mul(wm, 65553), 2))))
        return [(0, 1 << 16)]
    if n == "sigmoid":
        nb.exp((l, h) if l >= 0 else ((-h, -l) if h <= 0 else (0, max(-l, h))))
        nb.note("sig.num", (0, nb.mul(65553, 1 << 16)))
        return [(0, 1 << 16)]
    if n == "elu":
        nb.exp((max(0, -h), max(0, -l)))
        al = P["alpha"]
        p = nb.note("elu.p", nb.corner((al, al), (-(1 << 16), 17)))
        parts = []
        if h > 0:
            parts.append((max(l, 1), h))
        if l <= 0:
            parts.append((nb.shr(p[0], 16), nb.shr(p[1], 16)))
        return [hull(parts)]
    if n in ("concat", "subpos_merge"):
        return [hull([iv[t] for t in op["ins"]])]
    if n == "subpos_split":
        return [x] * 4
    if n == "fill":
        return [(P["value"], P["value"]) if P["kind"] == 0 else g["scalars"][P["value"]]]
    # split, pixel_shuffle, pixel_unshuffle, pad_rep, crop_to (its first input), crop.
    return [x]


def params32(op):
    """The parameter conditions of section 7.2: whether the node meets every one."""
    n, P = op["name"], op["p"]
    f32 = lambda v: I32_MIN <= v <= I32_MAX
    ok = True
    rq = P.get("rq")
    if rq is not None:
        ok = all(map(f32, rq["mult"])) and all(s <= 63 for s in rq["shift"]) and (rq["bias"] != 2 or all(map(f32, rq["biasv"])))
    if n in ("conv", "conv_transpose"):
        ok = ok and all(map(f32, P["w"]))
    elif n == "chscale":
        ok = all(map(f32, P["M"])) and P["shift"] <= 63
    elif n == "igdn":
        ok = all(map(f32, P["gamma"])) and P["gf"] <= 63 and P["frac"] <= 63
    elif n == "leaky":
        ok = f32(P["a"])
    return ok


def bounds(g):
    """The bounds pass (sections 8.1-8.4) and the verdict (section 8.5) of a graph
    that passed load steps 1-7: {"acc32": bool, "acc64": bool}. Refuses
    (model_mismatch) where the pass itself refuses the graph at load: a `check`
    interval that is empty (section 8.1), or lut16's input range (section 5.6)."""
    T = g["tensors"]
    ok = {"acc32": True, "acc64": True}
    iv = {}
    declared = [(T[i["tensor"]]["lo"], T[i["tensor"]]["hi"]) for i in g["inputs"]] + list(g["scalars"])
    for r in declared:
        for p in ok:
            ok[p] = ok[p] and fits_bits(r, WIDTH[p]["tensor"])
    for i in g["inputs"]:
        iv[i["tensor"]] = (T[i["tensor"]]["lo"], T[i["tensor"]]["hi"])
    for op in g["ops"]:
        nb = Node()
        results = node_bounds(nb, g, op, iv)
        fit = {p: not nb.saturated for p in ok}
        for p in ok:
            for k, r in nb.rec.items():
                fit[p] = fit[p] and fits_bits(r, WIDTH[p][CLASS[k]])
            for r in results:
                fit[p] = fit[p] and fits_bits(r, WIDTH[p]["tensor"])
        fit["acc32"] = fit["acc32"] and params32(op)
        for t, (a, b) in zip(op["outs"], results):
            rec = T[t]
            if rec["policy"] == 0:
                iv[t] = (a, b)
            elif rec["policy"] == 1:
                iv[t] = (min(max(a, rec["lo"]), rec["hi"]), min(max(b, rec["lo"]), rec["hi"]))
            else:
                fit["acc32"] = False
                iv[t] = (max(a, rec["lo"]), min(b, rec["hi"]))
                if iv[t][0] > iv[t][1]:
                    bad()
        for p in ok:
            ok[p] = ok[p] and fit[p]
    return ok


# --------------------------------------------------------------- the operators (section 5)


class T:
    """A tensor [c, h, w], elements in CHW order."""

    def __init__(self, c, h, w, d):
        self.c, self.h, self.w, self.d = c, h, w, d

    def at(self, c, y, x):
        return self.d[(c * self.h + y) * self.w + x]


def ev_conv(x, P, oh, ow):
    cin, H, W = x.c, x.h, x.w
    G = P["groups"]
    cg, og = cin // G, P["cout"] // G
    kh, kw = P["kh"], P["kw"]
    d = []
    for o in range(P["cout"]):
        base = (o // og) * cg
        for y in range(oh):
            for xx in range(ow):
                acc = 0
                for ci in range(cg):
                    for i in range(kh):
                        u = y * P["sh"] + i - P["ph"]
                        if not 0 <= u < H:
                            continue
                        for j in range(kw):
                            v = xx * P["sw"] + j - P["pw"]
                            if 0 <= v < W:
                                acc += P["w"][((o * cg + ci) * kh + i) * kw + j] * x.at(base + ci, u, v)
                d.append(rq_apply(P["rq"], o, acc))
    return T(P["cout"], oh, ow, d)


def ev_conv_t(x, P, oh, ow):
    """The gather form of section 5.3."""
    cin, H, W = x.c, x.h, x.w
    kh, kw, sh, sw, cout = P["kh"], P["kw"], P["sh"], P["sw"], P["cout"]
    d = []
    for o in range(cout):
        for y in range(oh):
            for xx in range(ow):
                acc = 0
                for i in range(kh):
                    ty = y + P["ph"] - i
                    if ty < 0 or ty % sh or ty // sh >= H:
                        continue
                    for j in range(kw):
                        tx = xx + P["pw"] - j
                        if tx < 0 or tx % sw or tx // sw >= W:
                            continue
                        for ci in range(cin):
                            acc += P["w"][((ci * cout + o) * kh + i) * kw + j] * x.at(ci, ty // sh, tx // sw)
                d.append(rq_apply(P["rq"], o, acc))
    return T(cout, oh, ow, d)


def per_channel(x, f):
    plane = x.h * x.w
    return T(x.c, x.h, x.w, [f(k // plane, v) for k, v in enumerate(x.d)])


def ev(g, op, vals, hw, scal):
    n, P = op["name"], op["p"]
    x = vals[op["ins"][0]]
    oh, ow = hw[op["outs"][0]]
    oc = [g["tensors"][t]["c"] for t in op["outs"]]
    if n == "conv":
        return [ev_conv(x, P, oh, ow)]
    if n == "conv_transpose":
        return [ev_conv_t(x, P, oh, ow)]
    if n == "rescale":
        return [per_channel(x, lambda c, v: rq_apply(P["rq"], c, v))]
    if n == "chscale":
        r = scal[P["scalar"]]
        return [per_channel(x, lambda c, v: R(v * P["M"][r * x.c + c], P["shift"], P["round"]))]
    if n == "mul":
        b = vals[op["ins"][1]]
        plane = x.h * x.w
        d = []
        for c in range(x.c):
            cb = c if b.c == x.c else 0
            for k in range(plane):
                d.append(rq_apply(P["rq"], c, x.d[c * plane + k] * b.d[cb * plane + k]))
        return [T(x.c, x.h, x.w, d)]
    if n == "lut16":
        return [T(x.c, x.h, x.w, [P["T"][v + 32768] if -32768 <= v <= 32767 else None for v in x.d])]
    if n == "add":
        b = vals[op["ins"][1]]
        return [T(x.c, x.h, x.w, [p + q for p, q in zip(x.d, b.d)])]
    if n == "concat":
        d = []
        for t in op["ins"]:
            d += vals[t].d
        return [T(oc[0], x.h, x.w, d)]
    if n == "chunk_add":
        half = (x.c // 2) * x.h * x.w
        return [T(x.c // 2, x.h, x.w, [x.d[k] + x.d[k + half] for k in range(half)])]
    if n == "split":
        plane = x.h * x.w
        return [T(P["count"], x.h, x.w, x.d[P["start"] * plane:(P["start"] + P["count"]) * plane])]
    if n == "pixel_shuffle":
        r = P["r"]
        co = x.c // (r * r)
        d = [0] * (co * oh * ow)
        for c in range(co):
            for i in range(r):
                for j in range(r):
                    for hh in range(x.h):
                        for ww in range(x.w):
                            d[(c * oh + hh * r + i) * ow + ww * r + j] = x.at(c * r * r + i * r + j, hh, ww)
        return [T(co, oh, ow, d)]
    if n == "pixel_unshuffle":
        r = P["r"]
        d = [0] * (x.c * r * r * oh * ow)
        for c in range(x.c):
            for i in range(r):
                for j in range(r):
                    for hh in range(oh):
                        for ww in range(ow):
                            d[((c * r * r + i * r + j) * oh + hh) * ow + ww] = x.at(c, hh * r + i, ww * r + j)
        return [T(x.c * r * r, oh, ow, d)]
    if n == "clamp":
        return [T(x.c, x.h, x.w, [min(max(v, P["lo"]), P["hi"]) for v in x.d])]
    if n == "relu":
        return [T(x.c, x.h, x.w, [max(v, 0) for v in x.d])]
    if n == "leaky":
        return [T(x.c, x.h, x.w, [v if v >= 0 else (v * P["a"]) >> P["f"] for v in x.d])]
    if n == "abs":
        return [T(x.c, x.h, x.w, [abs(v) for v in x.d])]
    if n == "pad_rep":
        return [T(x.c, oh, ow, [x.at(c, min(y, x.h - 1), min(xx, x.w - 1)) for c in range(x.c) for y in range(oh) for xx in range(ow)])]
    if n in ("crop_to", "crop"):
        y0 = x0 = 0
        if n == "crop":
            val = lambda o: o[1] if o[0] == "const" else scal[o[1]]
            y0, x0 = val(P["ops"][0]), val(P["ops"][1])
        return [T(x.c, oh, ow, [x.at(c, y0 + y, x0 + xx) for c in range(x.c) for y in range(oh) for xx in range(ow)])]
    if n == "subpos_split":
        outs = []
        for k in range(4):
            a, b2 = P["order"][k] >> 1, P["order"][k] & 1
            outs.append(T(x.c, oh, ow, [x.at(c, 2 * y + a, 2 * xx + b2) for c in range(x.c) for y in range(oh) for xx in range(ow)]))
        return outs
    if n == "subpos_merge":
        d = [0] * (x.c * oh * ow)
        for k in range(4):
            a, b2 = P["order"][k] >> 1, P["order"][k] & 1
            t = vals[op["ins"][k]]
            for c in range(x.c):
                for y in range(t.h):
                    for xx in range(t.w):
                        d[(c * oh + 2 * y + a) * ow + 2 * xx + b2] = t.at(c, y, xx)
        return [T(x.c, oh, ow, d)]
    if n == "fill":
        v = P["value"] if P["kind"] == 0 else scal[P["value"]]
        return [T(oc[0], x.h, x.w, [v] * (oc[0] * x.h * x.w))]
    if n == "igdn":
        C, plane = x.c, x.h * x.w
        out = [0] * (C * plane)
        for p in range(plane):
            xs = [x.d[j * plane + p] for j in range(C)]
            for i in range(C):
                N = P["beta"][i] + sum((P["gamma"][i * C + j] * xs[j] * xs[j]) // (1 << P["gf"]) for j in range(C))
                r = math.isqrt(max(0, N))
                out[i * plane + p] = (xs[i] * r) // (1 << P["frac"])
        return [T(C, x.h, x.w, out)]
    if n == "layernorm_c":
        C, plane = x.c, x.h * x.w
        out = [0] * (C * plane)
        for p in range(plane):
            xs = [x.d[c * plane + p] for c in range(C)]
            S = sum(xs)
            sg = (S > 0) - (S < 0)
            M = S + (C // 2) * sg
            m = trunc_div(M, C)
            ds = [v - m for v in xs]
            V = sum(d * d for d in ds) // C
            k = rsqrt_q16(V + P["eps"])
            for c in range(C):
                nn = (ds[c] * k) >> 16
                out[c * plane + p] = ((nn * P["gamma"][c]) >> 16) + P["beta"][c]
        return [T(C, x.h, x.w, out)]
    if n == "l2norm":
        plane = x.h * x.w
        out = []
        for c in range(x.c):
            vs = x.d[c * plane:(c + 1) * plane]
            k = rsqrt_q16(sum(v * v for v in vs) + P["eps"])
            out += [(v * k) >> 16 for v in vs]
        return [T(x.c, x.h, x.w, out)]
    if n == "softmax":
        out = []
        for row in range(x.c * x.h):
            vs = x.d[row * x.w:(row + 1) * x.w]
            m = max(vs)
            es = [exp_neg_q16(v - m) for v in vs]
            s = sum(es)
            out += [(e * 65536 + s // 2) // s for e in es]
        return [T(x.c, x.h, x.w, out)]
    if n == "sigmoid":
        def sig(v):
            if v >= 0:
                return (1 << 32) // (65536 + exp_neg_q16(-v))
            e = exp_neg_q16(v)
            return (e * 65536) // (65536 + e)
        return [T(x.c, x.h, x.w, [sig(v) for v in x.d])]
    if n == "elu":
        return [T(x.c, x.h, x.w, [v if v > 0 else (P["alpha"] * (exp_neg_q16(v) - 65536)) >> 16 for v in x.d])]
    if n == "gram":
        b = vals[op["ins"][1]]
        G = P["groups"]
        dd = x.c // G
        plane = x.h * x.w
        out = [0] * (x.c * dd)
        for gi in range(G):
            for i in range(dd):
                for j in range(dd):
                    s = sum(x.d[(gi * dd + i) * plane + p] * b.d[(gi * dd + j) * plane + p] for p in range(plane))
                    out[(gi * dd + i) * dd + j] = R(s, P["rq"]["shift"][0], P["rq"]["round"])
        return [T(x.c, 1, dd, out)]
    if n == "apply":
        V = vals[op["ins"][1]]
        G = P["groups"]
        dd = x.c // G
        plane = V.h * V.w
        out = [0] * (V.c * plane)
        for gi in range(G):
            for i in range(dd):
                for p in range(plane):
                    s = sum(x.d[(gi * dd + i) * dd + j] * V.d[(gi * dd + j) * plane + p] for j in range(dd))
                    out[(gi * dd + i) * plane + p] = R(s, P["rq"]["shift"][0], P["rq"]["round"])
        return [T(V.c, V.h, V.w, out)]
    raise Refuse("model_mismatch")


def apply_policy(rec, t):
    """Returns False when a `check` tensor holds an element outside its range."""
    if rec["policy"] == 1:
        t.d = [min(max(v, rec["lo"]), rec["hi"]) for v in t.d]
    elif rec["policy"] == 2:
        return all(rec["lo"] <= v <= rec["hi"] for v in t.d)
    return True


def run(g, data_inputs, scalars, state):
    """One decode (section 4): the data inputs in input-table order; each state input
    from `state` (section 10.3), which the decode updates when it completes."""
    allin = []
    di = 0
    si = 0
    if len(data_inputs) != sum(1 for i in g["inputs"] if i["kind"] == 0):
        raise Refuse("model_mismatch")
    for i in g["inputs"]:
        if i["kind"] == 0:
            allin.append(data_inputs[di])
            di += 1
        else:
            if state.get("s") is None:
                x0 = data_inputs[0]
                c = g["tensors"][i["tensor"]]["c"]
                allin.append(T(c, x0.h, x0.w, [i["init"]] * (c * x0.h * x0.w)))
            else:
                allin.append(state["s"][si])
            si += 1
    hw = {}
    # Step 1: each input's channel count, then its size.
    for i, t in zip(g["inputs"], allin):
        if t.c != g["tensors"][i["tensor"]]["c"] or t.h < 1 or t.w < 1:
            raise Refuse("model_mismatch")
        if i["max_h"] and (t.h > i["max_h"] or t.w > i["max_w"]):
            raise Refuse("too_large")
        if t.c * t.h * t.w > g["max_elems"]:
            raise Refuse("too_large")
        hw[i["tensor"]] = (t.h, t.w)
    # Step 2: each scalar against its range.
    if len(scalars) != len(g["scalars"]):
        raise Refuse("model_mismatch")
    for v, (lo, hi) in zip(scalars, g["scalars"]):
        if not lo <= v <= hi:
            raise Refuse("field_range")
    # Step 3: the run-time plan.
    walk(g, hw, scalars, False)
    # Step 4: evaluation, the inputs' policies first.
    vals = {}
    for i, t in zip(g["inputs"], allin):
        t = T(t.c, t.h, t.w, list(t.d))
        if not apply_policy(g["tensors"][i["tensor"]], t):
            raise Refuse("field_range")
        vals[i["tensor"]] = t
    for op in g["ops"]:
        outs = ev(g, op, vals, hw, scalars)
        for t, ten in zip(op["outs"], outs):
            # A value outside i64 is one load step 8 excludes; before the certifier
            # exists, such a decode is refused.
            if any(v is None or not I64_MIN <= v <= I64_MAX for v in ten.d):
                raise Refuse("model_mismatch")
            if not apply_policy(g["tensors"][t], ten):
                raise Refuse("field_range")
            vals[t] = ten
    outs = [vals[t] for t in g["outputs"]]
    state["s"] = [outs[i["src"]] for i in g["inputs"] if i["kind"] == 1]
    return outs


# --------------------------------------------------------------- canonical JSON (8.5, 14.1)


def cw(v):
    """The canonical form: ints as decimal strings, members sorted by key bytes."""
    if v is None:
        return "null"
    if v is True:
        return "true"
    if v is False:
        return "false"
    if isinstance(v, int):
        return '"%d"' % v
    if isinstance(v, str):
        return '"' + v.replace("\\", "\\\\").replace('"', '\\"') + '"'
    if isinstance(v, list):
        return "[" + ",".join(cw(x) for x in v) + "]"
    if isinstance(v, dict):
        keys = sorted(v, key=lambda k: k.encode())
        return "{" + ",".join(cw(k) + ":" + cw(v[k]) for k in keys) + "}"
    raise TypeError(v)


def dec(s):
    """A canonical decimal integer string."""
    if not isinstance(s, str) or not s or s == "-0":
        raise ValueError("not a canonical integer: %r" % (s,))
    d = s[1:] if s[0] == "-" else s
    if not d.isdigit() or (len(d) > 1 and d[0] == "0") or not d.isascii():
        raise ValueError("not a canonical integer: %r" % (s,))
    return int(s)


def tensor_in(v):
    c, h, w = (dec(x) for x in v["shape"])
    d = [dec(x) for x in v["data"]]
    if len(d) != c * h * w:
        raise ValueError("a tensor's data does not fill its shape")
    return T(c, h, w, d)


def tensor_out(t):
    return {"data": list(t.d), "shape": [t.c, t.h, t.w]}


def graph_line(v):
    """A graph vector, recomputed: `load` (steps 1-8), `profiles` (section 14.1: both
    when the load is refused, else exactly the profiles the verdict meets) and every
    decode's `expect`."""
    b = bytes.fromhex(v["graph"])
    out = {"graph": b.hex(), "id": v["id"], "decodes": []}
    try:
        g = load(b)
        out["load"] = "ok"
        out["profiles"] = [p for p in ("acc32", "acc64") if g["verdict"][p]]
    except Refuse as e:
        out["load"] = e.code
        out["profiles"] = ["acc32", "acc64"]
        g = None
    if g is not None:
        state = {}
        for d in v["decodes"]:
            ins = [tensor_in(t) for t in d["inputs"]]
            sc = [dec(s) for s in d["scalars"]]
            try:
                res = run(g, ins, sc, state)
                exp = {"outputs": [tensor_out(t) for t in res]}
            except Refuse as e:
                exp = {"refuse": e.code}
            out["decodes"].append({"expect": exp, "inputs": [tensor_out(t) for t in ins], "scalars": sc})
    return out


# --------------------------------------------------------------- stage operators (5E)


class Rans:
    """integer-payloads section 2's decoder state."""

    def __init__(self, b):
        self.b = b
        self.x = (int.from_bytes(b[4:8], "little") << 32) | int.from_bytes(b[0:4], "little")
        self.i = 2
        self.out = False

    def renorm(self):
        if self.x < (1 << 31):
            if 4 * self.i + 4 <= len(self.b):
                self.x = (self.x << 32) | int.from_bytes(self.b[4 * self.i:4 * self.i + 4], "little")
                self.i += 1
            else:
                self.out = True

    def bits(self, n):
        v = self.x & ((1 << n) - 1)
        self.x >>= n
        self.renorm()
        return v


def sat32(v):
    return max(I32_MIN, min(I32_MAX, v))


def rans64_decode(i, m):
    table_rule, escape, end = m["table"], m["escape"], m["end"]
    lo = hi = None
    if escape == "strict":
        lo, hi = dec(m["lo"]), dec(m["hi"])
        if not (I32_MIN <= lo <= hi <= I32_MAX):
            raise Refuse("model_mismatch", -1)
    tables = [([dec(c) for c in t["cdf"]], dec(t["offset"])) for t in i["tables"]]
    reads = [dec(r) for r in i["reads"]]
    for cdf, o in tables:
        n = len(cdf)
        ok_n = 3 <= n <= 4097 if table_rule == "t4097" else 2 <= n <= 65537
        if not ok_n or cdf[0] != 0 or cdf[-1] != 65536 or any(cdf[k] >= cdf[k + 1] for k in range(n - 1)) or not I32_MIN <= o <= I32_MAX:
            raise Refuse("model_mismatch", -1)
    if any(not 0 <= r < len(tables) for r in reads):
        raise Refuse("model_mismatch", -1)
    stream = bytes.fromhex(i["stream"])
    if len(stream) < 8 or len(stream) % 4:
        raise Refuse("stream_exhausted", 0)
    d = Rans(stream)
    syms = []
    for k, r in enumerate(reads):
        cdf, o = tables[r]
        n = len(cdf)
        cum = d.x & 0xFFFF
        s = max(j for j in range(n - 1) if cdf[j] <= cum)
        d.x = (cdf[s + 1] - cdf[s]) * (d.x >> 16) + cum - cdf[s]
        d.renorm()
        code = None
        if s < n - 2:
            sym = s + o
            if escape == "lenient":
                sym = sat32(sym)
            elif not lo <= sym <= hi:
                code = "field_range"
        elif escape == "lenient":
            v = d.bits(4)
            count = v
            groups = 0
            while v == 15 and groups < 8 and not d.out:
                v = d.bits(4)
                count += v
                groups += 1
            count = min(count, 8)
            raw = 0
            for j in range(count):
                raw |= d.bits(4) << (4 * j)
            value = -(raw >> 1) - 1 if raw & 1 else (raw >> 1) + (n - 2)
            sym = sat32(sat32(value) + o)
        else:
            count = d.bits(4)
            raw, last = 0, None
            if count <= 8:
                for j in range(count):
                    last = d.bits(4)
                    raw += last << (4 * j)
            value = -(raw // 2) - 1 if raw % 2 else raw // 2 + (n - 2)
            sym = value + o
            if count > 8 or (count >= 1 and last == 0):
                code = "stream_noncanonical"
            elif not lo <= sym <= hi:
                code = "field_range"
        if end == "canonical" and d.out:
            raise Refuse("stream_exhausted", k)
        if code:
            raise Refuse(code, k)
        syms.append(sym)
    if end == "canonical":
        if d.out or d.x != (1 << 31) or 4 * d.i != len(stream):
            raise Refuse("stream_noncanonical", len(reads))
    elif d.out:
        raise Refuse("stream_exhausted", len(reads))
    return {"symbols": syms}


class Win:
    """E8's window and states over memory M (8 zero bytes, then the stream)."""

    def __init__(self, stream):
        self.M = bytes(8) + stream
        self.p = len(stream)
        self.h = self.le64(self.p)
        self.pos = 0
        self.last_in = 0
        if self.h == 0:
            raise Refuse("stream_exhausted", 0)
        self.b = self.h.bit_length() - 1
        self.h -= 1 << self.b
        self.s = [0, 0]
        try:
            self.s[0] = self.read(8)
            self.s[1] = self.read(8)
            self.flush()
        except Refuse:
            raise Refuse("stream_exhausted", 0)

    def le64(self, i):
        return int.from_bytes(self.M[i:i + 8], "little")

    def read(self, k):
        if k > self.b:
            raise Refuse("stream_exhausted", self.pos)
        self.b -= k
        r = self.h >> self.b
        self.h -= r << self.b
        return r

    def flush(self):
        d = 7 ^ (self.b // 8)
        if d > self.p:
            raise Refuse("stream_exhausted", self.pos)
        self.p -= d
        self.h = self.le64(self.p) % (1 << (56 + self.b % 8))
        self.b = 56 + self.b % 8

    def step(self, which, t):
        self.s[which] = ((t >> 16) | self.read(t >> 24)) % 256

    def at(self, k):
        """The `in` or `out` step of position k runs: a refusal in it is at k."""
        self.pos = k

    def at_in(self, k):
        self.pos = k
        self.last_in = k

    def group_flush(self):
        """A flush that belongs to no single position (section 14.1, `at`): it reports
        the last position whose `in` step ran."""
        self.pos = self.last_in
        self.flush()


def wrap16(a):
    return ((a + 32768) % 65536) - 32768


def latent_tables(pmf):
    w = len(pmf[0]) if pmf else 0
    if w < 1 or any(len(r) != w or min(r) < 0 or sum(r) != 256 for r in pmf):
        raise Refuse("model_mismatch", -1)
    tabs = []
    for P in pmf:
        C = [0]
        for p in P:
            C.append(C[-1] + p)
        F = lambda k: C[k] - C[k] // 2
        G = lambda k: C[k] // 2 + 128
        sym = lambda k: k // 2 if k % 2 == 0 else -((k + 1) // 2)
        slot_sym, slot_next = [None] * 256, [None] * 256
        k = 0
        while C[k] < 256:
            for u in range(F(k), F(k + 1)):
                slot_sym[u], slot_next[u] = sym(k), P[k] + (u - F(k))
            for u in range(G(k), G(k + 1)):
                slot_sym[u], slot_next[u] = sym(k), P[k] + (F(k + 1) - F(k)) + (u - G(k))
            k += 1
        t = []
        for u in range(256):
            nb = 0
            while slot_next[u] * (1 << nb) < 256:
                nb += 1
            base = (slot_next[u] * (1 << nb)) ^ 256
            t.append(nb * (1 << 24) + base * (1 << 16) + slot_sym[u] % 65536)
        tabs.append(t)
    return tabs


def mans_sgm(i):
    pmf = [[dec(x) for x in r] for r in i["pmf"]]
    bounds = [dec(x) for x in i["bounds"]]
    spans = [([bool(b) for b in s["mask"]], [dec(x) for x in s["sigma"]]) for s in i["spans"]]
    tabs = latent_tables(pmf)
    if any(not 0 <= b <= 255 for b in bounds):
        raise Refuse("model_mismatch", -1)
    nd = min(len(tabs), len(bounds))
    if any(len(m) != len(sg) or any(not 0 <= x < nd for x in sg) for m, sg in spans):
        raise Refuse("model_mismatch", -1)
    w = Win(bytes.fromhex(i["memory"]))
    out = []
    base = 0
    for mask, sigma in spans:
        L = len(mask)
        sy = [0] * L

        def inb(which, k):
            w.at_in(base + k)
            if mask[k]:
                t = tabs[sigma[k]][w.s[which]]
                w.step(which, t)
                sy[k] = wrap16(t % 65536)

        def outb(k):
            w.at(base + k)
            if mask[k]:
                beta = bounds[sigma[k]]
                if sy[k] + beta == 0:
                    f = w.read(1)
                    v = w.read(3 if f == 1 else 16)
                    sy[k] = wrap16(v // 2 + beta) if v % 2 == 0 else wrap16(-(v // 2) - beta)
                    w.flush()

        k = 0
        while k + 3 < L:
            inb(0, k); inb(1, k + 1); inb(0, k + 2); inb(1, k + 3)
            w.group_flush()
            for q in range(4):
                outb(k + q)
            k += 4
        if L % 4 >= 2:
            inb(0, k); outb(k); inb(1, k + 1); outb(k + 1)
            k += 2
            w.group_flush()
        if L % 2 == 1:
            inb(0, k); outb(k)
            w.group_flush()
        out.append(sy)
        base += L
    return {"symbols": out}


def mans_fac(i):
    rows = [[dec(x) for x in r] for r in i["rows"]]
    lens = [dec(x) for x in i["lengths"]]
    if len(rows) != len(lens) or any(len(r) != 63 or r[0] < 0 or r[62] > 255 or any(r[k] > r[k + 1] for k in range(62)) for r in rows):
        raise Refuse("model_mismatch", -1)
    Zp = [0] * 512
    nb = 0
    for i_ in range(511, 0, -1):
        if i_ * (1 << nb) < 256:
            nb += 1
        Zp[i_] = ((i_ * (1 << nb)) ^ 256) + 256 * nb
    w = Win(bytes.fromhex(i["memory"]))
    Tz = [0] * 256
    out = []
    base = 0
    for r, L in zip(rows, lens):
        prev = 0
        for a in range(63):
            k = r[a] - prev
            for j in range(prev, r[a]):
                Tz[j] = Zp[k + j - prev] * 65536 + a
            prev = r[a]
        Tz[255] = Zp[1] * 65536 + 63
        sy = [0] * L

        def inb(which, k):
            w.at_in(base + k)
            t = Tz[w.s[which]]
            w.step(which, t)
            sy[k] = t % 256

        def outb(k):
            w.at(base + k)
            if sy[k] == 63:
                sy[k] = w.read(6)

        k = 0
        while k + 3 < L:
            inb(0, k); inb(1, k + 1); inb(0, k + 2); inb(1, k + 3)
            for q in range(4):
                outb(k + q)
            w.group_flush()
            k += 4
        if L % 4 >= 2:
            inb(0, k); outb(k); inb(1, k + 1); outb(k + 1)
            k += 2
        if L % 2 == 1:
            inb(0, k); outb(k)
        w.group_flush()
        out.append(sy)
        base += L
    return {"symbols": out}


def rdiv(n, d):
    if d < 1:
        return 0
    sg = (n > 0) - (n < 0)
    return sg * ((2 * abs(n) + d) // (2 * d))


def stage_expect(op, i, m):
    D = lambda k: dec(i[k])
    if op == "rans64_decode":
        return rans64_decode(i, m)
    if op == "dequantize":
        f, s, b = D("f"), D("sym"), D("base")
        if not 0 <= f <= 62:
            raise Refuse("field_range", 0)
        v = s * (1 << f) + b
        if not I64_MIN <= v <= I64_MAX:
            raise Refuse("field_range", 0)
        return {"value": v}
    if op == "bucket":
        sigma, e = D("sigma"), [dec(x) for x in i["edges"]]
        if m["bucket"] == "ceil":
            if not e:
                raise Refuse("model_mismatch", -1)
            s2 = max(sigma, e[0])
            return {"index": (len(e) - 1) - sum(1 for k in range(len(e) - 1) if s2 <= e[k])}
        if not e or any(e[k] > e[k + 1] for k in range(len(e) - 1)):
            raise Refuse("model_mismatch", -1)
        return {"index": sum(1 for x in e if x <= sigma)}
    if op == "coded":
        return {"coded": D("v") > D("thr")}
    if op == "scale_index":
        n, s = D("n"), D("s")
        if n < 1 or not 0 <= s <= 127:
            raise Refuse("model_mismatch", -1)
        return {"index": min(max(R(D("v"), s, 1), 0), n - 1)}
    if op == "block_stat":
        x = tensor_in(i["x"])
        b, s, lo, hi = D("b"), D("s"), D("lo"), D("hi")
        if b < 1 or not 0 <= s <= 127 or lo > hi:
            raise Refuse("model_mismatch", -1)
        d = list(x.d)
        for c in range(x.c):
            for by in range(0, x.h, b):
                for bx in range(0, x.w, b):
                    cells = [(y, xx) for y in range(by, min(by + b, x.h)) for xx in range(bx, min(bx + b, x.w))]
                    v = min(max(R(sum(x.at(c, y, xx) for y, xx in cells), s, 1), lo), hi)
                    for y, xx in cells:
                        d[(c * x.h + y) * x.w + xx] = v
        return {"out": tensor_out(T(x.c, x.h, x.w, d))}
    if op == "band_lookup":
        e = [dec(x) for x in i["edges"]]
        table = [[dec(x) for x in r] for r in i["table"]]
        row = D("row")
        if not 0 <= row < len(table) or len(table[row]) != len(e) + 1:
            raise Refuse("model_mismatch", -1)
        return {"value": table[row][sum(1 for k in e if D("v") >= k)]}
    if op == "mans_decode":
        return mans_sgm(i) if i["mode"] == "sgm" else mans_fac(i)
    if op == "residual_step":
        pred, r, q, lo, hi = D("pred"), D("r"), D("q"), D("lo"), D("hi")
        if lo >= hi or not lo <= pred <= hi or not 1 <= q <= 65535 or not I32_MIN <= r <= I32_MAX:
            raise Refuse("model_mismatch", -1)
        return {"value": min(max(pred + r * q, lo), hi)}
    if op == "mc_predict":
        f = tensor_in(i["frame"])
        y, x, c, dy, dx, pr = D("y"), D("x"), D("c"), D("dy"), D("dx"), D("precision")
        if not (0 <= c < f.c and 0 <= y < f.h and 0 <= x < f.w) or pr not in (1, 2):
            raise Refuse("model_mismatch", -1)
        p = lambda Y, X: f.at(c, min(max(Y, 0), f.h - 1), min(max(X, 0), f.w - 1))
        if pr == 1:
            return {"value": p(y + dy, x + dx)}
        iy, fy, ix, fx = y + (dy >> 1), dy & 1, x + (dx >> 1), dx & 1
        if not fy and not fx:
            v = p(iy, ix)
        elif not fy:
            v = (p(iy, ix) + p(iy, ix + 1) + 1) >> 1
        elif not fx:
            v = (p(iy, ix) + p(iy + 1, ix) + 1) >> 1
        else:
            v = (p(iy, ix) + p(iy, ix + 1) + p(iy + 1, ix) + p(iy + 1, ix + 1) + 2) >> 2
        return {"value": v}
    if op == "up2_predict":
        prev = tensor_in(i["prev"])
        d0, d1 = D("d0"), D("d1")
        hp, wp = -(-d0 // 2), -(-d1 // 2)
        if d0 < 1 or d1 < 1 or (prev.h, prev.w) != (hp, wp):
            raise Refuse("geometry_mismatch", -1)
        p = lambda c, Y, X: prev.at(c, min(Y, hp - 1), min(X, wp - 1))
        d = []
        for c in range(prev.c):
            for y in range(d0):
                for x in range(d1):
                    iy, fy, ix, fx = y >> 1, y & 1, x >> 1, x & 1
                    if not fy and not fx:
                        v = p(c, iy, ix)
                    elif not fy:
                        v = (p(c, iy, ix) + p(c, iy, ix + 1) + 1) >> 1
                    elif not fx:
                        v = (p(c, iy, ix) + p(c, iy + 1, ix) + 1) >> 1
                    else:
                        v = (p(c, iy, ix) + p(c, iy, ix + 1) + p(c, iy + 1, ix) + p(c, iy + 1, ix + 1) + 2) >> 2
                    d.append(v)
        return {"out": tensor_out(T(prev.c, d0, d1, d))}
    if op == "rate_predict":
        prev = tensor_in(i["prev"])
        num, den, d0 = D("rate_num"), D("rate_den"), D("d0")
        n = prev.w
        if prev.h != 1 or num < 1 or den < 1 or n * num // den != d0:
            raise Refuse("geometry_mismatch", -1)
        d = []
        for c in range(prev.c):
            for k in range(d0):
                s = k * den
                lo, rho = s // num, s % num
                hi = min(lo + 1, n - 1)
                a, b = prev.at(c, 0, lo), prev.at(c, 0, hi)
                d.append(rdiv(a * num + (b - a) * rho, num))
        return {"out": tensor_out(T(prev.c, 1, d0, d))}
    if op == "coding_order":
        c, h, w = (dec(x) for x in i["shape"])
        cb = None if i["order"] == "raster" else i["anchor"]
        idx = [(cc * h + y) * w + x for cc in range(c) for y in range(h) for x in range(w) if cb is None or (((y + x) % 2 == 0) == cb)]
        return {"indices": idx}
    raise ValueError("unknown stage operator %r" % op)


def stage_line(v):
    try:
        exp = stage_expect(v["op"], v["in"], v["modes"])
    except Refuse as e:
        exp = {"at": e.at, "refuse": e.code}
    return {"expect": exp, "id": v["id"], "in": v["in"], "modes": v["modes"], "op": v["op"]}


# --------------------------------------------------------------- lowerings (section 13)


class WB:
    """A WHG1 graph under construction with section 13.0's numbering."""

    def __init__(self):
        self.scalars, self.inputs, self.tensors, self.ops = [], [], [], []

    def tensor(self, c, policy=0, lo=0, hi=0):
        self.tensors.append((c, policy, lo, hi))
        return len(self.tensors) - 1

    def input(self, c, lo, hi, mh=0, mw=0):
        t = self.tensor(c, 2, lo, hi)
        self.inputs.append((t, mh, mw))
        return t

    def op(self, code, ins, out_rec, params):
        t = self.tensor(*out_rec)
        self.ops.append((code, ins, [t], params))
        return t

    def bytes(self, outputs):
        b = bytearray(b"WHG1")
        b += bytes([1, 1, 0, 0])
        b += struct.pack("<IIHHHHII", len(self.tensors), len(self.ops), len(self.inputs), len(outputs), len(self.scalars), 0, CAP, 0)
        for lo, hi in self.scalars:
            b += struct.pack("<qq", lo, hi)
        for t, mh, mw in self.inputs:
            b += struct.pack("<IIIB3xI4xq", t, mh, mw, 0, 0, 0)
        for t in outputs:
            b += struct.pack("<I", t)
        for c, pol, lo, hi in self.tensors:
            b += struct.pack("<Ibb2xqq", c, -128, pol, lo, hi)
        for code, ins, outs, params in self.ops:
            b += struct.pack("<BBBx", code, len(ins), len(outs))
            for t in ins + outs:
                b += struct.pack("<I", t)
            b += struct.pack("<I", len(params)) + params
        return bytes(b)


def ints_le(v, width):
    return b"".join(x.to_bytes(width, "little", signed=True) for x in v)


def rq_block(round_, bias, mult, mkind, mbytes, shifts, skind, biasv, bbytes):
    """The RQ block of section 10.4."""
    b = struct.pack("<BBBBBBH", round_, bias, mkind, skind, mbytes, bbytes, 0)
    b += ints_le(mult, mbytes) if mkind else b""
    b += bytes(shifts)
    b += ints_le(biasv, bbytes) if bias else b""
    return b


def smallest(v, allowed):
    for w in allowed:
        lo, hi = -(1 << (8 * w - 1)), (1 << (8 * w - 1)) - 1
        if all(lo <= x <= hi for x in v):
            return w
    return 8


class LowerRefuse(Exception):
    pass


def lref():
    raise LowerRefuse()


def p_conv(cout, kh, kw, sh, sw, ph, pw, groups, w_bytes, rq, w):
    return struct.pack("<8IB3x", cout, kh, kw, sh, sw, ph, pw, groups, w_bytes) + rq + ints_le(w, w_bytes)


def p_convt(cout, kh, kw, sh, sw, ph, pw, oph, opw, w_bytes, rq, w):
    return struct.pack("<9IB3x", cout, kh, kw, sh, sw, ph, pw, oph, opw, w_bytes) + rq + ints_le(w, w_bytes)


def p_clamp(lo, hi):
    return struct.pack("<qq", lo, hi)


def int_synth_lower(js):
    """int_synth/1 (section 13.1)."""
    try:
        root = json.loads(js)
    except ValueError:
        lref()
    layers = root.get("layers") if isinstance(root, dict) else None
    if not isinstance(layers, list) or not layers:
        lref()

    def u(l, k, default=None):
        v = l.get(k, default)
        if v is None or isinstance(v, bool) or not isinstance(v, int) or not 0 <= v <= 0xFFFFFFFF:
            lref()
        return v

    def i32s(l, k):
        a = l.get(k)
        if not isinstance(a, list) or any(isinstance(x, bool) or not isinstance(x, int) or not I32_MIN <= x <= I32_MAX for x in a):
            lref()
        return a

    if not isinstance(layers[0], dict):
        lref()
    c = u(layers[0], "in_ch")
    if c == 0:
        lref()
    wb = WB()
    x = wb.input(c, -128, 127)
    for l in layers:
        if not isinstance(l, dict):
            lref()
        ic, oc, k, s, p = u(l, "in_ch"), u(l, "out_ch"), u(l, "k"), u(l, "stride"), u(l, "pad")
        op_, sh = u(l, "out_pad", 0), u(l, "shift", 0)
        relu = l.get("relu", False)
        if not isinstance(relu, bool):
            lref()
        w, b = i32s(l, "weight"), i32s(l, "bias")
        if ic != c or oc == 0 or k == 0 or s == 0 or sh > 127 or len(w) != ic * oc * k * k or len(b) != oc:
            lref()
        rq = rq_block(0, 1, [], 0, 0, [sh], 1, b, 4)
        x = wb.op(0x11, [x], (oc,), p_convt(oc, k, k, s, s, p, p, op_, op_, 4, rq, w))
        if relu:
            x = wb.op(0x0C, [x], (oc,), b"")
        c = oc
    x = wb.op(0x0B, [x], (c,), p_clamp(0, 255))
    return [("synth", wb.bytes([x]))]


def parse_net(b, magic, audio):
    """WHM1 / WAM1 (integer-payloads section 9)."""
    if b[:4] != magic:
        lref()
    rd = Reader(b, 4)
    try:
        afrac, gf, shift, n = rd.u(4), rd.u(4), rd.u(4), rd.u(4)
        layers = []
        for _ in range(n):
            tag = rd.u(1)
            if tag == 0:
                inc, out, k, s, p, op_ = (rd.u(4) for _ in range(6))
                w = rd.ints(inc * out * k * (1 if audio else k), 2)
                layers.append(("convT", inc, out, k, s, p, op_, w, rd.ints(out, 8), rd.ints(out, 8)))
            elif tag == 1 and not audio:
                inc, out, k, s, p = (rd.u(4) for _ in range(5))
                w = rd.ints(out * inc * k * k, 2)
                layers.append(("conv", inc, out, k, s, p, 0, w, rd.ints(out, 8), rd.ints(out, 8)))
            elif tag == 2:
                c = rd.u(4)
                layers.append(("igdn", c, rd.ints(c * c, 4), rd.ints(c, 8)))
            elif tag == 3 and not audio:
                layers.append(("relu",))
            elif tag == 4 and not audio:
                layers.append(("leaky", rd.i(8), rd.u(4)))
            else:
                lref()
    except Refuse:
        lref()
    if rd.p != len(b):
        lref()
    return {"afrac": afrac, "gf": gf, "shift": shift, "layers": layers, "audio": audio}


def largest_x(pred):
    """The largest X in [0, 2^63] with pred(X), pred monotone; None if pred(0) fails."""
    if not pred(0):
        return None
    lo, hi = 0, 1 << 63
    if pred(hi):
        return hi
    while hi - lo > 1:
        mid = (lo + hi) // 2
        if pred(mid):
            lo = mid
        else:
            hi = mid
    return lo


def x_l(net, layer, c):
    """X_l of section 13.2 for a layer of tag 0, 1, 2 or 4 reading c channels; the set
    is refused when P_l(0) fails."""
    I64, I128 = (1 << 63) - 1, (1 << 127) - 1
    kind = layer[0]
    if kind in ("conv", "convT"):
        k = layer[3]
        terms = c * k * (1 if net["audio"] else k)
        Wm, Mm, Bm = max(map(abs, layer[7]), default=0), max(map(abs, layer[8]), default=0), max(map(abs, layer[9]), default=0)
        p = lambda X: terms * Wm * X <= I64 and (terms * Wm * X * Mm) // (1 << net["shift"]) + 1 + Bm <= I64
    elif kind == "igdn":
        G, Bb = max(map(abs, layer[2]), default=0), max(map(abs, layer[3]), default=0)

        def p(X):
            if G * X * X > I128:
                return False
            N = Bb + c * ((G * X * X) // (1 << net["gf"]) + 1)
            if N > I128:
                return False
            r = math.isqrt(N)
            return r <= I64 and (X * r) // (1 << net["afrac"]) + 1 <= I64
    else:
        a = abs(layer[1])
        p = lambda X: X * a <= I64
    x = largest_x(p)
    if x is None:
        lref()
    return x


def x_record(net, k, c):
    """The tensor record of x_k on c channels (section 13.2, Range checks). x_k is read
    by layer k, if there is one. When that layer is of tag 0, 1, 2 or 4, x_k is
    check [-X_k, min(X_k, I64)]; that includes the network's input. A network input
    that no such layer reads is check [-2^63, I64]. Every other tensor is derived."""
    L = net["layers"]
    if k < len(L) and L[k][0] in ("convT", "conv", "igdn", "leaky"):
        x = x_l(net, L[k], c)
        return (c, 2, -x, min(x, I64_MAX))
    if k == 0:
        return (c, 2, I64_MIN, I64_MAX)
    return (c,)


def net_layers(wb, net, x):
    L = net["layers"]
    for k, l in enumerate(L):
        c = wb.tensors[x][0]
        oc = l[2] if l[0] in ("conv", "convT") else c
        rec = x_record(net, k + 1, oc)
        if l[0] in ("conv", "convT"):
            _, inc, out, kk, s, p, op_, w, mult, bias = l
            if inc != c or out == 0 or kk == 0 or s == 0:
                lref()
            rq = rq_block(0, 2, mult, 2, 8, [net["shift"]], 1, bias, 8)
            if l[0] == "conv":
                x = wb.op(0x01, [x], rec, p_conv(out, kk, kk, s, s, p, p, 1, 2, rq, w))
            elif net["audio"]:
                x = wb.op(0x11, [x], rec, p_convt(out, 1, kk, 1, s, 0, p, 0, op_, 2, rq, w))
            else:
                x = wb.op(0x11, [x], rec, p_convt(out, kk, kk, s, s, p, p, op_, op_, 2, rq, w))
        elif l[0] == "igdn":
            if l[1] != c:
                lref()
            x = wb.op(0x12, [x], rec, struct.pack("<IIB7x", net["gf"], net["afrac"], 4) + ints_le(l[2], 4) + ints_le(l[3], 8))
        elif l[0] == "relu":
            x = wb.op(0x0C, [x], rec, b"")
        else:
            if l[2] > 63:
                lref()
            x = wb.op(0x0D, [x], rec, struct.pack("<qI", l[1], l[2]))
    return x


def net_check(net):
    if net["shift"] > 127 or net["gf"] > 127 or not 1 <= net["afrac"] <= 30 or not net["layers"]:
        lref()
    for l in net["layers"]:
        if l[0] in ("conv", "convT"):
            return l[1]
        if l[0] == "igdn":
            return l[1]
    lref()


def rgb8(wb, x, afrac):
    c = wb.tensors[x][0]
    t = wb.op(0x0B, [x], (c,), p_clamp(0, 1 << afrac))
    return wb.op(0x02, [t], (c,), rq_block(1, 0, [255], 1, 4, [afrac], 1, [], 0))


def whm1_lower(b, output_step):
    net = parse_net(b, b"WHM1", False)
    c0 = net_check(net)
    wb = WB()
    r = x_record(net, 0, c0)
    x = wb.input(r[0], r[2], r[3])
    x = net_layers(wb, net, x)
    if output_step:
        x = rgb8(wb, x, net["afrac"])
    return wb.bytes([x])


def wam1_lower(b):
    net = parse_net(b, b"WAM1", True)
    c0 = net_check(net)
    wb = WB()
    r = x_record(net, 0, c0)
    x = wb.input(r[0], r[2], r[3])
    wb.scalars.append((1, 1 << 26))
    x = net_layers(wb, net, x)
    c = wb.tensors[x][0]
    lo = -((1 << 63) // 32767)
    hi = (I64_MAX - (1 << (net["afrac"] - 1))) // 32767
    crop = struct.pack("<B3xIB3xIB3xIB3xI", 0, 0, 0, 0, 0, 1, 1, 0)
    t = wb.op(0x14, [x], (c, 2, lo, hi), crop)
    t = wb.op(0x02, [t], (c,), rq_block(1, 0, [32767], 1, 4, [net["afrac"]], 1, [], 0))
    t = wb.op(0x0B, [t], (c,), p_clamp(-32768, 32767))
    return wb.bytes([t])


def wmb1_lower(b):
    if b[:4] != b"WMB1":
        lref()
    rd = Reader(b, 4)
    try:
        secs = [rd.take(rd.u(4)) for _ in range(4)]
    except Refuse:
        lref()
    if rd.p != len(b):
        lref()
    nets = [parse_net(s, b"WHM1", False) for s in secs[:3]]
    if len({n["afrac"] for n in nets}) != 1:
        lref()
    try:
        tables = json.loads(secs[3])
    except ValueError:
        lref()
    a = tables.get("afrac") if isinstance(tables, dict) else None
    if isinstance(a, bool) or not isinstance(a, int) or a != nets[2]["afrac"]:
        lref()
    return [("h_s", net_graph(nets[0], False)), ("entropy", net_graph(nets[1], False)), ("g_s", net_graph(nets[2], True))]


def net_graph(net, step):
    c0 = net_check(net)
    wb = WB()
    r = x_record(net, 0, c0)
    x = wb.input(r[0], r[2], r[3])
    x = net_layers(wb, net, x)
    if step:
        x = rgb8(wb, x, net["afrac"])
    return wb.bytes([x])


def world_lower(tb, gb, c):
    tn, gn = parse_net(tb, b"WHM1", False), parse_net(gb, b"WHM1", False)
    cin = net_check(tn)
    net_check(gn)
    if tn["afrac"] != gn["afrac"] or c < 1 or cin <= c:
        lref()
    A = cin - c
    wb = WB()
    z = wb.input(c, I64_MIN, I64_MAX)
    parts = [z]
    for _ in range(A):
        wb.scalars.append((I64_MIN, I64_MAX))
        s = len(wb.scalars) - 1
        parts.append(wb.op(0x17, [z], (1,), struct.pack("<B7xq", 1, s)))
    x0 = wb.op(0x06, parts, x_record(tn, 0, cin), b"")
    x = net_layers(wb, tn, x0)
    one = 1 << tn["afrac"]
    out = wb.op(0x0B, [x], (wb.tensors[x][0],), p_clamp(-one, one))
    return [("T", wb.bytes([out])), ("g_s", net_graph(gn, True))]


def jpegai_lower(norm_b, hy_b, huv_b):
    """jpegai-int/1, section 13.6.1."""
    try:
        norm, hy, huv = json.loads(norm_b), json.loads(hy_b), json.loads(huv_b)
    except ValueError:
        lref()
    isint = lambda v: isinstance(v, int) and not isinstance(v, bool)
    ints = lambda v: v if isinstance(v, list) and all(isint(x) for x in v) else lref()
    if not isinstance(norm, dict):
        lref()
    pmf = norm.get("pmf")
    if not isinstance(pmf, list) or len(pmf) < 32:
        lref()
    widths = set()
    for row in pmf:
        r = ints(row)
        if not r or min(r) < 0 or sum(r) != 256:
            lref()
        widths.add(len(r))
    if len(widths) != 1:
        lref()
    bounds = ints(norm.get("bounds"))
    if len(bounds) < 32 or any(not 0 <= v <= 255 for v in bounds):
        lref()
    if not isint(norm.get("per_channel")) or norm["per_channel"] != 16:
        lref()
    for k, rows in (("cdfs_y", 160), ("cdfs_uv", 96)):
        v = ints(norm.get(k))
        if len(v) != rows * 63 or any(not 0 <= x <= 255 for x in v) or any(v[r * 63 + j] > v[r * 63 + j + 1] for r in range(rows) for j in range(62)):
            lref()
    gy, guv = ints(norm.get("gain_y")), ints(norm.get("gain_uv"))
    if len(gy) != 160 or len(guv) != 96 or any(not I64_MIN <= v <= I64_MAX for v in gy + guv):
        lref()

    def member(o, k, dims):
        m = o.get(k)
        if not isinstance(m, dict) or m.get("dims") != dims:
            lref()
        v = ints(m.get("v"))
        n = 1
        for d in dims:
            n *= d
        if len(v) != n or any(not I64_MIN <= x <= I64_MAX for x in v):
            lref()
        return v

    def plane(o, cz, gains):
        if not isinstance(o, dict):
            lref()

        def d0(k):
            m = o.get(k)
            dims = m.get("dims") if isinstance(m, dict) else None
            if not isinstance(dims, list) or not dims or not isint(dims[0]) or not 1 <= dims[0] <= 0xFFFFFFFF:
                lref()
            return dims[0]

        c1, c2 = d0("conv1.weight"), d0("depthwise.weight")
        c3 = 16 * cz
        wb = WB()
        x = wb.input(cz, -31, 32, 4, 4)
        blocks = (("conv1", c1, cz, 1, 0, "onnx::Pow_85"), ("depthwise", c2, c1, 3, 1, "onnx::Pow_91"), ("pointwise", c3, c2, 1, 0, "onnx::Pow_97"))
        for bi, (name, cout, cin, k, pad, sh) in enumerate(blocks):
            w = member(o, name + ".weight", [cout, cin, k, k])
            bias = member(o, name + ".bias", [cout])
            shifts = member(o, sh, [1, cout, 1, 1])
            if any(not 0 <= s <= 127 for s in shifts):
                lref()
            x = wb.op(0x0B, [x], (wb.tensors[x][0],), p_clamp(-128, 127))
            rq = rq_block(0, 1, [], 0, 0, shifts, 2, bias, smallest(bias, (4, 8)))
            wbytes = smallest(w, (1, 2, 4, 8))
            x = wb.op(0x01, [x], (cout,), p_conv(cout, k, k, 1, 1, pad, pad, 1, wbytes, rq, w))
            if bi < 2:
                x = wb.op(0x0C, [x], (cout,), b"")
        x = wb.op(0x09, [x], (cz,), struct.pack("<I", 4))
        x = wb.op(0x13, [x], (cz,), b"")
        x = wb.op(0x0B, [x], (cz,), p_clamp(0, 3967))
        x = wb.op(0x02, [x], (cz,), rq_block(0, 1, [], 0, 0, [0], 1, gains, smallest(gains, (4, 8))))
        return wb.bytes([x])

    return [("hs_y", plane(hy, 160, gy)), ("hs_uv", plane(huv, 96, guv))]


def lower_set(entry, files):
    low, cap = entry["lowering"], entry["capability"]
    try:
        if low == "int_synth/1":
            roles = int_synth_lower(files["int_synth.json"])
        elif low == "whm1/1" and cap == "wai.splat.int_codec":
            roles = [("g_s", whm1_lower(files["g_s.whm"], True))]
        elif low == "whm1/1":
            roles = [("h_s", whm1_lower(files["h_s.whm"], False)), ("g_s", whm1_lower(files["g_s.whm"], True))]
        elif low == "wam1/1":
            roles = [("g_s", wam1_lower(files["model.wam"]))]
        elif low == "wmb1/1":
            roles = wmb1_lower(files["model.wmb"])
        elif low == "world/1":
            roles = world_lower(files["T.whm"], files["g_s.whm"], dec(entry["c"]))
        elif low == "jpegai-int/1":
            roles = jpegai_lower(files["norm.json"], files["weights/hs_y.json"], files["weights/hs_uv.json"])
        else:
            raise ValueError("no lowering %r" % low)
    except LowerRefuse:
        raise ValueError("the lowering refused the set")
    # Every lowered graph loads.
    for _, b in roles:
        load(b)
    return roles


def set_digest(files):
    if len(files) == 1:
        return hashlib.sha256(next(iter(files.values()))).hexdigest()
    listing = "".join("%s  %s\n" % (hashlib.sha256(files[n]).hexdigest(), n) for n in sorted(files, key=lambda s: s.encode()))
    return hashlib.sha256(listing.encode()).hexdigest()


# --------------------------------------------------------------- main


def check(root, write):
    wai = os.path.dirname(os.path.abspath(root))
    failures = 0
    nfiles = nvec = 0
    for sub, fn in (("vectors", graph_line), ("stage", stage_line), ("graphs", graph_line)):
        d = os.path.join(root, sub)
        for name in sorted(os.listdir(d)):
            if not name.endswith(".jsonl"):
                continue
            raw = open(os.path.join(d, name), "rb").read()
            out = []
            try:
                for line in raw.decode("ascii").split("\n")[:-1]:
                    out.append(cw(fn(json.loads(line))) + "\n")
            except Exception as e:
                failures += 1
                print("FAIL %s/%s: %s: %s" % (sub, name, type(e).__name__, e))
                continue
            got = "".join(out).encode()
            nfiles += 1
            nvec += len(out)
            if write:
                os.makedirs(os.path.join(write, sub), exist_ok=True)
                open(os.path.join(write, sub, name), "wb").write(got)
            if got != raw:
                failures += 1
                a, b = got.split(b"\n"), raw.split(b"\n")
                k = next((k for k in range(min(len(a), len(b))) if a[k] != b[k]), min(len(a), len(b)))
                print("DIFF %s/%s: line %d differs" % (sub, name, k + 1))
    print("rewrote %d vector files, %d vectors" % (nfiles, nvec))
    idx = json.load(open(os.path.join(root, "certificates", "index.json")))
    nsets = ngraphs = 0
    lowered = {}
    for e in idx["sets"]:
        try:
            files = {m: open(os.path.join(wai, p), "rb").read() for m, p in e["members"].items()}
            if set_digest(files) != e["set_sha256"]:
                raise ValueError("the set's digest differs")
            roles = lower_set(e, files)
            if len(roles) != len(e["roles"]):
                raise ValueError("the roles differ")
            for (role, b), path in zip(roles, e["roles"]):
                if not path.endswith("/%s/%s.whg1" % (e["set_sha256"], role)):
                    raise ValueError("role %s is at %s" % (role, path))
                if write:
                    p = os.path.join(write, os.path.relpath(os.path.join(wai, path), root))
                    os.makedirs(os.path.dirname(p), exist_ok=True)
                    open(p, "wb").write(b)
                if open(os.path.join(wai, path), "rb").read() != b:
                    raise ValueError("%s: the lowered graph differs" % role)
                if e.get("synthetic"):
                    lowered[(FORM_FILE[e["lowering"]], e["name"], role)] = b
                ngraphs += 1
            nsets += 1
        except Exception as ex:
            failures += 1
            print("FAIL lowering %s (%s): %s" % (e.get("name"), e.get("capability"), ex))
    print("re-lowered %d sets, %d graphs" % (nsets, ngraphs))
    failures += check_lower_vectors(root, lowered)
    return failures


FORM_FILE = {
    "int_synth/1": "lower_int_synth", "whm1/1": "lower_whm1", "wam1/1": "lower_wam1",
    "wmb1/1": "lower_wmb1", "world/1": "lower_world", "jpegai-int/1": "lower_jpegai",
}


def check_lower_vectors(root, lowered):
    """Each vector of graphs/lower_<form>.jsonl has the id <file>/<set>/<role>/...: its
    graph must be that role of that synthetic set as re-lowered here, and every role
    of every synthetic set must have one. `lowered` maps (file, set, role) to bytes."""
    failures = 0
    seen = set()
    d = os.path.join(root, "graphs")
    for name in sorted(os.listdir(d)):
        if not (name.startswith("lower_") and name.endswith(".jsonl")):
            continue
        f = name[:-len(".jsonl")]
        for k, line in enumerate(open(os.path.join(d, name), "rb").read().decode("ascii").split("\n")[:-1]):
            v = json.loads(line)
            parts = v["id"].split("/")
            key = tuple(parts[:3])
            if len(parts) < 4 or parts[0] != f or key not in lowered:
                failures += 1
                print("FAIL graphs/%s line %d: %s names no lowered role" % (name, k + 1, v["id"]))
            elif bytes.fromhex(v["graph"]) != lowered[key]:
                failures += 1
                print("FAIL graphs/%s line %d: %s is not the role as lowered" % (name, k + 1, v["id"]))
            seen.add(key)
    for key in sorted(set(lowered) - seen):
        failures += 1
        print("FAIL %s: no lowering vector for this role" % "/".join(key))
    print("matched %d lowered roles to their lowering vectors" % len(seen & set(lowered)))
    return failures


def main(argv):
    if len(argv) not in (2, 4) or (len(argv) == 4 and argv[2] != "--write"):
        print(__doc__)
        return 2
    write = argv[3] if len(argv) == 4 else None
    failures = check(argv[1], write)
    if failures:
        print("%d failures" % failures)
        return 1
    print("ops-conformance: every vector file and every lowered graph rewritten byte for byte")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))