graph_net.py

"""MageZero's graph network (Will Wroble, WillWroble/MageZero branch graph-encoder), for DraftZero's
imitation training and search (docs/022).

The network is MageZero's `NetGraph` (src/magezero/model.py at de225045, 2026-10-04), copied here
rather than imported: that branch's modules import each other as top-level names (`from vocab import
...`), and DraftZero pins MageZero v0.2 for the flat networks. The code below is the upstream model with
four changes, none of which changes what a default network computes:

  * its sizes are constructor arguments (`arch`), so a sweep can vary them; the defaults are
    upstream's constants, and a default network's state dict is upstream's, so it loads in
    MageZero's own graph server (`load_checkpoint` there);
  * `leaf_dropout`: drop a share of leaf edges while training, the graph analogue of the flat
    networks' token dropout (0 = upstream);
  * `local_depth`: LocalLayers per type within each pass (1 = upstream). Each extra layer re-attends
    from the node's updated embedding to the same children; extra layers live in `local_extra`, so a
    depth-1 network's state dict is upstream's;
  * `forward` also returns the value head's input to its Tanh (`value_x`), which the imitation
    trainer's cross-entropy needs, as supervised.value_logit does for the flat networks.

Input: the graph encoder's state graph (graph_tables), nodes typed ROOT, PLAYER, ZONE, STACK_OBJECT,
PERMANENT, CARD, ABILITY or LEAF, edges child -> parent with a label:

  embeddings: typed node = type embedding; leaf = leaf vocab row + numeric value bucket
  local layers: `passes` bottom-up passes, each `local_depth` LocalLayers per type, run in order
      ABILITY -> CARD -> PERMANENT -> STACK_OBJECT -> ZONE -> PLAYER -> ROOT
      every node attends over [itself, children + edge label embedding]
  global layers: TransformerEncoder over [CLS, the state's internal nodes]
      ├── MLP per node  priority   (read at ABILITY nodes: a play is its ability's node; Pass has one)
      ├── MLP per node  target     (read at CARD / PERMANENT / STACK_OBJECT / PLAYER nodes; Stop Choosing is a CARD)
      ├── MLP on CLS    choose_use (false, true)
      └── MLP on CLS    value      (tanh, -1..1)
"""
from __future__ import annotations

import math
from enum import IntEnum
from typing import NamedTuple

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

UPSTREAM = "WillWroble/MageZero@de225045c82a2ddc107b99ac833e21ab6d566279 (graph-encoder)"
GLOBAL_MAX = 2 ** 31 - 1

# upstream's constants (model.py): the default architecture
ARCH_DEFAULT = {"type": "graph", "d_model": 512, "heads": 4, "ff": 1024, "dropout": 0.25, "passes": 2,
                "global_layers": 2, "head_hidden": 256, "leaf_dropout": 0.0, "local_depth": 1}

# numeric value buckets: exact -5..20 (P/T, counters, mana, costs, small counts), then 21-25, 26-30,
# 31-40, 41-60, 61+ (life totals, library and zone sizes); everything below -5 is one bucket
VALUE_BOUNDS = list(range(-5, 22)) + [26, 31, 41, 61]


class NodeType(IntEnum):
    """Java FeatureGraph.Node.Type, plus the root."""
    ROOT = 0
    PLAYER = 1
    ZONE = 2
    STACK_OBJECT = 3
    PERMANENT = 4
    CARD = 5
    ABILITY = 6
    LEAF = 7


STAGES = (NodeType.ABILITY, NodeType.CARD, NodeType.PERMANENT, NodeType.STACK_OBJECT,
          NodeType.ZONE, NodeType.PLAYER, NodeType.ROOT)
PRIORITY_TYPES = (NodeType.ABILITY,)
TARGET_TYPES = (NodeType.CARD, NodeType.PERMANENT, NodeType.STACK_OBJECT, NodeType.PLAYER)


# ------------------------------------------------------------------------------------------------
# ids: the Java encoder's hash (upstream vocab.py), to tell typed nodes from leaves
# ------------------------------------------------------------------------------------------------

_MASK = (1 << 64) - 1
_GOLDEN = 0x9E3779B185EBCA87


def _mix64(z: int) -> int:
    z = ((z ^ (z >> 30)) * 0xBF58476D1CE4E5B9) & _MASK
    z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & _MASK
    return z ^ (z >> 31)


def feature_id(name: str) -> int:
    """The id XMage stores for a leaf name or edge label: StateEncoder.indexFor(hash64(name))."""
    data = name.encode("utf-8")
    h = _mix64(_GOLDEN ^ (len(data) * _GOLDEN & _MASK))
    full = len(data) // 8 * 8
    for i in range(0, full, 8):
        h ^= _mix64(int.from_bytes(data[i:i + 8], "little"))
        h = (((h << 27 | h >> 37) & _MASK) * _GOLDEN + 0x165667B19E3779F9) & _MASK
    h ^= _mix64(int.from_bytes(data[full:], "little"))
    h = ((h ^ (h >> 33)) * 0xFF51AFD7ED558CCD) & _MASK
    h = ((h ^ (h >> 33)) * 0xC4CEB9FE1A85EC53) & _MASK
    h ^= h >> 33
    if h == 1 << 63:  # Long.MIN_VALUE: Java's -h overflows and the remainder keeps the sign
        return -((1 << 63) % GLOBAL_MAX)
    return (h if h < 1 << 63 else (1 << 64) - h) % GLOBAL_MAX


# a typed node's id is the hash of its type name and the root's id is 0; every other node is a leaf
TYPE_IDS = {0: NodeType.ROOT}
for _t in NodeType:
    if _t not in (NodeType.ROOT, NodeType.LEAF):
        TYPE_IDS[feature_id(_t.name)] = _t


def node_types(ids: np.ndarray) -> np.ndarray:
    types = np.full(ids.shape, NodeType.LEAF, dtype=np.int8)
    for type_id, t in TYPE_IDS.items():
        types[ids == type_id] = t
    return types


def full_arch(arch: dict | None) -> dict:
    a = {**ARCH_DEFAULT, **(arch or {})}
    unknown = set(a) - set(ARCH_DEFAULT)
    if unknown:
        raise ValueError(f"unknown graph arch keys {sorted(unknown)} (known: {sorted(ARCH_DEFAULT)})")
    if a["type"] != "graph":
        raise ValueError(f"arch.type {a['type']!r}: graph_net builds only 'graph'")
    if a["d_model"] % a["heads"]:
        raise ValueError(f"d_model {a['d_model']} is not a multiple of heads {a['heads']}")
    return a


def upstream_loadable(arch: dict) -> bool:
    """A network MageZero's own graph server loads as it is (upstream's constants)."""
    a = full_arch(arch)
    return all(a[k] == ARCH_DEFAULT[k] for k in ("d_model", "heads", "ff", "passes", "global_layers", "head_hidden",
                                                 "local_depth"))


# ------------------------------------------------------------------------------------------------
# the network
# ------------------------------------------------------------------------------------------------

class Graphs(NamedTuple):
    """A batch of state graphs, concatenated: the network's input."""
    node_type: torch.Tensor     # [N] NodeType
    node_row: torch.Tensor      # [N] leaf vocab row; -1 for typed nodes and leaves outside the vocab
    node_value: torch.Tensor    # [N] numeric value (0 for non-numeric)
    node_offsets: torch.Tensor  # [B+1] first node of each state
    edge_child: torch.Tensor    # [E] node index (batch-wide)
    edge_parent: torch.Tensor   # [E] node index (batch-wide)
    edge_label: torch.Tensor    # [E] edge vocab row + 1; 0 for labels outside the vocab

    def to(self, device):
        return Graphs(*(t.to(device, non_blocking=True) for t in self))


class Out(NamedTuple):
    priority: torch.Tensor      # [N] per-node priority score (-inf at leaves)
    target: torch.Tensor        # [N] per-node target score (-inf at leaves)
    use: torch.Tensor           # [B, 2] choose_use logits (false, true)
    value: torch.Tensor         # [B] tanh(value_x)
    value_x: torch.Tensor       # [B] the value head's input to its Tanh


def graph_index(offsets: torch.Tensor) -> torch.Tensor:
    """State of each node, from a batch's node offsets."""
    return torch.repeat_interleave(torch.arange(offsets.numel() - 1, device=offsets.device), offsets.diff())


def segment_logsumexp(x: torch.Tensor, seg: torch.Tensor, n: int) -> torch.Tensor:
    """Log-sum-exp over the rows of x that share an index in seg (values 0..n-1); returns n rows
    (-inf for an index with no row)."""
    index = seg.view(-1, *([1] * (x.dim() - 1))).expand_as(x)
    with torch.no_grad():  # the max only stabilizes; the result and its gradient don't depend on it
        m = x.new_full((n, *x.shape[1:]), float("-inf")).scatter_reduce(0, index, x, "amax")
        m = torch.where(torch.isfinite(m), m, torch.zeros_like(m))
    return m + x.new_zeros((n, *x.shape[1:])).index_add(0, seg, (x - m[seg]).exp()).log()


class LocalLayer(nn.TransformerEncoderLayer):
    """A pre-norm transformer layer run over each node's [node, children...] sequence, keeping only
    the node's output token (upstream's LocalLayer)."""

    def __init__(self, d_model: int, heads: int, ff: int, dropout: float):
        super().__init__(d_model, heads, ff, dropout, activation="gelu", batch_first=True, norm_first=True)

    def forward(self, node: torch.Tensor, child: torch.Tensor, seg: torch.Tensor) -> torch.Tensor:
        """node: [S, d] this stage's nodes; child: [E, d] child tokens (child embedding + edge label
        embedding); seg: [E] position in `node` of each child's parent. Returns the nodes' new embeddings."""
        S, d = node.shape
        H, dh = self.self_attn.num_heads, self.self_attn.head_dim
        w, b = self.self_attn.in_proj_weight, self.self_attn.in_proj_bias
        q, k, v = F.linear(self.norm1(node), w, b).chunk(3, dim=-1)
        k_child, v_child = F.linear(self.norm1(child), w[d:], b[d:]).chunk(2, dim=-1)
        seg = torch.cat([torch.arange(S, device=seg.device), seg])  # each node also attends to itself
        k = torch.cat([k, k_child]).view(-1, H, dh)
        v = torch.cat([v, v_child]).view(-1, H, dh)
        score = (q.reshape(S, H, dh)[seg] * k).sum(-1).float() / math.sqrt(dh)  # [S+E, H]
        attn = (score - segment_logsumexp(score, seg, S)[seg]).exp()
        attn = F.dropout(attn, self.self_attn.dropout, self.training)
        out = torch.zeros(S, H, dh, device=node.device).index_add_(
            0, seg, (attn.to(v.dtype).unsqueeze(-1) * v).float())
        node = node + self.dropout1(self.self_attn.out_proj(out.view(S, d).to(node.dtype)))
        return node + self._ff_block(self.norm2(node))


def _mlp(d: int, hidden: int, out_dim: int, *tail: nn.Module) -> nn.Sequential:
    return nn.Sequential(nn.Linear(d, hidden), nn.ReLU(), nn.Linear(hidden, out_dim), *tail)


class NetGraph(nn.Module):
    def __init__(self, num_embeddings: int = 0, num_edge_labels: int = 0, arch: dict | None = None):
        super().__init__()
        a = self.arch = full_arch(arch)
        d = a["d_model"]
        self.embedding = nn.Embedding(num_embeddings, d)                         # leaves
        self.edge_embedding = nn.Embedding(num_edge_labels + 1, d, padding_idx=0)  # row 0: unknown label
        self.type_embedding = nn.Embedding(len(NodeType), d)
        self.value_embedding = nn.Embedding(len(VALUE_BOUNDS) + 1, d)
        self.register_buffer("value_bounds", torch.tensor(VALUE_BOUNDS), persistent=False)
        self.local_layers = nn.ModuleList(
            nn.ModuleDict({t.name: LocalLayer(d, a["heads"], a["ff"], a["dropout"]) for t in STAGES})
            for _ in range(a["passes"]))
        if a["local_depth"] < 1:
            raise ValueError(f"local_depth {a['local_depth']}: at least 1")
        if a["local_depth"] > 1:  # the layers after each type's first, per pass (absent at depth 1: upstream's layout)
            self.local_extra = nn.ModuleList(
                nn.ModuleDict({t.name: nn.ModuleList(LocalLayer(d, a["heads"], a["ff"], a["dropout"])
                                                     for _ in range(a["local_depth"] - 1)) for t in STAGES})
                for _ in range(a["passes"]))
        self.cls = nn.Parameter(torch.randn(d))
        global_layer = nn.TransformerEncoderLayer(d, a["heads"], a["ff"], a["dropout"], activation="gelu",
                                                  batch_first=True, norm_first=True)
        self.global_layers = nn.TransformerEncoder(global_layer, a["global_layers"], norm=nn.LayerNorm(d),
                                                   enable_nested_tensor=False)
        h = a["head_hidden"]
        self.priority_head = _mlp(d, h, 1)
        self.target_head = _mlp(d, h, 1)
        self.binary_head = _mlp(d, h, 2)
        self.value_head = _mlp(d, h, 1, nn.Tanh())

    def forward(self, g: Graphs) -> Out:
        B, N = g.node_offsets.numel() - 1, g.node_type.numel()

        # initial embeddings: typed nodes their type, leaves their name and numeric value
        leaves = (g.node_row >= 0).nonzero().squeeze(1)
        value = torch.bucketize(g.node_value[leaves], self.value_bounds, right=True)
        h = self.type_embedding(g.node_type).index_add(
            0, leaves, self.embedding(g.node_row[leaves]) + self.value_embedding(value))

        # leaves outside the vocab have no embedding, so their edges are dropped (and, while
        # training, a leaf_dropout share of the others)
        child_leaf = g.node_type[g.edge_child] == NodeType.LEAF
        keep = ~child_leaf | (g.node_row[g.edge_child] >= 0)
        if self.training and self.arch["leaf_dropout"] > 0:
            keep = keep & ~(child_leaf & (torch.rand(keep.shape, device=keep.device) < self.arch["leaf_dropout"]))
        child, parent, label = g.edge_child[keep], g.edge_parent[keep], g.edge_label[keep]

        # one stage per type: every node of that type, its children, and each child's parent position
        parent_type = g.node_type[parent]
        stages = []
        for t in STAGES:
            nodes = (g.node_type == t).nonzero().squeeze(1)
            if len(nodes):
                edges = (parent_type == t).nonzero().squeeze(1)
                stages.append((t.name, nodes, child[edges], label[edges], torch.searchsorted(nodes, parent[edges])))

        # local layers, bottom-up passes: children earlier in the order arrive updated in this pass,
        # children of the same type or later in the order (attachments, targets, linked exile) from the last one
        extra = getattr(self, "local_extra", None)
        for p, layers in enumerate(self.local_layers):
            for name, nodes, c, lbl, seg in stages:
                kids = h[c] + self.edge_embedding(lbl)
                x = layers[name](h[nodes], kids, seg).to(h.dtype)
                for layer in (extra[p][name] if extra is not None else ()):
                    x = layer(x, kids, seg).to(h.dtype)
                h = h.index_copy(0, nodes, x)

        # global layers: self-attention over [CLS, internal nodes], one padded sequence per state
        internal = (g.node_type != NodeType.LEAF).nonzero().squeeze(1)
        gi = graph_index(g.node_offsets)[internal]
        counts = torch.bincount(gi, minlength=B)
        pos = torch.arange(len(internal), device=h.device) - (counts.cumsum(0) - counts)[gi] + 1
        x = h.new_zeros(B, int(counts.max()) + 1, h.shape[1])
        x[:, 0] = self.cls.to(h.dtype)
        x[gi, pos] = h[internal]
        pad = torch.arange(x.shape[1], device=h.device) > counts.unsqueeze(1)
        x = self.global_layers(x, src_key_padding_mask=pad)
        cls, z = x[:, 0], x[gi, pos]

        def per_node(head):
            return torch.full((N,), float("-inf"), device=h.device).index_copy(
                0, internal, head(z).squeeze(-1).float())

        vx = self.value_head[:-1](cls).squeeze(-1).float()
        return Out(per_node(self.priority_head), per_node(self.target_head), self.binary_head(cls).float(),
                   torch.tanh(vx), vx)


def parameter_counts(model: NetGraph) -> dict:
    out = {"total": 0}
    for name, p in model.named_parameters():
        part = name.split(".")[0]
        out[part] = out.get(part, 0) + p.numel()
        out["total"] += p.numel()
    return out


# ------------------------------------------------------------------------------------------------
# batches: graph rows -> the network's input, and options -> logits
# ------------------------------------------------------------------------------------------------

def gather(ptr: np.ndarray, sel: np.ndarray, *arrays: np.ndarray) -> tuple:
    """The CSR slices of rows `sel` (ptr over `arrays`), concatenated, and their new ptr."""
    lens = (ptr[sel + 1] - ptr[sel]).astype(np.int64)
    new = np.zeros(len(sel) + 1, np.int64)
    np.cumsum(lens, out=new[1:])
    pos = np.repeat(ptr[sel] - new[:-1], lens) + np.arange(new[-1])
    return (new,) + tuple(a[pos] for a in arrays)


class OptionBatch(NamedTuple):
    """Per option of a batch: its row, its nodes (batch-wide) and which node belongs to which option."""
    opt_row: torch.Tensor       # [O] row of each option
    node: torch.Tensor          # [K] batch-wide node index of each option node
    node_opt: torch.Tensor      # [K] option of each option node
    n_opt: int

    def to(self, device):
        return OptionBatch(self.opt_row.to(device), self.node.to(device), self.node_opt.to(device), self.n_opt)


def option_logits(out: Out, ob: OptionBatch, graph_type: torch.Tensor, opt_index: torch.Tensor) -> torch.Tensor:
    """[O] logit of each option: log-sum-exp of its nodes' scores under the head its row's decision type
    reads (priority at PRIORITY rows, target at CHOOSE_TARGET rows; two copies of a card are one option,
    so their probabilities add), or the choose_use logit at CHOOSE_USE rows (options [no, yes]).
    opt_index: [O] each option's position within its row."""
    gt = graph_type[ob.opt_row[ob.node_opt]]
    s = torch.where(gt == 0, out.priority[ob.node], out.target[ob.node])
    lg = segment_logsumexp(s.float(), ob.node_opt, ob.n_opt)
    is_use = graph_type[ob.opt_row] == 5
    if is_use.any():
        u = out.use.float()[ob.opt_row, opt_index.clamp(max=1)]
        lg = torch.where(is_use, u, lg)
    return lg