Blame for code/graph/model.py.orig

"""The unified GRM model: encode, reason, select, verify, answer.

Pipeline (PLAN sections 5 and 48)::

    graph --> GraphEncoder --------+
                                   >-- ReasoningCore --> PathSelection --> Verifier --> ReasoningResult
    query --> QueryEncoder --> PAT +

Two behaviours are worth calling out because they were the substantive defects in
the previous version.

``reason`` now *selects*.  It used to return every path the verifier accepted,
which made the answer set "everything reachable within max_hops" and put exact
match out of reach on any task with distractors.  Selection runs through
:class:`~graph.heads.PathSelectionHead` with a learned threshold, so the model
commits to a subset and is measured on it.

``forward_example`` drives PAT and the reasoning core on every training step.
PAT used to be absent from the forward pass entirely (the old
``GRMForwardOutput.pat`` was hard-coded to ``None``), so its only gradient came
from a side call in the trainer.
"""

from __future__ import annotations

import itertools
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, replace
from itertools import islice

import torch
from torch import Tensor, nn

from graph.config import GraphConfig
from graph.constraints import ConstraintState
from graph.graph import GraphStore
from graph.graph_encoder import GraphEncoder, GraphEncoderOutput
from graph.heads import (
    ConfidenceHead,
    ContradictionHead,
    CountHead,
    EntityAnswerHead,
    NextHopHead,
    PathSelectionHead,
    RelationAnswerHead,
    SelectionOutput,
    SupportHead,
    VerificationStack,
)
from graph.pat import PathAwareTransformer, PATOutput
from graph.query_encoder import QueryEncoder, QueryEncoding
from graph.reasoner import BeamSearchReasoner, SearchOutput
from graph.reasoning_core import ReasoningCore
from graph.text_features import build_featuriser
from graph.types import (
    EntityAnswer,
    GraphPath,
    GraphQuery,
    GraphTensors,
    ReasoningResult,
    ReasoningStatus,
)
from graph.verification import (
    classify_evidence,
    detect_conflicting_paths,
    verify_path,
)

STATUS_INDEX: Mapping[str, int] = {
    "supported": 0,
    "unsupported": 1,
    "contradicted": 2,
    "insufficient_evidence": 3,
    "error": 3,
}
INDEX_STATUS: tuple[ReasoningStatus, ...] = (
    ReasoningStatus.SUPPORTED,
    ReasoningStatus.UNSUPPORTED,
    ReasoningStatus.CONTRADICTED,
    ReasoningStatus.INSUFFICIENT_EVIDENCE,
)


@dataclass(frozen=True)
class EncodedGraph:
    """Graph encoding plus the id-keyed views the search and PAT consume."""

    output: GraphEncoderOutput
    node_states: dict[str, Tensor]
    relation_states: dict[str, Tensor]
    memory: Tensor
    memory_mask: Tensor


@dataclass(frozen=True)
class GRMForwardOutput:
    """Every tensor the multi-task loss needs from one example."""

    graph: GraphEncoderOutput
    query: QueryEncoding
    pat: PATOutput | None
    entity_logits: Tensor
    relation_logits: Tensor
    support_logit: Tensor
    contradiction_logits: Tensor
    count_logits: Tensor
    confidence_logit: Tensor
    selection_logits: Tensor
    selection_labels: Tensor
    #: The threshold the model would apply, so the set-level objective can train it.
    selection_threshold: Tensor | None
    next_hop_logits: Tensor
    next_hop_labels: Tensor
    pat_positive: Tensor | None
    pat_negative: Tensor | None
    node_states: Tensor
    entity_order: tuple[str, ...]


class GraphModel(nn.Module):
    """Graph-native reasoning model with an explicit query-conditioned core."""

    def __init__(self, config: GraphConfig) -> None:
        super().__init__()
        self.config = config
        checkpointing = config.training.gradient_checkpointing
        # One featuriser is shared by the query encoder and the graph tensorisation, so
        # a relation string gets the same vector wherever the model meets it.  It holds
        # no trainable parameters, so it is deliberately not a submodule: registering a
        # frozen pretrained encoder would put it in the checkpoint and in the parameter
        # budget, neither of which it belongs in.
        self.featuriser = (
            build_featuriser(
                config.text_features,
                config.text_feature_dim,
                # The graph model builds its featuriser for training as well as
                # evaluation, so it asks for an accelerator when there is one.
                device="auto",
                model_name=config.text_feature_model,
            )
            if config.text_features
            else None
        )
        if self.featuriser is not None and self.featuriser.dimension != (
            config.graph_input_dim
        ):
            raise ValueError(
                f"text_features produces {self.featuriser.dimension}-dimensional "
                f"vectors but graph_input_dim is {config.graph_input_dim}; set them "
                "equal so the encoder receives the features it was configured for"
            )
        self.query_encoder = QueryEncoder(config, self.featuriser)
        self.graph_encoder = GraphEncoder(config)
        self.pat = PathAwareTransformer(
            config.pat,
            gradient_checkpointing=checkpointing,
            logit_cap=config.attention_logit_cap,
            qk_norm=config.attention_qk_norm,
        )
        self.reasoning_core = ReasoningCore(
            config.hidden_dim,
            config.num_reasoning_layers,
            config.num_reasoning_heads,
            dropout=config.dropout,
            max_hops=config.max_hops,
            ffn_multiplier=config.reasoning_ffn_multiplier,
            gradient_checkpointing=checkpointing,
            logit_cap=config.attention_logit_cap,
            qk_norm=config.attention_qk_norm,
        )
        self.next_hop_head = NextHopHead(config.hidden_dim, config.dropout)
        self.path_selection_head = PathSelectionHead(config.hidden_dim, config.dropout)
        self.entity_answer_head = EntityAnswerHead(config.hidden_dim)
        self.relation_answer_head = RelationAnswerHead(
            config.hidden_dim, config.num_relations
        )
        self.verification_stack = VerificationStack(
            config.hidden_dim,
            depth=config.verifier_layers,
            width_multiplier=config.verifier_width_multiplier,
            dropout=config.dropout,
            num_signals=config.verification_signals,
        )
        self.support_head = SupportHead(config.hidden_dim)
        self.contradiction_head = ContradictionHead(config.hidden_dim)
        self.count_head = CountHead(config.hidden_dim, config.max_count)
        self.confidence_head = ConfidenceHead(config.hidden_dim)
        self.reasoner = BeamSearchReasoner(
            self.next_hop_head,
            config.beam_width,
            config.candidate_width,
            config.max_hops,
            config.relation_top_k,
        )

    # -- encoding -----------------------------------------------------------

    def graph_to_example(self, graph: GraphStore):
        """Tensorise ``graph`` at the widths this model's encoder expects.

        Single point of truth for the node/edge feature widths so a config with
        ``graph_input_dim != edge_input_dim`` cannot produce a shape mismatch.
        """
        return graph.to_example(
            feature_dim=self.config.graph_input_dim,
            edge_feature_dim=self.config.edge_input_dim,
            featuriser=self.featuriser,
        )

    def encode_graph(
        self,
        graph_tensors: GraphTensors,
        query_state: Tensor,
        node_ids: Sequence[str],
        relation_ids: Sequence[str],
        *,
        start_ids: Sequence[str],
        use_types: bool = True,
    ) -> EncodedGraph:
        """Encode a graph and expose id-keyed node/relation states.

        ``start_ids`` are the query's start entities, marked in the encoder's
        input when the model was built with ``start_marking``.
        """
        if not use_types:
            # The untyped ablation must blank the identity the encoder actually
            # reads, not only the local bookkeeping index; zeroing just the latter
            # would leave the type embedding fully informed and report the
            # unablated model.
            graph_tensors = replace(
                graph_tensors,
                node_type_ids=torch.zeros_like(graph_tensors.node_type_ids),
                node_type_hashes=torch.zeros_like(graph_tensors.node_type_identity),
            )
        encoded = self.graph_encoder(
            graph_tensors,
            query_state,
            start_mask=_start_mask(node_ids, start_ids, query_state.device),
        )
        node_states = {
            node_id: encoded.node_states[index]
            for index, node_id in enumerate(node_ids)
        }
        relation_bank = self._pool_relation_states(
            encoded.relation_states, graph_tensors.edge_type_ids, len(relation_ids)
        )
        relation_states = {
            relation_id: relation_bank[index]
            for index, relation_id in enumerate(relation_ids)
        }
        memory, memory_mask = self.reasoning_core.build_memory(
            encoded.node_states, relation_bank
        )
        return EncodedGraph(encoded, node_states, relation_states, memory, memory_mask)

    @staticmethod
    def _pool_relation_states(
        edge_states: Tensor, edge_type_ids: Tensor, num_relations: int
    ) -> Tensor:
        """Mean-pool per-edge states into one state per relation type.

        Vectorised with ``index_add_`` rather than accumulated in a Python loop.
        The loop built one autograd node per edge, so a graph with a few hundred
        edges produced a backward chain hundreds of nodes deep and the backward
        pass cost roughly twenty times the forward pass.  Scatter-add is a single
        node and the ratio returns to normal.
        """
        dim = edge_states.shape[-1]
        if not num_relations or edge_states.shape[0] == 0:
            return edge_states.new_zeros((max(num_relations, 0), dim))
        safe_ids = edge_type_ids.clamp(0, num_relations - 1)
        sums = edge_states.new_zeros((num_relations, dim))
        sums = sums.index_add(0, safe_ids, edge_states)
        counts = edge_states.new_zeros((num_relations,))
        counts = counts.index_add(
            0, safe_ids, torch.ones_like(safe_ids, dtype=edge_states.dtype)
        )
        return sums / counts.clamp_min(1.0).unsqueeze(-1)

    # -- training -----------------------------------------------------------

    def forward_example(
        self,
        graph: GraphStore,
        query: GraphQuery,
        *,
        supporting_paths: Sequence[GraphPath],
        competing_paths: Sequence[GraphPath],
        answers: Sequence[str],
        status: ReasoningStatus,
        feature_dim: int | None = None,
    ) -> GRMForwardOutput:
        """Compute every training signal for one example in a single pass."""
        example = self.graph_to_example(graph)
        device = next(self.parameters()).device
        graph_tensors = example.graph.to(device)
        query_encoding = self.query_encoder.encode(query)
        query_state = query_encoding.pooled[0]
        encoded = self.encode_graph(
            graph_tensors,
            query_state,
            example.node_external_ids,
            example.relation_external_ids,
            start_ids=_query_starts(query),
        )
        constraints = ConstraintState.from_query(query)

        positives = tuple(supporting_paths)[: self.config.max_supervised_paths]
        if self.config.on_policy_selection:
            positives, negatives = self._on_policy_candidates(
                graph, query, encoded, query_state, positives, competing_paths
            )
        else:
            negatives = self._selection_negatives(
                graph, query, positives, competing_paths
            )
        labelled = (*positives, *negatives)

        pat_output: PATOutput | None = None
        pat_positive: Tensor | None = None
        pat_negative: Tensor | None = None
        if labelled:
            pat_output = self.pat.encode_paths(
                labelled,
                encoded.node_states,
                encoded.relation_states,
                query_state,
                constraints,
            )
            path_states = pat_output.path_states
            if positives:
                pat_positive = pat_output.compatibility[: len(positives)]
            if negatives:
                pat_negative = pat_output.compatibility[len(positives) :]
            core = self.reasoning_core(
                query_state,
                encoded.memory,
                encoded.memory_mask,
                path_states.unsqueeze(0),
                hop=1,
            )
            selection = self.path_selection_head(
                core.candidate_states[0], path_states, query_state
            )
            selection_logits = selection.logits
            selection_threshold = selection.threshold
            selection_labels = torch.cat(
                (
                    selection_logits.new_ones(len(positives)),
                    selection_logits.new_zeros(len(negatives)),
                )
            )
            reasoning_state = core.state[0]
            evidence_state = path_states.mean(0)
        else:
            zeros = query_state.new_zeros((0,))
            selection_logits, selection_labels = zeros, zeros
            selection_threshold = None
            core = self.reasoning_core(
                query_state,
                encoded.memory,
                encoded.memory_mask,
                query_state.new_zeros((1, 1, query_state.shape[-1])),
                hop=1,
            )
            reasoning_state = core.state[0]
            evidence_state = torch.zeros_like(reasoning_state)

        entity_logits = self.entity_answer_head(
            reasoning_state, encoded.output.node_states
        )
        relation_context = (
            torch.stack(list(encoded.relation_states.values())).mean(0)
            if encoded.relation_states
            else torch.zeros_like(reasoning_state)
        )
        relation_logits = self.relation_answer_head(
            reasoning_state.unsqueeze(0), relation_context.unsqueeze(0)
        )
        verified_state = self.verification_stack(
            reasoning_state.unsqueeze(0),
            evidence_state.unsqueeze(0),
            query_state.unsqueeze(0),
            self._evidence_signals(
                pat_output.compatibility if pat_output is not None else None,
                len(labelled),
                query_state.device,
            ),
        )
        support_logit = self.support_head(verified_state)
        contradiction_logits = self.contradiction_head(verified_state)
        count_logits = self.count_head(verified_state)
        selected_fraction = (
            (selection_logits > selection_threshold).float().mean()
            if selection_logits.numel() and selection_threshold is not None
            else torch.zeros((), device=device)
        )
        confidence_logit = self.confidence_head(
            verified_state,
            support_logit.detach().reshape(1),
            selected_fraction.reshape(1),
        )
        next_hop_logits, next_hop_labels = self._next_hop_supervision(
            graph,
            query,
            positives,
            encoded,
            core.candidate_states[0] if labelled else None,
        )
        return GRMForwardOutput(
            graph=encoded.output,
            query=query_encoding,
            pat=pat_output,
            entity_logits=entity_logits,
            relation_logits=relation_logits,
            support_logit=support_logit,
            contradiction_logits=contradiction_logits,
            count_logits=count_logits,
            confidence_logit=confidence_logit,
            selection_logits=selection_logits,
            selection_labels=selection_labels,
            selection_threshold=selection_threshold,
            next_hop_logits=next_hop_logits,
            next_hop_labels=next_hop_labels,
            pat_positive=pat_positive,
            pat_negative=pat_negative,
            node_states=encoded.output.node_states,
            entity_order=example.node_external_ids,
        )

    def forward_batch(
        self,
        graphs: Sequence[GraphStore],
        queries: Sequence[GraphQuery],
        supporting_paths: Sequence[Sequence[GraphPath]],
        competing_paths: Sequence[Sequence[GraphPath]],
    ) -> list[GRMForwardOutput]:
        """:meth:`forward_example` for a batch, with each network run once.

        Per-example training launched every layer once per example: graphs have
        different shapes, and the model's layers are small, so a step was bound by
        kernel launches and Python overhead with the GPU about 15% busy.  Here the
        graphs are encoded as one disjoint union, the queries, PAT paths and
        reasoning tokens are padded and masked, and the heads run over flattened
        rows.  Each example's outputs equal :meth:`forward_example`'s up to
        floating-point reordering (and dropout draws); the equivalence is tested.
        """
        if self.config.on_policy_selection:
            # The on-policy search is sequential per example; nothing to batch.
            raise ValueError("forward_batch does not support on_policy_selection")
        count = len(graphs)
        if not count:
            return []
        device = next(self.parameters()).device
        hidden = self.config.hidden_dim

        # -- tensorise (text features encoded in one call) ----------------------
        if self.featuriser is not None:
            self.featuriser.prefetch_identifiers(
                [
                    text
                    for graph in graphs
                    for text in (
                        *(
                            f"{node_id} {graph.nodes[node_id]}"
                            for node_id in graph.nodes
                        ),
                        *(
                            f"{edge.relation_id} {edge.relation_name or ''}"
                            for edge in graph.edges
                        ),
                    )
                ]
            )
        examples = [self.graph_to_example(graph) for graph in graphs]
        tensors = [example.graph for example in examples]
        node_counts = [item.num_nodes for item in tensors]
        edge_counts = [item.num_edges for item in tensors]
        relation_counts = [len(example.relation_external_ids) for example in examples]
        node_starts = [0, *itertools.accumulate(node_counts)]
        edge_starts = [0, *itertools.accumulate(edge_counts)]
        relation_starts = [0, *itertools.accumulate(relation_counts)]
        union = GraphTensors(
            node_ids=torch.arange(node_starts[-1], dtype=torch.long),
            node_type_ids=torch.cat([item.node_type_ids for item in tensors]),
            node_features=torch.cat([item.node_features for item in tensors]),
            edge_index=torch.cat(
                [
                    item.edge_index + start
                    for item, start in zip(tensors, node_starts, strict=False)
                ],
                dim=1,
            ),
            edge_type_ids=torch.cat(
                [
                    item.edge_type_ids + start
                    for item, start in zip(tensors, relation_starts, strict=False)
                ]
            ),
            edge_features=torch.cat([item.edge_features for item in tensors]),
            graph_id=torch.repeat_interleave(
                torch.arange(count), torch.tensor(node_counts)
            ),
            node_type_hashes=torch.cat([item.node_type_identity for item in tensors]),
            edge_type_hashes=torch.cat([item.edge_type_identity for item in tensors]),
        ).to(device)
        edge_owner = torch.repeat_interleave(
            torch.arange(count, device=device),
            torch.tensor(edge_counts, device=device),
        )
        node_owner = union.graph_id
        relation_owner = torch.repeat_interleave(
            torch.arange(count, device=device),
            torch.tensor(relation_counts, device=device),
        )

        # -- query and graph encoding ------------------------------------------
        query_encodings = self.query_encoder.encode_batch(queries)
        query_states = torch.cat([encoding.pooled for encoding in query_encodings])
        encoded = self.graph_encoder(
            union,
            query_states,
            edge_query=query_states[edge_owner],
            start_mask=torch.cat(
                [
                    _start_mask(ex.node_external_ids, _query_starts(query), device)
                    for ex, query in zip(examples, queries, strict=True)
                ]
            ),
            node_query=query_states[node_owner],
        )
        # ``edge_type_ids`` were offset per graph, so pooling over the union
        # yields every graph's relation bank, stacked.
        relation_bank = self._pool_relation_states(
            encoded.relation_states, union.edge_type_ids, relation_starts[-1]
        )
        node_rows = [
            {
                node_id: node_starts[i] + j
                for j, node_id in enumerate(ex.node_external_ids)
            }
            for i, ex in enumerate(examples)
        ]
        relation_rows = [
            {
                relation_id: relation_starts[i] + j
                for j, relation_id in enumerate(ex.relation_external_ids)
            }
            for i, ex in enumerate(examples)
        ]
        memory, memory_mask = self._batch_memory(
            encoded.node_states,
            relation_bank,
            node_counts,
            relation_counts,
            node_starts,
            relation_starts,
        )

        # -- candidate paths through PAT ---------------------------------------
        positives_per: list[tuple[GraphPath, ...]] = []
        negatives_per: list[tuple[GraphPath, ...]] = []
        for graph, query, supporting, competing in zip(
            graphs, queries, supporting_paths, competing_paths, strict=True
        ):
            positives = tuple(supporting)[: self.config.max_supervised_paths]
            positives_per.append(positives)
            negatives_per.append(
                self._selection_negatives(graph, query, positives, competing)
            )
        labelled_per = [
            (*positives, *negatives)
            for positives, negatives in zip(positives_per, negatives_per, strict=True)
        ]
        path_counts = [len(paths) for paths in labelled_per]
        owners = [i for i, paths in enumerate(labelled_per) for _ in paths]
        path_starts = [0, *itertools.accumulate(path_counts)]
        # Example and within-example position of every labelled path.
        path_owner = torch.tensor(owners, dtype=torch.long, device=device)
        slots = torch.tensor(
            [j for n in path_counts for j in range(n)], dtype=torch.long, device=device
        )
        pat_output: PATOutput | None = None
        if owners:
            pat_input = self.pat.build_batch_input(
                [path for paths in labelled_per for path in paths],
                owners,
                node_rows,
                relation_rows,
                encoded.node_states,
                relation_bank,
                query_states,
                [ConstraintState.from_query(query) for query in queries],
            )
            pat_output = self.pat(pat_input, query_states[path_owner])

        # -- reasoning core over each example's candidates ---------------------
        width = max(1, *path_counts)
        candidates = query_states.new_zeros((count, width, hidden))
        candidate_mask = torch.zeros((count, width), dtype=torch.long, device=device)
        # An example without candidates is given one zero candidate, exactly as
        # forward_example does.
        candidate_mask[:, 0] = 1
        if pat_output is not None:
            candidates = candidates.index_put(
                (path_owner, slots), pat_output.path_states.to(candidates.dtype)
            )
            candidate_mask[path_owner, slots] = 1
        core = self.reasoning_core(
            query_states,
            memory,
            memory_mask,
            candidates,
            hop=1,
            candidate_mask=candidate_mask,
        )
        reasoning_states = core.state

        # -- heads -------------------------------------------------------------
        path_totals = torch.tensor(path_counts, device=device).clamp_min(1)
        selection_logits_all: Tensor | None = None
        thresholds: Tensor | None = None
        evidence = torch.zeros_like(reasoning_states)
        if pat_output is not None:
            path_states = pat_output.path_states
            candidate_rows = core.candidate_states[path_owner, slots]
            head = self.path_selection_head
            selection_logits_all = head.score(
                torch.cat(
                    (candidate_rows, path_states, query_states[path_owner]), dim=-1
                )
            ).squeeze(-1)
            evidence = (
                path_states.new_zeros((count, path_states.shape[-1])).index_add(
                    0, path_owner, path_states
                )
                / path_totals.unsqueeze(-1).to(path_states.dtype)
            ).to(reasoning_states.dtype)
            thresholds = head.threshold + head.threshold_offset(
                torch.cat((query_states, evidence.to(query_states.dtype)), dim=-1)
            ).squeeze(-1)

        projected = self.entity_answer_head.project(reasoning_states)
        entity_logits_all = self.entity_answer_head.scale * (
            encoded.node_states * projected[node_owner]
        ).sum(-1)
        relation_context = relation_bank.new_zeros((count, hidden)).index_add(
            0, relation_owner, relation_bank
        ) / torch.tensor(relation_counts, device=device).clamp_min(1).unsqueeze(-1).to(
            relation_bank.dtype
        )
        relation_logits = self.relation_answer_head(
            reasoning_states, relation_context.to(reasoning_states.dtype)
        )
        signals = self._batch_evidence_signals(pat_output, owners, path_counts, device)
        verified = self.verification_stack(
            reasoning_states, evidence, query_states, signals
        )
        support_logits = self.support_head(verified)
        contradiction_logits = self.contradiction_head(verified)
        count_logits = self.count_head(verified)
        if selection_logits_all is not None and thresholds is not None:
            kept = (selection_logits_all > thresholds[path_owner]).to(verified.dtype)
            selected_fraction = verified.new_zeros((count,)).index_add(
                0, path_owner, kept
            ) / path_totals.to(verified.dtype)
        else:
            selected_fraction = verified.new_zeros((count,))
        confidence_logits = self.confidence_head.score(
            torch.cat(
                (
                    verified,
                    support_logits.detach().unsqueeze(-1),
                    selected_fraction.unsqueeze(-1),
                ),
                dim=-1,
            )
        ).squeeze(-1)
        next_hop = self._batch_next_hop_supervision(
            graphs,
            positives_per,
            path_counts,
            node_rows,
            relation_rows,
            encoded.node_states,
            relation_bank,
            core.candidate_states,
        )

        # -- per-example views -------------------------------------------------
        outputs: list[GRMForwardOutput] = []
        for i in range(count):
            nodes = slice(node_starts[i], node_starts[i + 1])
            edges = slice(edge_starts[i], edge_starts[i + 1])
            paths = slice(path_starts[i], path_starts[i + 1])
            n_pos = len(positives_per[i])
            n_paths = path_counts[i]
            if n_paths and selection_logits_all is not None and thresholds is not None:
                selection_logits = selection_logits_all[paths]
                selection_labels = torch.cat(
                    (
                        selection_logits.new_ones(n_pos),
                        selection_logits.new_zeros(n_paths - n_pos),
                    )
                )
                threshold: Tensor | None = thresholds[i]
            else:
                selection_logits = query_states.new_zeros((0,))
                selection_labels = selection_logits
                threshold = None
            compatibility = (
                pat_output.compatibility[paths]
                if pat_output is not None and n_paths
                else None
            )
            outputs.append(
                GRMForwardOutput(
                    graph=GraphEncoderOutput(
                        encoded.node_states[nodes],
                        encoded.relation_states[edges],
                        encoded.edge_states[edges],
                        None,
                    ),
                    query=query_encodings[i],
                    pat=_slice_pat(pat_output, paths) if n_paths else None,
                    entity_logits=entity_logits_all[nodes],
                    relation_logits=relation_logits[i : i + 1],
                    support_logit=support_logits[i : i + 1],
                    contradiction_logits=contradiction_logits[i : i + 1],
                    count_logits=count_logits[i : i + 1],
                    confidence_logit=confidence_logits[i : i + 1],
                    selection_logits=selection_logits,
                    selection_labels=selection_labels,
                    selection_threshold=threshold,
                    next_hop_logits=next_hop[i][0],
                    next_hop_labels=next_hop[i][1],
                    pat_positive=compatibility[:n_pos]
                    if compatibility is not None and n_pos
                    else None,
                    pat_negative=compatibility[n_pos:]
                    if compatibility is not None and n_paths > n_pos
                    else None,
                    node_states=encoded.node_states[nodes],
                    entity_order=examples[i].node_external_ids,
                )
            )
        return outputs

    def _batch_memory(
        self,
        node_states: Tensor,
        relation_bank: Tensor,
        node_counts: Sequence[int],
        relation_counts: Sequence[int],
        node_starts: Sequence[int],
        relation_starts: Sequence[int],
    ) -> tuple[Tensor, Tensor]:
        """Padded ``[B, M, D]`` graph memory: each example's nodes, then relations.

        Matches :meth:`ReasoningCore.build_memory` row for row; the padding is
        masked out of the cross-attention.
        """
        core = self.reasoning_core
        device = node_states.device
        count = len(node_counts)
        lengths = [n + r for n, r in zip(node_counts, relation_counts, strict=True)]
        width = max(lengths)
        roles = core.memory_role_embedding(torch.arange(2, device=device))
        nodes = core.memory_projection(node_states) + roles[0]
        destination = [
            i * width + j for i, n in enumerate(node_counts) for j in range(n)
        ]
        rows = [nodes]
        if relation_bank.shape[0]:
            rows.append(core.memory_projection(relation_bank) + roles[1])
            destination += [
                i * width + node_counts[i] + j
                for i, r in enumerate(relation_counts)
                for j in range(r)
            ]
        values = torch.cat(rows)
        memory = values.new_zeros((count * width, values.shape[-1])).index_copy(
            0, torch.tensor(destination, dtype=torch.long, device=device), values
        )
        mask = torch.zeros((count, width), dtype=torch.long)
        for i, length in enumerate(lengths):
            mask[i, :length] = 1
        return memory.view(count, width, -1), mask.to(device)

    def _batch_evidence_signals(
        self,
        pat_output: PATOutput | None,
        owners: Sequence[int],
        path_counts: Sequence[int],
        device: torch.device,
    ) -> Tensor | None:
        """:meth:`_evidence_signals` for every example in a batch."""
        if not self.config.verification_signals:
            return None
        count = len(path_counts)
        maximum = torch.zeros((count,), device=device)
        mean = torch.zeros((count,), device=device)
        if pat_output is not None:
            owner = torch.tensor(owners, dtype=torch.long, device=device)
            compatibility = pat_output.compatibility.reshape(-1).float()
            maximum = maximum.scatter_reduce(
                0, owner, compatibility, reduce="amax", include_self=False
            )
            mean = (
                mean.index_add(0, owner, compatibility)
                / torch.tensor(path_counts, device=device).clamp_min(1).float()
            )
        sizes = torch.log1p(torch.tensor(path_counts, device=device).float()) / 10.0
        signals = torch.stack((maximum, mean, sizes), dim=-1)
        return signals[:, : self.config.verification_signals]

    def _batch_next_hop_supervision(
        self,
        graphs: Sequence[GraphStore],
        positives_per: Sequence[Sequence[GraphPath]],
        path_counts: Sequence[int],
        node_rows: Sequence[Mapping[str, int]],
        relation_rows: Sequence[Mapping[str, int]],
        node_states: Tensor,
        relation_bank: Tensor,
        candidate_states: Tensor,
    ) -> list[tuple[Tensor, Tensor]]:
        """:meth:`_next_hop_supervision` for a batch, scored in one head call."""
        device = node_states.device
        empty = torch.zeros((0,), device=device)
        width = candidate_states.shape[1]
        # Row 0 of each lookup table is the zero vector a missing id falls back to.
        state_index: list[int] = []
        source_index: list[int] = []
        relation_index: list[int] = []
        target_index: list[int] = []
        labels: list[float] = []
        spans: list[tuple[int, int] | None] = []
        for i, (graph, positives) in enumerate(zip(graphs, positives_per, strict=True)):
            if not positives or not positives[0].steps or not path_counts[i]:
                spans.append(None)
                continue
            gold = positives[0].steps[0]
            siblings = graph.outgoing(gold.source_id)
            if len(siblings) < 2:
                spans.append(None)
                continue
            start = len(labels)
            nodes, relations = node_rows[i], relation_rows[i]
            for edge in siblings[: self.config.candidate_width]:
                state_index.append(i * width)
                source = nodes.get(edge.source_id)
                relation = relations.get(edge.relation_id)
                target = nodes.get(edge.target_id)
                source_index.append(0 if source is None else source + 1)
                relation_index.append(0 if relation is None else relation + 1)
                target_index.append(0 if target is None else target + 1)
                labels.append(
                    float(
                        edge.relation_id == gold.relation_id
                        and edge.target_id == gold.target_id
                    )
                )
            spans.append((start, len(labels)))
        if not labels:
            return [(empty, empty) for _ in graphs]

        def table(values: Tensor) -> Tensor:
            return torch.cat((values.new_zeros((1, values.shape[-1])), values))

        def rows(values: Tensor, index: list[int]) -> Tensor:
            return values[torch.tensor(index, dtype=torch.long, device=device)]

        node_table = table(node_states)
        logits = self.next_hop_head(
            rows(candidate_states.reshape(-1, candidate_states.shape[-1]), state_index),
            rows(node_table, source_index),
            rows(table(relation_bank), relation_index),
            rows(node_table, target_index),
        )
        label_tensor = torch.tensor(labels, device=device)
        return [
            (empty, empty)
            if span is None
            else (logits[span[0] : span[1]], label_tensor[span[0] : span[1]])
            for span in spans
        ]

    def _selection_negatives(
        self,
        graph: GraphStore,
        query: GraphQuery,
        positives: Sequence[GraphPath],
        competing_paths: Sequence[GraphPath],
    ) -> tuple[GraphPath, ...]:
        """Negatives drawn from the paths the search will actually surface.

        Training on the generator's curated ``competing_paths`` alone leaves a
        train/inference mismatch: those are same-length routes from the same start,
        while the beam returns paths of *every* admissible length.  Worse, an
        example with no competing paths has an all-positive label vector, so both
        the per-path and set-level objectives are minimised by keeping everything —
        and the measured result was a model that selected all 14 candidates when
        the gold set averaged 2.

        Curated competitors come first, because they are the hardest negatives;
        the remainder is filled from the reachable-path distribution the beam sees.
        """
        budget = max(
            self.config.max_supervised_paths,
            len(positives) * self.config.selection_negative_ratio,
        )
        gold = {_path_key(path) for path in positives}
        chosen: list[GraphPath] = []
        seen: set[tuple[tuple[str, str, str], ...]] = set()
        for path in competing_paths:
            key = _path_key(path)
            if key in gold or key in seen:
                continue
            seen.add(key)
            chosen.append(path)
            if len(chosen) >= budget:
                return tuple(chosen)
        starts = query.start_entity_ids or query.start_entities
        horizon = min(query.max_hops, self.config.max_hops)
        for path in islice(
            graph.iter_paths(starts, horizon), self.config.reachable_negative_pool
        ):
            if path.hop_count < query.min_hops:
                continue
            key = _path_key(path)
            if key in gold or key in seen:
                continue
            seen.add(key)
            chosen.append(path)
            if len(chosen) >= budget:
                break
        return tuple(chosen)

    def _on_policy_candidates(
        self,
        graph: GraphStore,
        query: GraphQuery,
        encoded: EncodedGraph,
        query_state: Tensor,
        positives: Sequence[GraphPath],
        competing_paths: Sequence[GraphPath],
    ) -> tuple[tuple[GraphPath, ...], tuple[GraphPath, ...]]:
        """Label the candidates the search itself produces (on-policy selection).

        Training on curated negatives leaves the selection head facing a different
        distribution at inference than the one it was fitted on.  The beam picks
        candidates with ``next_hop_head``, which is trained to rank the gold *first*
        hop highly, so its output is dominated by paths that start correctly and
        diverge later -- precisely the negatives an enumeration-order sample misses.
        Off-policy training therefore produced a threshold that kept ten candidates
        where the gold set held two, no matter how the base rate was balanced.

        The search runs without gradients (path enumeration needs none); the chosen
        paths are then re-encoded with gradients by the caller.  Gold paths the beam
        missed are appended, because a batch with no positive carries no signal for
        what to keep.
        """
        budget = self.config.max_supervised_paths
        gold_keys = {_path_key(path) for path in positives}
        with torch.no_grad():
            search = self.reasoner.search(
                graph,
                query,
                query_state,
                encoded.node_states,
                encoded.relation_states,
                self.pat,
                self.reasoning_core,
                encoded.memory,
                encoded.memory_mask,
                use_pat=self.config.pat.enabled,
            )
        found_positive: list[GraphPath] = []
        found_negative: list[GraphPath] = []
        for path in search.paths:
            if _path_key(path) in gold_keys:
                found_positive.append(path)
            else:
                found_negative.append(path)
        # Keep every gold path so the objective always has something to keep.
        missed = [
            path
            for path in positives
            if _path_key(path) not in {_path_key(p) for p in found_positive}
        ]
        chosen_positive = (*found_positive[:budget], *missed[:budget])[:budget]
        if not chosen_positive:
            chosen_positive = tuple(positives[:budget])
        limit = max(budget, len(chosen_positive) * self.config.selection_negative_ratio)
        chosen_negative = tuple(found_negative[:limit])
        if not chosen_negative:
            chosen_negative = self._selection_negatives(
                graph, query, chosen_positive, competing_paths
            )
        return tuple(chosen_positive), chosen_negative

    def _next_hop_supervision(
        self,
        graph: GraphStore,
        query: GraphQuery,
        positives: Sequence[GraphPath],
        encoded: EncodedGraph,
        candidate_states: Tensor | None,
    ) -> tuple[Tensor, Tensor]:
        """Label the first transition of the gold path against its siblings.

        This is the signal that teaches the beam which relation to take out of
        the start entity -- the decision the leaked composition used to make.
        """
        device = encoded.output.node_states.device
        if not positives or not positives[0].steps or candidate_states is None:
            empty = torch.zeros((0,), device=device)
            return empty, empty
        gold = positives[0].steps[0]
        start = gold.source_id
        siblings = graph.outgoing(start)
        if len(siblings) < 2:
            empty = torch.zeros((0,), device=device)
            return empty, empty
        zero = torch.zeros_like(encoded.output.node_states[0])
        state = candidate_states[0]
        rows, labels = [], []
        for edge in siblings[: self.config.candidate_width]:
            rows.append(
                (
                    state,
                    encoded.node_states.get(edge.source_id, zero),
                    encoded.relation_states.get(edge.relation_id, zero),
                    encoded.node_states.get(edge.target_id, zero),
                )
            )
            labels.append(
                float(
                    edge.relation_id == gold.relation_id
                    and edge.target_id == gold.target_id
                )
            )
        logits = self.next_hop_head(
            torch.stack([row[0] for row in rows]),
            torch.stack([row[1] for row in rows]),
            torch.stack([row[2] for row in rows]),
            torch.stack([row[3] for row in rows]),
        )
        return logits, torch.tensor(labels, device=device)

    # -- inference ----------------------------------------------------------

    @torch.no_grad()
    def reason(
        self,
        graph: GraphStore,
        query: GraphQuery,
        *,
        use_pat: bool | None = None,
        verify: bool | None = None,
        beam_width: int | None = None,
        use_types: bool = True,
        query_conditioned: bool = True,
        select: bool = True,
        abstain: bool = True,
    ) -> ReasoningResult:
        """Answer ``query`` against ``graph`` and return a verified typed result.

        ``abstain=False`` disables the learned support gate, leaving the deterministic
        verifier as the only thing that can empty the answer set.  It exists because the
        gate is distribution-sensitive: its evidence signals come from PAT
        query-compatibility, which saturates low on unseen relation vocabulary, and on
        real typed KGs it fired on every query and drove every metric to 0.000.
        """
        was_training = self.training
        self.eval()
        try:
            return self._reason(
                graph,
                query,
                use_pat=self.config.pat.enabled if use_pat is None else use_pat,
                verify_enabled=query.verify if verify is None else verify,
                beam_width=beam_width,
                use_types=use_types,
                query_conditioned=query_conditioned,
                select=select,
                abstain=abstain,
            )
        finally:
            if was_training:
                self.train()

    def _reason(
        self,
        graph: GraphStore,
        query: GraphQuery,
        *,
        use_pat: bool,
        verify_enabled: bool,
        beam_width: int | None,
        use_types: bool,
        query_conditioned: bool,
        select: bool,
        abstain: bool = True,
    ) -> ReasoningResult:
        example = self.graph_to_example(graph)
        device = next(self.parameters()).device
        graph_tensors = example.graph.to(device)
        query_encoding = self.query_encoder.encode(query)
        query_state = query_encoding.pooled[0]
        if not query_conditioned:
            # The ablation must remove query conditioning from *every* path the
            # query takes -- graph encoding, the reasoning core and path
            # selection.  Zeroing it in only one of them leaves the other two
            # conditioned and the ablation reports the unablated model.
            query_state = torch.zeros_like(query_state)
        encoded = self.encode_graph(
            graph_tensors,
            query_state,
            example.node_external_ids,
            example.relation_external_ids,
            use_types=use_types,
            start_ids=_query_starts(query),
        )
        search = self.reasoner.search(
            graph,
            query,
            query_state,
            encoded.node_states,
            encoded.relation_states,
            self.pat,
            self.reasoning_core,
            encoded.memory,
            encoded.memory_mask,
            use_pat=use_pat,
            beam_width=beam_width,
            query_conditioned=query_conditioned,
        )

        candidates = search.paths
        selection_scores: list[float] = []
        kept: tuple[GraphPath, ...] = candidates
        predicted_count: int | None = None
        if candidates and search.candidate_states is not None:
            selection = self.path_selection_head(
                search.candidate_states, search.path_states, query_state
            )
            selection_scores = selection.logits.detach().cpu().tolist()
            scored = tuple(
                replace(
                    path,
                    score=float(score),
                    query_compatibility=_compat(search.pat_output, index),
                )
                for index, (path, score) in enumerate(
                    zip(candidates, selection_scores, strict=True)
                )
            )
            if select:
                # Abstention is decided first and independently of the selection
                # rule.  Count-guided selection always returns at least one path,
                # so gating inside it would make INSUFFICIENT_EVIDENCE unreachable
                # exactly as the old unconditional fallback did.  Whether the graph
                # supports *any* answer is the verifier's question (PLAN 19/20);
                # how many paths to commit to is the selection rule's.
                if abstain and not self._supports_any_answer(search, query_state):
                    kept = ()
                else:
                    predicted_count = self._predicted_answer_count(search, query_state)
                    kept = self._select_paths(scored, selection, predicted_count)
                    if not kept:
                        kept = (max(scored, key=lambda item: item.score),)
            else:
                kept = scored

        verified = (
            tuple(verify_path(path, graph, query) for path in kept)
            if verify_enabled
            else kept
        )
        supporting = tuple(
            path for path in verified if path.verified or not verify_enabled
        )
        conflicts = tuple(
            path
            for group in detect_conflicting_paths(supporting, graph)
            for path in group
        )
        status = self._status(
            supporting,
            conflicts,
            graph,
            query,
            verify_enabled,
            search,
            encoded,
            query_state,
        )
        answers_by_id: dict[str, float] = {}
        for path in supporting:
            if path.endpoint_id is not None:
                answers_by_id[path.endpoint_id] = max(
                    answers_by_id.get(path.endpoint_id, float("-inf")), path.score
                )
        ranked = sorted(answers_by_id.items(), key=lambda item: item[1], reverse=True)
        answers = tuple(
            EntityAnswer(entity_id, float(score))
            for entity_id, score in ranked[: self.config.answer_top_k]
        )
        relation_answers = tuple(
            dict.fromkeys(
                path.steps[-1].relation_id for path in supporting if path.steps
            )
        )
        boolean_answer, aggregate_value = self._mode_answers(
            query, status, answers_by_id, supporting
        )
        confidence = self._confidence(supporting, status)
        if (
            query.target_entity_ids
            and status is ReasoningStatus.SUPPORTED
            and not set(query.target_entity_ids) & set(answers_by_id)
        ):
            status = ReasoningStatus.UNSUPPORTED
            boolean_answer = False if query.answer_mode == "boolean" else boolean_answer
        return ReasoningResult(
            status=status,
            answers=answers
            if status is ReasoningStatus.SUPPORTED and query.answer_mode == "entity"
            else (),
            supporting_paths=supporting if query.return_paths else (),
            conflicting_paths=conflicts if query.return_paths else (),
            confidence=confidence,
            calibrated=False,
            model_version=self.config.model_version,
            graph_version=graph.graph_version,
            reasoning_steps=max((path.hop_count for path in supporting), default=0),
            metadata={
                "pat_enabled": use_pat,
                "verifier_enabled": verify_enabled,
                "selection_enabled": select,
                "query_conditioned": query_conditioned,
                "typed_graph": use_types,
                "beam_width": beam_width or self.config.beam_width,
                "candidates_considered": len(candidates),
                "candidates_selected": len(kept),
                "answers_committed": len(
                    {path.endpoint_id for path in kept if path.endpoint_id}
                ),
                "predicted_answer_count": predicted_count,
                # ``max_count`` bounds what the count head can express, and real KG
                # queries can have more true answers than that.  At the cap the
                # prediction is a floor, not an estimate, and count-guided selection
                # truncates the commit.  Surfacing it distinguishes "the model
                # thinks there are 32 answers" from "it cannot count past 32".
                "count_head_saturated": predicted_count is not None
                and predicted_count >= self.config.max_count,
                "reasoning_state_count": len(search.reasoning_states),
                "pat_path_scores": [path.query_compatibility for path in supporting],
            },
            relation_answers=relation_answers
            if query.answer_mode == "relation"
            else (),
            boolean_answer=boolean_answer,
            aggregate_value=aggregate_value,
        )

    def _select_paths(
        self,
        scored: Sequence[GraphPath],
        selection: SelectionOutput,
        predicted: int | None,
    ) -> tuple[GraphPath, ...]:
        """Choose the committed subset of candidate paths.

        Two rules are available.  A learned threshold is the obvious one, but it is
        fitted on the training candidate distribution and does not transfer: it kept
        roughly nine paths where the gold set held two, which is what capped exact
        match and made a wider beam score worse than greedy.

        Count-guided selection instead asks the count head how many answers the
        query has and takes that many highest-scoring paths.  Cardinality then comes
        from a head trained directly on ``len(answers)`` rather than from a decision
        boundary that has to be calibrated against a distribution it never saw.
        """
        threshold = float(selection.threshold.item())
        if not self.config.count_guided_selection or predicted is None:
            return tuple(path for path in scored if path.score > threshold)
        ranked = sorted(scored, key=lambda item: item.score, reverse=True)
        # Paths are the evidence, answers are the endpoints: keep ranked paths
        # until the requested number of distinct endpoints is covered, so several
        # routes to one answer are not charged against the budget.  A path with no
        # endpoint names no answer, so it is evidence for nothing and is dropped
        # rather than kept free of charge against that budget.
        kept: list[GraphPath] = []
        endpoints: set[str] = set()
        for path in ranked:
            endpoint = path.endpoint_id
            if endpoint is None:
                continue
            if endpoint not in endpoints:
                if len(endpoints) >= predicted:
                    break
                endpoints.add(endpoint)
            kept.append(path)
        return tuple(kept)

    def _predicted_answer_count(
        self, search: object, query_state: Tensor
    ) -> int | None:
        """How many answers the count head expects, clamped to a sane range."""
        pooled = getattr(search, "pooled_state", None)
        path_states = getattr(search, "path_states", None)
        if pooled is None:
            return None
        evidence = (
            path_states.mean(0)
            if path_states is not None and path_states.numel()
            else torch.zeros_like(pooled)
        )
        verified = self.verification_stack(
            pooled.unsqueeze(0),
            evidence.unsqueeze(0),
            query_state.unsqueeze(0),
            self._evidence_signals(
                getattr(getattr(search, "pat_output", None), "compatibility", None),
                len(getattr(search, "paths", ()) or ()),
                query_state.device,
            ),
        )
        logits = self.count_head(verified).reshape(-1)
        return max(1, int(logits.argmax().item()))

    def _evidence_signals(
        self, compatibility: Tensor | None, count: int, device: torch.device
    ) -> Tensor | None:
        """Scalar summaries of how well *any* candidate matches the query.

        Computed the same way from training candidates and from beam candidates, which
        is the whole point: the pooled evidence mean the trunk otherwise sees comes
        from different distributions in the two regimes, and the support head learned
        a rule that inverted at inference.  PAT query-compatibility does not have that
        problem -- it is the same function of a path and a query wherever it is called,
        and it is already trained to discriminate (0.959).

        The maximum is the informative one: "is there any candidate compatible with
        this query at all" is exactly the abstention question.
        """
        if not self.config.verification_signals:
            return None
        if compatibility is None or compatibility.numel() == 0:
            values = [
                torch.zeros((), device=device),
                torch.zeros((), device=device),
            ]
        else:
            flat = compatibility.reshape(-1).float()
            values = [flat.max(), flat.mean()]
        # Candidate count, compressed: "the beam found nothing" is itself evidence.
        values.append(torch.log1p(torch.tensor(float(count), device=device)) / 10.0)
        signals = torch.stack(values)[: self.config.verification_signals]
        return signals.reshape(1, -1)

    def _supports_any_answer(self, search: SearchOutput, query_state: Tensor) -> bool:
        """Neural verification of whether the graph supports any answer.

        Consumes the pooled reasoning state, the pooled candidate evidence and
        the query, exactly as the training-time support objective does, so the
        head is used at inference in the same shape it was trained in.
        """
        pooled = getattr(search, "pooled_state", None)
        path_states = getattr(search, "path_states", None)
        if pooled is None:
            return True
        evidence = (
            path_states.mean(0)
            if path_states is not None and path_states.numel()
            else torch.zeros_like(pooled)
        )
        verified = self.verification_stack(
            pooled.unsqueeze(0),
            evidence.unsqueeze(0),
            query_state.unsqueeze(0),
            self._evidence_signals(
                getattr(getattr(search, "pat_output", None), "compatibility", None),
                len(getattr(search, "paths", ()) or ()),
                query_state.device,
            ),
        )
        probability = torch.sigmoid(self.support_head(verified)).reshape(-1)[0]
        return bool(probability.item() > self.config.abstention_threshold)

    def _status(
        self,
        supporting: Sequence[GraphPath],
        conflicts: Sequence[GraphPath],
        graph: GraphStore,
        query: GraphQuery,
        verify_enabled: bool,
        search: object,
        encoded: EncodedGraph,
        query_state: Tensor,
    ) -> ReasoningStatus:
        """Combine the neural status classifier with deterministic evidence.

        The classifier may propose ``CONTRADICTED``; it may never promote an
        answer the verifier rejected, which is why the deterministic branch wins
        whenever there is no verified support (PLAN section 20).
        """
        if not supporting:
            if query.target_entity_ids or query.answer_mode == "boolean":
                return ReasoningStatus.UNSUPPORTED
            return ReasoningStatus.INSUFFICIENT_EVIDENCE
        if conflicts:
            return ReasoningStatus.CONTRADICTED
        pooled = getattr(search, "pooled_state", None)
        if pooled is not None:
            evidence = getattr(search, "path_states", None)
            evidence_state = (
                evidence.mean(0)
                if evidence is not None and evidence.numel()
                else torch.zeros_like(pooled)
            )
            verified_state = self.verification_stack(
                pooled.unsqueeze(0),
                evidence_state.unsqueeze(0),
                query_state.unsqueeze(0),
                self._evidence_signals(
                    getattr(getattr(search, "pat_output", None), "compatibility", None),
                    len(getattr(search, "paths", ()) or ()),
                    query_state.device,
                ),
            )
            logits = self.contradiction_head(verified_state)
            predicted = INDEX_STATUS[int(logits.argmax(dim=-1).item())]
            if predicted is ReasoningStatus.CONTRADICTED:
                return ReasoningStatus.CONTRADICTED
        if verify_enabled:
            return classify_evidence(
                candidate_paths=tuple(supporting), graph=graph, query=query
            )
        return ReasoningStatus.SUPPORTED

    def _mode_answers(
        self,
        query: GraphQuery,
        status: ReasoningStatus,
        answers_by_id: Mapping[str, float],
        supporting: Sequence[GraphPath],
    ) -> tuple[bool | None, int | float | None]:
        boolean_answer: bool | None = None
        aggregate_value: int | float | None = None
        if query.answer_mode == "boolean":
            if status is ReasoningStatus.SUPPORTED:
                boolean_answer = not query.target_entity_ids or bool(
                    set(query.target_entity_ids) & set(answers_by_id)
                )
            elif status in {ReasoningStatus.UNSUPPORTED, ReasoningStatus.CONTRADICTED}:
                boolean_answer = False
        elif query.answer_mode in {"count", "aggregate"}:
            aggregate_value = len(answers_by_id)
        return boolean_answer, aggregate_value

    def _confidence(
        self, supporting: Sequence[GraphPath], status: ReasoningStatus
    ) -> float:
        if not supporting:
            return float(torch.sigmoid(torch.tensor(-4.0)).item())
        best = max(path.score for path in supporting)
        return float(torch.sigmoid(torch.tensor(best)).item())

    # -- accounting ---------------------------------------------------------

    def parameter_buckets(self) -> Mapping[str, int]:
        """Partition parameters into the five PLAN section 6 budget categories.

        The buckets are disjoint and sum exactly to the model total.  Ownership
        follows the section 6 wording: graph embeddings belong to the graph
        encoder and PAT's own embeddings belong to PAT, while the separate
        ``embeddings`` budget covers the query vocabulary and positional tables.
        The previous implementation counted query embeddings inside
        ``reasoning_transformer`` *and* again under ``embeddings``, so the
        reported buckets summed to 108% of the model.
        """
        buckets = dict.fromkeys(
            ("graph_encoder", "pat", "reasoning_transformer", "heads", "embeddings"), 0
        )
        query_embedding_prefixes = (
            "query_encoder.token_embedding",
            "query_encoder.position_embedding",
            "query_encoder.structure_embedding",
        )
        for name, parameter in self.named_parameters():
            count = parameter.numel()
            if name.startswith(query_embedding_prefixes):
                buckets["embeddings"] += count
            elif name.startswith("graph_encoder"):
                buckets["graph_encoder"] += count
            elif name.startswith("pat"):
                buckets["pat"] += count
            elif name.startswith(("query_encoder", "reasoning_core")):
                buckets["reasoning_transformer"] += count
            else:
                buckets["heads"] += count
        total = sum(parameter.numel() for parameter in self.parameters())
        if sum(buckets.values()) != total:
            raise RuntimeError(
                f"parameter buckets sum to {sum(buckets.values()):,} but the model has {total:,}"
            )
        return buckets


def _path_key(path: GraphPath) -> tuple[tuple[str, str, str], ...]:
    return tuple(
        (step.source_id, step.relation_id, step.target_id) for step in path.steps
    )


def _compat(pat_output: PATOutput | None, index: int) -> float | None:
    if pat_output is None:
        return None
    values = pat_output.compatibility
    if index >= values.shape[0]:
        return None
    return float(values[index].item())


def _slice_pat(output: PATOutput | None, rows: slice) -> PATOutput | None:
    """One example's paths out of a batched PAT pass."""
    if output is None:
        return None
    return PATOutput(
        *(getattr(output, field)[rows] for field in PATOutput.__dataclass_fields__)
    )


def _query_starts(query: GraphQuery) -> tuple[str, ...]:
    """The node ids a query starts from, as the search resolves them."""
    return tuple(query.start_entity_ids or query.start_entities)


def _start_mask(
    node_ids: Sequence[str], start_ids: Sequence[str], device: torch.device
) -> Tensor:
    """Boolean mask over ``node_ids`` flagging the query's start entities."""
    starts = set(start_ids)
    return torch.tensor(
        [node_id in starts for node_id in node_ids], dtype=torch.bool, device=device
    )