def attention(X, Wq, Wk, Wv, mask=None, bias=None):
    """Single-head scaled dot-product attention. mask[i, j] = False blocks j → i; bias[i, j] is added to the score."""
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    S = Q @ K.T / K.size(1) ** 0.5
    if bias is not None:
        S = S + bias
    if mask is not None:
        S = S.masked_fill(~mask, float("-inf"))
    return torch.softmax(S, dim=1) @ V

karate = nx.karate_club_graph()
A_k = torch.tensor(
    nx.to_numpy_array(karate, weight=None), dtype=torch.bool
)
torch.manual_seed(SEED)
_X = torch.randn(34, 8)
_Wq, _Wk, _Wv = (torch.randn(8, 8) for _ in range(3))
_masked = attention(
    _X, _Wq, _Wk, _Wv, mask=A_k | torch.eye(34, dtype=torch.bool)
)

# the same thing as explicit message passing: for every node, softmax over its neighbours (+ itself)
_loop = torch.zeros(34, 8)
for _v in range(34):
    _nbrs = [_v] + list(karate.neighbors(_v))
    _s = torch.stack(
        [(_X[_v] @ _Wq) @ (_X[_u] @ _Wk) / 8**0.5 for _u in _nbrs]
    )
    _loop[_v] = (
        torch.softmax(_s, 0)[:, None]
        * torch.stack([_X[_u] @ _Wv for _u in _nbrs])
    ).sum(0)

_full = attention(
    _X, _Wq, _Wk, _Wv
)  # no mask: the graph is not an input at all
pd.DataFrame(
    [
        {
            "check": "masked attention vs. per-node message-passing loop",
            "max |difference|": (_masked - _loop).abs().max().item(),
        },
        {
            "check": "masked attention vs. full attention",
            "max |difference|": (_masked - _full).abs().max().item(),
        },
    ]
)

class DenseGCN(nn.Module):
    def __init__(self, n_layers, d=32, d_in=2):
        super().__init__()
        self.inp = nn.Linear(d_in, d)
        self.layers = nn.ModuleList(
            nn.Linear(d, d) for _ in range(n_layers)
        )
        self.out = nn.Linear(d, 2)

    def forward(self, b, pe=None):
        A_t = b["A"] + torch.eye(b["A"].size(-1))
        d = A_t.sum(-1).pow(-0.5)
        A_hat = d[..., :, None] * A_t * d[..., None, :]
        h = F.relu(self.inp(b["X"]))
        for lin in self.layers:
            h = F.relu(A_hat @ lin(h))
        return self.out(h)

class GraphTransformer(nn.Module):
    """Pre-built PyTorch multi-head attention + FFN blocks; optional PE (concatenated) and shortest-path bias."""

    def __init__(
        self,
        n_layers=2,
        d=32,
        heads=4,
        d_pe=0,
        spd_bias=False,
        max_dist=N_TREE,
    ):
        super().__init__()
        self.heads = heads
        self.inp = nn.Linear(2 + d_pe, d)
        self.att = nn.ModuleList(
            nn.MultiheadAttention(d, heads, batch_first=True)
            for _ in range(n_layers)
        )
        self.ff = nn.ModuleList(
            nn.Sequential(
                nn.Linear(d, 2 * d), nn.ReLU(), nn.Linear(2 * d, d)
            )
            for _ in range(n_layers)
        )
        self.n1 = nn.ModuleList(nn.LayerNorm(d) for _ in range(n_layers))
        self.n2 = nn.ModuleList(nn.LayerNorm(d) for _ in range(n_layers))
        self.spd = (
            nn.Embedding(max_dist, heads) if spd_bias else None
        )  # b(d(i, j)) per head
        self.out = nn.Linear(d, 2)

    def forward(self, b, pe=None):
        h = self.inp(b["X"] if pe is None else torch.cat([b["X"], pe], -1))
        B, n, _ = h.shape
        bias = None
        if self.spd is not None:
            bias = (
                self.spd(b["D"])
                .permute(0, 3, 1, 2)
                .reshape(B * self.heads, n, n)
            )  # float mask = additive bias
        for att, ff, n1, n2 in zip(self.att, self.ff, self.n1, self.n2):
            h = n1(h + att(h, h, h, attn_mask=bias, need_weights=False)[0])
            h = n2(h + ff(h))
        return self.out(h)

def sign_flip(pe):
    return pe * (torch.randint(0, 2, (pe.size(0), 1, pe.size(2))) * 2 - 1)

def fit_trees(model, pe_key=None, epochs=150, augment=None, seed=SEED):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    tr = trees_train
    for _ in range(epochs):
        for i in range(0, tr["X"].size(0), 50):
            batch = {k: v[i : i + 50] for k, v in tr.items()}
            pe = None if pe_key is None else batch[pe_key]
            if pe is not None and augment:
                pe = augment(pe)
            opt.zero_grad()
            F.cross_entropy(
                model(batch, pe).reshape(-1, 2), batch["y"].reshape(-1)
            ).backward()
            opt.step()
    with torch.no_grad():
        out = model(
            trees_test, None if pe_key is None else trees_test[pe_key]
        )
    return (out.argmax(-1) == trees_test["y"]).float().mean().item()

tree_run = mo.ui.run_button(label="Train the 6 models (≈ 10 s)")
tree_run

_rev = reveal(
    r"""
Measured (one seed): **GCN-2 0.87**, since it cannot reach 3 hops, so nodes at exactly distance 3 are lost; **GCN-4 0.99**.
**Transformer without structure 0.65**, near "always say no" (0.57): attention can find the marked node anywhere but has no idea
how far away it is. **Random-walk PE barely helps (0.67)**: return probabilities describe a node's own neighbourhood, not its
distance to *another* node. **Laplacian PE 0.97**: eigenvectors are global coordinates, so distances are recoverable.
**Shortest-path bias 1.00**: the needed quantity is injected straight into the attention scores.
Lesson: *what* structural information you give a transformer must match the task.

""",
    predict="Rank GCN-2, GCN-4, a structure-blind transformer, + random-walk PE, + Laplacian PE and + shortest-path bias. Which structural signal actually encodes the distance between *two* nodes?",
)
_post = mo.md(r"""
## 4 · PEARL: positional encodings from random node IDs

Eigenvectors cost $O(n^3)$ and are sign/basis-ambiguous. **PEARL** (Kanatsoulis et al., ICLR 2025): give every node $M$ random
numbers (random "IDs"), run the **same GNN** on each random sample, and **average** over samples. Each sample breaks symmetry
like a one-hot ID; the average restores permutation equivariance, and the GNN costs $O(|E|)$ per sample. Trained end-to-end:
""")
mo.vstack([_rev, _post])

_rev = reveal(
    r"""
Measured: PEARL lifts the structure-blind transformer from 0.65 to **≈ 0.81** (0.82 with fresh random IDs at test time, so
the result does not depend on one random draw). It sits below the Laplacian PE (0.97) and shortest-path bias (1.00) on this
small, distance-defined task: those inject *exactly* the needed quantity, while PEARL must learn it. Its selling points are
cost ($O(M|E|)$ instead of $O(n^3)$) and no sign or basis ambiguity.

""",
    predict="Where will PEARL land between the structure-blind transformer (0.65) and the Laplacian PE (0.97)?",
)
_post = mo.md(r"""
## 5 · On Cora: do graph transformers beat a GCN?

Full attention over all 2,708 papers (7.3M attention scores per head per layer), compared with a 2-layer GCN and a **GPS-style**
layer (Rampášek et al. 2022): $\mathbf h \leftarrow \mathbf h + \text{GCN}(\mathbf h) + \text{Attention}(\mathbf h)$, local and global together.
""")
mo.vstack([_rev, _post])

cora = planetoid("Cora")[0]
cora_lap = T.AddLaplacianEigenvectorPE(
    k=16, attr_name="pe", is_undirected=True
)(cora.clone()).pe

class CoraModel(nn.Module):
    def __init__(self, kind, d=64, heads=4, d_pe=0):
        super().__init__()
        self.kind = kind
        self.inp = nn.Linear(cora.num_features + d_pe, d)
        self.gcn = (
            nn.ModuleList(GCNConv(d, d) for _ in range(2))
            if kind in ("GCN", "GPS")
            else None
        )
        self.att = (
            nn.ModuleList(
                nn.MultiheadAttention(d, heads, batch_first=True)
                for _ in range(2)
            )
            if kind != "GCN"
            else None
        )
        self.norm = nn.ModuleList(nn.LayerNorm(d) for _ in range(2))
        self.out = nn.Linear(d, 7)

    def forward(self, x, edge_index):
        h = F.dropout(F.relu(self.inp(x)), 0.5, self.training)
        for i in range(2):
            upd = 0
            if self.gcn is not None:
                upd = upd + self.gcn[i](h, edge_index)
            if self.att is not None:
                upd = (
                    upd
                    + self.att[i](
                        h[None], h[None], h[None], need_weights=False
                    )[0][0]
                )
            h = self.norm[i](
                h + F.dropout(F.relu(upd), 0.5, self.training)
            )
        return self.out(h)

def fit_cora(model, x, epochs=150, seed=SEED):
    torch.manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)
    best_val, best_test = 0, 0
    for _ in range(epochs):
        model.train()
        opt.zero_grad()
        F.cross_entropy(
            model(x, cora.edge_index)[cora.train_mask],
            cora.y[cora.train_mask],
        ).backward()
        opt.step()
        model.eval()
        with torch.no_grad():
            pred = model(x, cora.edge_index).argmax(1)
        val = (
            (pred[cora.val_mask] == cora.y[cora.val_mask])
            .float()
            .mean()
            .item()
        )
        if val >= best_val:
            best_val, best_test = (
                val,
                (pred[cora.test_mask] == cora.y[cora.test_mask])
                .float()
                .mean()
                .item(),
            )
    return best_test

cora_run = mo.ui.run_button(label="Train on Cora (≈ 40 s)")
cora_run

_rev = reveal(
    r"""
Measured (one seed, 150 epochs): **GCN 0.74 · GPS 0.75 · transformer 0.56 · transformer + LapPE 0.57**. On a small,
homophilous citation graph with 140 labels, the useful signal is *local* (cited papers share topics). Global attention alone
behaves like an MLP over the bag-of-words (≈ L03's 0.54) and costs 5× the time; Laplacian PEs don't rescue it here. GPS keeps
the local GCN path and matches it. Graph transformers pay off on **long-range** tasks (§3) and on molecules, where graphs are
small and global interactions matter; they are not a free upgrade.

""",
    predict="On Cora, will a full graph transformer beat a 2-layer GCN? Where does GPS (local + global) land?",
)
_post = mo.md(r"""
---

## 6 · Exercises

**Exercise 1 (by hand).** Graphormer bias: two keys with scaled scores 14 and 12, and the second gets bias $c_{12}=-2$.
What attention weight goes to token 2? → `ex1_alpha` (3 decimals).

**Exercise 2 (code).** Write `my_sdpa(Q, K, V, mask=None)` (shapes `[n, d]`; boolean `mask[i, j]=True` means *allowed*) that
matches `F.scaled_dot_product_attention`.

**Exercise 3 (cost).** How many attention scores does one head of full attention compute per layer on a 1M-node social graph? → `ex3_scores`.
""")
mo.vstack([_rev, _post])

ex1_alpha = None
ex3_scores = None

def my_sdpa(Q, K, V, mask=None):
    return None

def _verdict(ok, answered):
    if not answered:
        return mo.callout(mo.md("Not answered yet."), kind="neutral")
    return mo.callout(
        mo.md("Correct!" if ok else "Not yet — try again or open a hint."),
        kind="success" if ok else "warn",
    )

def _check_sdpa():
    try:
        g = torch.Generator().manual_seed(1)
        Q, K, V = (torch.randn(10, 8, generator=g) for _ in range(3))
        mask = torch.rand(10, 10, generator=g) > 0.3
        mask |= torch.eye(10, dtype=torch.bool)
        got = my_sdpa(Q, K, V, mask)
        if got is None:
            return _verdict(False, False)
        ref = F.scaled_dot_product_attention(
            Q[None], K[None], V[None], attn_mask=mask[None]
        )[0]
        ok = torch.allclose(got, ref, atol=1e-5) and torch.allclose(
            my_sdpa(Q, K, V),
            F.scaled_dot_product_attention(Q[None], K[None], V[None])[0],
            atol=1e-5,
        )
        return _verdict(ok, True)
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_alpha is not None
            and abs(ex1_alpha - 1 / (1 + np.e**4)) < 2e-3,
            ex1_alpha is not None,
        ),
        mo.md("**Exercise 2**"),
        _check_sdpa(),
        mo.md("**Exercise 3**"),
        _verdict(ex3_scores == 10**12, ex3_scores is not None),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"Scores 14 and $12-2=10$: $\alpha_{12}=e^{10}/(e^{14}+e^{10})=1/(1+e^4)\approx 0.018$."
                ),
                "Solution 2": mo.md("""```python
def my_sdpa(Q, K, V, mask=None):
    S = Q @ K.T / Q.size(-1) ** 0.5
    if mask is not None:
        S = S.masked_fill(~mask, float("-inf"))
    return torch.softmax(S, dim=-1) @ V
```"""),
                "Solution 3": mo.md(
                    r"$n^2=10^{12}$ scores, about 4 TB in float32 per head per layer. Hence sparse / linear attention (Exphormer, NodeFormer) or GPS with sampled global attention at social-network scale."
                ),
            }
        ),
    ]
)