diff --git a/gigl/transforms/graph_transformer.py b/gigl/transforms/graph_transformer.py index 3856836c9..0e4c6809a 100644 --- a/gigl/transforms/graph_transformer.py +++ b/gigl/transforms/graph_transformer.py @@ -71,6 +71,7 @@ from torch_geometric.utils import to_torch_sparse_tensor from gigl.src.common.types.graph_data import EdgeType as GiGLEdgeType +from gigl.transforms.sequence_layout import select_sequences TokenInputData = dict[str, Tensor] @@ -1277,15 +1278,9 @@ def _build_sequence_layout_from_ppr_edges( The sequence order is: 1. Anchor node first, when ``include_anchor_first`` is True. 2. Destination nodes reachable by outgoing ``"ppr"`` edges from that anchor, - sorted by descending PPR weight. + sorted by descending PPR weight, with equal weights in edge order. """ batch_size = anchor_indices.size(0) - node_index_sequences = torch.full( - (batch_size, max_seq_len), - fill_value=-1, - dtype=torch.long, - device=device, - ) valid_mask = torch.zeros( (batch_size, max_seq_len), dtype=torch.bool, @@ -1301,13 +1296,20 @@ def _build_sequence_layout_from_ppr_edges( ppr_relation_feature_sequences = None if include_anchor_first and max_seq_len > 0: - node_index_sequences[:, 0] = anchor_indices - valid_mask[:, 0] = True start_pos = 1 else: start_pos = 0 if start_pos >= max_seq_len: + node_index_sequences = torch.full( + (batch_size, max_seq_len), + fill_value=-1, + dtype=torch.long, + device=device, + ) + if start_pos > 0: + node_index_sequences[:, 0] = anchor_indices + valid_mask[:, 0] = True return ( node_index_sequences, valid_mask, @@ -1330,13 +1332,6 @@ def _build_sequence_layout_from_ppr_edges( f"Requested '{PPR_RELATION_FEATURES_NAME}' but PPR edge_attr only " "contains the scalar weight column." ) - ppr_relation_feature_sequences = torch.zeros( - (batch_size, max_seq_len, edge_features.size(1) - 1), - dtype=torch.float, - device=device, - ) - edge_weights = edge_features[:, 0] - relation_features = edge_features[:, 1:] if return_relation_features else None anchor_batch_index_by_homo_idx = torch.full( (num_nodes,), @@ -1347,97 +1342,52 @@ def _build_sequence_layout_from_ppr_edges( anchor_batch_index_by_homo_idx[anchor_indices] = torch.arange( batch_size, device=device ) - - src_idx = edge_index[0] - dst_idx = edge_index[1] - anchor_batch_idx = anchor_batch_index_by_homo_idx[src_idx] + anchor_batch_idx = anchor_batch_index_by_homo_idx[edge_index[0]] keep = anchor_batch_idx >= 0 - if not keep.any(): - return ( - node_index_sequences, - valid_mask, - ppr_weight_sequences, - ppr_relation_feature_sequences, - ) - - all_anchor_batch_idx = anchor_batch_idx[keep] - all_dst_idx = dst_idx[keep] - all_weights = edge_weights[keep] - all_relation_features = ( - relation_features[keep] if relation_features is not None else None + anchor_batch_idx = anchor_batch_idx[keep] + dst_idx = edge_index[1][keep] + edge_features = edge_features[keep] + + # Pad each anchor's edges, in edge order, into one row of a ``[batch_size, + # num_candidates]`` table; at least one column, so the gathers below have + # something to index. + batch_order = torch.argsort(anchor_batch_idx, stable=True) + sorted_batch_idx = anchor_batch_idx[batch_order] + edge_counts = torch.bincount(sorted_batch_idx, minlength=batch_size) + num_candidates = max(1, int(edge_counts.max())) + group_starts = torch.cumsum(edge_counts, dim=0) - edge_counts + column = ( + torch.arange(sorted_batch_idx.numel(), device=device) + - group_starts[sorted_batch_idx].long() ) - - if include_anchor_first: - keep = all_dst_idx != anchor_indices[all_anchor_batch_idx] - if not keep.any(): - return ( - node_index_sequences, - valid_mask, - ppr_weight_sequences, - ppr_relation_feature_sequences, - ) - all_anchor_batch_idx = all_anchor_batch_idx[keep] - all_dst_idx = all_dst_idx[keep] - all_weights = all_weights[keep] - all_relation_features = ( - all_relation_features[keep] if all_relation_features is not None else None - ) - - # Flattened COO edges can be laid out in one pass by sorting first on weight - # and then stably on anchor batch id, which preserves descending-weight order - # within each anchor group without a Python loop. Relation features, when - # present, are reordered by the same permutation to stay aligned with edges. - weight_order = torch.argsort(all_weights, descending=True, stable=True) - all_anchor_batch_idx = all_anchor_batch_idx[weight_order] - all_dst_idx = all_dst_idx[weight_order] - all_weights = all_weights[weight_order] - if all_relation_features is not None: - all_relation_features = all_relation_features[weight_order] - - batch_order = torch.argsort(all_anchor_batch_idx, stable=True) - sorted_batch_idx = all_anchor_batch_idx[batch_order] - sorted_dst_idx = all_dst_idx[batch_order] - sorted_weights = all_weights[batch_order] - sorted_relation_features = ( - all_relation_features[batch_order] - if all_relation_features is not None - else None + candidates = torch.full( + (batch_size, num_candidates), -1, dtype=torch.long, device=device ) - - n = sorted_batch_idx.size(0) - is_group_start = torch.zeros(n, dtype=torch.long, device=device) - is_group_start[0] = 1 - if n > 1: - is_group_start[1:] = (sorted_batch_idx[1:] != sorted_batch_idx[:-1]).long() - - group_id = is_group_start.cumsum(0) - 1 - group_starts = torch.nonzero(is_group_start, as_tuple=True)[0] - positions = torch.arange(n, device=device) - group_starts[group_id] + start_pos - - valid = positions < max_seq_len - valid_batch_idx = sorted_batch_idx[valid] - valid_positions = positions[valid] - valid_dst_idx = sorted_dst_idx[valid] - valid_weights = sorted_weights[valid] - valid_relation_features = ( - sorted_relation_features[valid] - if sorted_relation_features is not None - else None + candidates[sorted_batch_idx, column] = dst_idx[batch_order] + padded_features = edge_features.new_zeros( + (batch_size, num_candidates, edge_features.size(1)) ) - - node_index_sequences[valid_batch_idx, valid_positions] = valid_dst_idx - valid_mask[valid_batch_idx, valid_positions] = True - if ppr_weight_sequences is not None: - ppr_weight_sequences[valid_batch_idx, valid_positions, 0] = ( - valid_weights.float() - ) - if ( - ppr_relation_feature_sequences is not None - and valid_relation_features is not None - ): - ppr_relation_feature_sequences[valid_batch_idx, valid_positions] = ( - valid_relation_features.float() - ) + padded_features[sorted_batch_idx, column] = edge_features[batch_order] + + node_index_sequences, source, lengths = select_sequences( + anchor_indices, + candidates, + max_seq_len, + priority=-padded_features[:, :, 0], + ties_by_column=True, + include_anchor=include_anchor_first, + ) + valid_mask = torch.arange(max_seq_len, device=device) < lengths.unsqueeze(1) + # Anchor and padding tokens have no source edge; their payload is zero. + has_edge = (source >= 0).unsqueeze(-1) + token_features = padded_features.gather( + 1, + source.clamp(min=0).unsqueeze(-1).expand(-1, -1, padded_features.size(-1)), + ).masked_fill(~has_edge, 0.0) + if return_edge_weights: + ppr_weight_sequences = token_features[:, :, :1] + if return_relation_features: + ppr_relation_feature_sequences = token_features[:, :, 1:] return ( node_index_sequences, diff --git a/gigl/transforms/sequence_layout.py b/gigl/transforms/sequence_layout.py new file mode 100644 index 000000000..ca6e90d43 --- /dev/null +++ b/gigl/transforms/sequence_layout.py @@ -0,0 +1,232 @@ +"""Shared pieces for building per-anchor token sequences. + +``select_sequences`` picks each anchor's tokens from its candidates into a padded +``[num_sequences, max_seq_len]`` grid, whatever generated the candidates (k-hop +expansion, PPR edges). ``PackedLayout`` holds the offsets +``gigl.nn.packed_attention`` needs to run on real tokens only, and +``pack_sequences`` computes it from the grid's sequence lengths. +""" + +from typing import NamedTuple, Optional + +import torch +import torch.nn.functional as F +from jaxtyping import Int, Int32, Int64 +from torch import Tensor + + +class PackedLayout(NamedTuple): + """Where each sequence lives in the packed token buffer. + + Attributes: + cu_seqlens: ``[num_sequences + 1]`` int32 cumulative sequence lengths, + on the same device as the tokens. Sequence ``i`` occupies + ``[cu_seqlens[i], cu_seqlens[i + 1])``. The attention kernel + requires int32, so the total token count must fit in int32. + max_seq_len: Python int upper bound on any sequence length. It must be at + least the longest sequence, or the attention kernel skips rows. The + attention clamps it to the token count. + """ + + cu_seqlens: Int32[Tensor, "num_sequences_plus_one"] + max_seq_len: int + + +class SelectedSequences(NamedTuple): + """The tokens ``select_sequences`` picked for each anchor. + + Attributes: + nodes: ``[num_anchors, max_seq_len]`` long node ids, the anchor in + column 0 if ``include_anchor``, -1 from the sequence length on. + source: ``[num_anchors, max_seq_len]`` long column of ``candidates`` + each token came from, so a per-candidate payload follows with + ``payload.gather(1, source.clamp(min=0))``. It is -1 at the anchor + and at padding. + lengths: ``[num_anchors]`` long count of valid tokens, anchor included. + It is at most ``max_seq_len``, except that it counts the anchor even + when ``max_seq_len`` is 0. + """ + + nodes: Int64[Tensor, "num_anchors max_seq_len"] + source: Int64[Tensor, "num_anchors max_seq_len"] + lengths: Int64[Tensor, "num_anchors"] + + +def pack_sequences( + lengths: Int64[Tensor, "num_sequences"], max_seq_len: int +) -> tuple[Int64[Tensor, "num_tokens"], PackedLayout]: + """Computes what is needed to drop the padding from a padded grid of sequences. + + Sequence ``i`` is valid in its first ``lengths[i]`` positions. Any + ``[num_sequences, max_seq_len]`` per-position tensor gathers into packed + order with ``grid.flatten()[flat_index]``, and ``cu_seqlens`` marks where + each packed sequence starts. This costs one host read, of the total token + count; it is why the packed buffer holds exactly the real tokens. + + A length above ``max_seq_len`` is clamped to it, because the grid has no + positions beyond that. Checking for it would cost another host read. + + Example: + >>> flat_index, packed = pack_sequences(torch.tensor([2, 1, 3]), max_seq_len=3) + >>> flat_index.tolist() + [0, 1, 3, 6, 7, 8] + >>> packed.cu_seqlens.tolist() + [0, 2, 3, 6] + + Args: + max_seq_len: Width of the padded grid. + + Returns: + - ``[num_tokens]`` long indices of the valid positions in the flattened + grid, sequence by sequence. + - The packed layout. + """ + lengths = lengths.clamp(max=max_seq_len) + valid = torch.arange(max_seq_len, device=lengths.device) < lengths.unsqueeze(1) + cu_seqlens = torch.cat([lengths.new_zeros(1), torch.cumsum(lengths, 0)]) + flat_index = valid.flatten().nonzero().squeeze(1) + return flat_index, PackedLayout(cu_seqlens.to(torch.int32), max_seq_len) + + +# Bits of a sort key that hold the candidate's column; a node id takes the bits +# above. A fixed width keeps shapes, not the column count, out of the arithmetic, so +# ``torch.compile`` does not recompile when the candidate count changes. +_COLUMN_BITS = 26 +_COLUMN_MASK = (1 << _COLUMN_BITS) - 1 + + +def select_sequences( + anchor_indices: Int[Tensor, "num_anchors"], + candidates: Int64[Tensor, "num_anchors num_candidates"], + max_seq_len: int, + priority: Optional[Tensor] = None, + ties_by_column: bool = False, + include_anchor: bool = True, +) -> SelectedSequences: + """Picks each anchor's tokens from its candidates, with fixed shapes. + + Every anchor's sequence is the anchor, then its distinct candidate nodes other + than the anchor, ordered by ``priority`` (ties by ascending node id, or by + column with ``ties_by_column``) and cut at ``max_seq_len``. A node that + appears in several columns keeps the lowest column, so put the preferred copy + of a node first. This is the selection step shared by the k-hop and PPR + sequence builders: they differ in how they generate the candidates and in + what ``priority`` ranks. + + Example: + Two anchors, 0 and 1. Node 3 is a duplicate candidate of anchor 0, anchor 0 + is among its own candidates, and -1 is padding:: + + >>> anchors = torch.tensor([0, 1]) + >>> candidates = torch.tensor([[5, 3, 0, 3, -1], [4, 1, 6, 4, 2]]) + >>> nodes, source, lengths = select_sequences(anchors, candidates, 4) + >>> nodes.tolist() + [[0, 3, 5, -1], [1, 2, 4, 6]] + >>> lengths.tolist() + [3, 4] + + Gather a per-candidate payload with ``source``, zeroing the anchor and + padding tokens:: + + >>> payload = torch.arange(10, 20).reshape(2, 5) + >>> (payload.gather(1, source.clamp(min=0)) * (source >= 0)).tolist() + [[0, 11, 10, 0], [0, 19, 15, 17]] + + A ``priority`` ranks the candidates (smaller first); equal priorities go by + node id, or by column with ``ties_by_column``:: + + >>> priority = torch.tensor([[0, 1, 0, 1, 0], [2, 9, 0, 5, 0]]) + >>> select_sequences(anchors, candidates, 4, priority).nodes.tolist() + [[0, 5, 3, -1], [1, 2, 6, 4]] + >>> select_sequences( + ... anchors, candidates, 4, priority, ties_by_column=True + ... ).nodes.tolist() + [[0, 5, 3, -1], [1, 6, 2, 4]] + + Without ``include_anchor`` the candidates are the whole sequence, so anchor + 0 stays in its sequence as the candidate it is:: + + >>> select_sequences(anchors, candidates, 4, include_anchor=False).source.tolist() + [[2, 1, 0, -1], [1, 4, 0, 2]] + + Args: + anchor_indices: ``[num_anchors]`` node ids of the anchors. + candidates: ``[num_anchors, num_candidates]`` long candidate node ids, + -1 at padding. The anchor itself may appear. There can be at most + ``2 ** 26`` columns. + max_seq_len: Tokens per sequence, anchor included. + priority: ``[num_anchors, num_candidates]`` integer or float rank of each + candidate; smaller comes first. If None, candidates are ordered by + ascending node id, which costs one sort instead of two. Pass the + column index to keep the candidates' own order. + ties_by_column: Whether candidates with equal ``priority`` stay in column + order instead of ascending node id. It has no effect without + ``priority``. + include_anchor: Whether column 0 holds the anchor, which is then dropped + from the candidates. If False, the sequence is only candidates, and + the anchor stays in it if it is one of them. + + Returns: + The ``SelectedSequences``: ``nodes``, ``source`` and ``lengths``. + """ + anchor_indices = anchor_indices.long() + num_candidates = candidates.size(1) + padding_key = torch.iinfo(torch.long).max + + # Sort ``(node << _COLUMN_BITS) | column``: each node's copies are adjacent + # with the lowest column first, and padding sorts last. + columns = torch.arange(num_candidates, device=candidates.device) + keys = torch.where( + candidates >= 0, (candidates << _COLUMN_BITS) | columns, padding_key + ) + keys = torch.sort(keys, dim=1).values + nodes = keys >> _COLUMN_BITS + keep = (keys != padding_key) & F.pad(nodes.diff(dim=1) != 0, (1, 0), value=True) + if include_anchor: + keep = keep & (nodes != anchor_indices.unsqueeze(1)) + if priority is not None: + if ties_by_column: + # Put the kept candidates back in column order, so the stable sort + # below leaves priority ties there. + by_column = torch.sort( + torch.where(keep, keys & _COLUMN_MASK, num_candidates), dim=1 + ).indices + keys = keys.gather(1, by_column) + keep = keep.gather(1, by_column) + # Dropped candidates rank last, so the kept ones come first, in + # ``priority`` order and then ascending node id. Padding decodes to an + # out-of-range column; its priority is never used. + columns = (keys & _COLUMN_MASK).clamp(max=num_candidates - 1) + if priority.is_floating_point(): + worst = torch.inf + else: + worst = torch.iinfo(priority.dtype).max + order = torch.sort( + torch.where(keep, priority.gather(1, columns), worst), + dim=1, + stable=True, + ).indices + keys = keys.gather(1, order) + keep = keep.gather(1, order) + # With the anchor, it takes column 0 and the kept candidates are numbered from + # 1. The kept candidates are cut at ``max_seq_len``. + first_position = int(include_anchor) + positions = torch.cumsum(keep, dim=1) - 1 + first_position + keep = keep & (positions < max_seq_len) + # Dropped candidates scatter into an extra last column, which is sliced off. + packed = torch.full( + (anchor_indices.numel(), max_seq_len + 1), + -1, + dtype=torch.long, + device=keys.device, + ) + packed.scatter_(1, torch.where(keep, positions, max_seq_len), keys) + packed = packed[:, :max_seq_len] + if include_anchor and max_seq_len > 0: + packed[:, 0] = anchor_indices << _COLUMN_BITS + source = torch.where(packed >= 0, packed & _COLUMN_MASK, -1) + if include_anchor and max_seq_len > 0: + source[:, 0] = -1 + return SelectedSequences( + packed >> _COLUMN_BITS, source, first_position + keep.sum(dim=1) + ) diff --git a/gigl/transforms/static_khop.py b/gigl/transforms/static_khop.py new file mode 100644 index 000000000..9023cf157 --- /dev/null +++ b/gigl/transforms/static_khop.py @@ -0,0 +1,323 @@ +"""Builds each anchor's k-hop token sequence with fixed-shape tensor ops. + +This is the layout half of the packed backend of ``GraphTransformerEncoder``. For +every anchor it picks which nodes become its tokens (the anchor, then its +neighbours within ``hop_distance`` hops), and ``pack_sequences`` (in ``gigl.transforms.sequence_layout``) turns the padded +result into the offsets that ``gigl.nn.packed_attention`` needs to run the +transformer on real tokens only. + +Why this exists: + - Sequences are short and uneven (about 90 real tokens out of a 256-token + budget in our models), so a padded ``[num_sequences, max_seq_len]`` batch + spends most of its compute and memory on padding. Packing the real tokens + into one buffer, with ``cu_seqlens`` marking the boundaries, avoids that. + - The sparse k-hop path in ``gigl.transforms.graph_transformer`` (sparse + matmul, ``coalesce``, ``unique``, ``nonzero``, bool-mask indexing) produces + tensors whose sizes depend on the data. It needs many host syncs and blocks + ``torch.compile``. This module instead uses a dense neighbour table and only + sort, gather and scatter, whose output shapes are known up front. It + produces the same layout as that path's default: anchor first, then + neighbours in ascending node id. + - A dense table has one cost: the candidates before deduplication grow as + ``num_anchors * width ** hop_distance``, and one hub node sets ``width`` for + the whole batch. ``khop_layout`` counts the candidates before it builds + anything and raises ``KhopLayoutTooLargeError`` past ``max_candidates``. + ``GraphTransformerEncoder`` catches that and runs the batch on the padded + path. + - The whole path makes exactly two host reads (each is a GPU sync): the table + width (the largest per-node degree, all edge types combined) and, in + ``pack_sequences``, the total token count. + +How a layout is built:: + + edge_index [2, num_edges] + | sort edges by source; an edge's slot is its rank in its source's group + v + neighbour table [num_nodes + 1, width] width = largest degree (host read 1) + | row i = nodes one hop from i, -1 pads; + | the extra last row is all -1 + | hop h: gather the rows of the hop h - 1 frontier -> [num_anchors, width ** h] + v + candidates [num_anchors, 1 + width + ... + width ** hop_distance] + | select_sequences: the first copy of a node is its first-reach hop, so + | drop repeats, padding and the anchor itself + v + anchor first, then neighbours in ascending node id (or by hop, then node id, + or in expansion order, per ``ordering``), cut at max_seq_len + | + v + layout [num_anchors, max_seq_len] of (node << hop_bits) | hop (-1 at padding), lengths + | + | pack_sequences(lengths, max_seq_len) (host read 2: total tokens) + v + flat_index + PackedLayout(cu_seqlens, max_seq_len) for gigl.nn.packed_attention + +Constraints: + - Memory is ``O(num_anchors * width ** hop_distance)``; see + ``MAX_LAYOUT_CANDIDATES`` for the limit. + - ``cu_seqlens`` is int32, so the total token count must fit in int32. + +Example: + >>> from gigl.transforms.sequence_layout import pack_sequences + >>> edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]]) + >>> layout, lengths, hop_bits = khop_layout( + ... edge_index, + ... num_nodes=5, + ... anchor_indices=torch.arange(2), + ... hop_distance=2, + ... max_seq_len=4, + ... sampling_direction="out", + ... ) + >>> hop_bits + 2 + >>> (layout >> hop_bits).tolist() + [[0, 1, 2, -1], [1, 2, 3, -1]] + >>> (layout & 3).tolist() + [[0, 1, 2, 3], [0, 1, 2, 3]] + >>> lengths.tolist() + [3, 3] + >>> flat_index, packed = pack_sequences(lengths, max_seq_len=4) + >>> (layout.flatten()[flat_index] >> hop_bits).tolist() + [0, 1, 2, 1, 2, 3] + >>> packed.cu_seqlens.tolist() + [0, 3, 6] +""" + +from typing import Literal, NamedTuple + +import torch +from jaxtyping import Int, Int64 +from torch import Tensor + +from gigl.transforms.sequence_layout import select_sequences + +# Maximum number of int64 cells ``khop_layout`` builds: the candidates, +# ``num_anchors * (1 + width + ... + width ** hop_distance)``, plus the neighbour +# table, ``(num_nodes + 1) * width``. The layout keeps about ten int64 tensors of +# the candidate size alive at once, so the limit is about 2.5 GiB. The width is +# the largest degree in the whole batch, so one hub node can push every anchor +# over the limit and send the batch to the padded path. +MAX_LAYOUT_CANDIDATES = 2**25 + +KhopOrdering = Literal["node_id", "hop", "input"] + + +class KhopLayout(NamedTuple): + """The k-hop token layout of every anchor, as built by ``khop_layout``. + + Attributes: + layout: ``[num_anchors, max_seq_len]`` long. Each valid cell packs a + token's local node id and first-reach hop as + ``(node << hop_bits) | hop``; decode with ``cell >> hop_bits`` and + ``cell & ((1 << hop_bits) - 1)``. Column 0 is the anchor (hop 0), + followed by its neighbours within ``hop_distance`` hops in the + ``ordering`` order. The hop is the fewest hops from the anchor to the + node. Cells at and after the sequence length are -1, so only decode + the valid ones. + lengths: ``[num_anchors]`` long count of valid tokens per anchor, anchor + included, so at least 1 and at most ``max_seq_len``. + hop_bits: Python int, ``hop_distance.bit_length()``. + """ + + layout: Int64[Tensor, "num_anchors max_seq_len"] + lengths: Int64[Tensor, "num_anchors"] + hop_bits: int + + +class KhopLayoutTooLargeError(ValueError): + """The k-hop layout would have more candidates than the allowed maximum.""" + + +def _build_neighbor_table( + edge_index: Int[Tensor, "2 num_edges"], + num_nodes: int, + num_anchors: int, + hop_distance: int, + max_candidates: int, + sampling_direction: Literal["in", "out"] = "out", +) -> Int64[Tensor, "num_nodes_plus_one width"]: + """Builds the dense neighbour table that lets k-hop expansion run as gathers. + + Row ``i`` lists the nodes one hop from node ``i``, padded with -1. With + ``sampling_direction="in"`` (GLT in-edge sampling, where edges point from + neighbour to target) reachability follows the reversed edges, as in + ``heterodata_to_graph_transformer_input``. + + The table is as wide as the largest row (at least 1), which costs one host + read. The layout's size is checked right after that read and before the table + is allocated, because one hub node makes the table itself large. Duplicate + edges take one slot each. + + The table has an extra all -1 row at index ``num_nodes``. Looking up a -1 + (padding) id selects that row, so it yields -1 again. + + Args: + num_anchors: Number of anchors the layout will be built for. + hop_distance: Number of hops the layout will expand. + max_candidates: Maximum of the table cells plus the layout's candidates; see + ``MAX_LAYOUT_CANDIDATES``. + sampling_direction: ``"in"`` reverses the edges, ``"out"`` keeps them. + + Returns: + ``[num_nodes + 1, width]`` long neighbour table. + + Raises: + KhopLayoutTooLargeError: If the table cells plus candidates exceed ``max_candidates``. + """ + if sampling_direction == "in": + source, target = edge_index[1], edge_index[0] + else: + source, target = edge_index[0], edge_index[1] + # Scatter indices must be int64. + source, target = source.long(), target.long() + device = edge_index.device + # Build the table with sort and scatter rather than ``to_dense_batch``, whose + # ``fill_value`` host-to-device copy is an extra sync on CUDA. + # A stable sort groups each node's edges, and an edge's slot is its rank in + # that group. Degrees are counted with ``scatter_add_`` because ``bincount`` + # would sync with the host on CUDA. + sorted_source, order = torch.sort(source, stable=True) + degree = torch.zeros(num_nodes + 1, dtype=torch.long, device=device) + degree.scatter_add_(0, source, torch.ones_like(source)) + width = max(int(degree.max()), 1) + num_candidates = num_anchors * sum(width**hop for hop in range(hop_distance + 1)) + num_cells = num_candidates + (num_nodes + 1) * width + if num_cells > max_candidates: + raise KhopLayoutTooLargeError( + f"The k-hop layout needs {num_candidates} candidates for {num_anchors} " + f"anchors at width {width} and {hop_distance} hops, plus a " + f"{num_nodes + 1} x {width} neighbour table, more than the maximum of " + f"{max_candidates}." + ) + group_starts = torch.cumsum(degree, 0) - degree + slot = torch.arange(source.numel(), device=device) - group_starts[sorted_source] + table = torch.full((num_nodes + 1, width), -1, dtype=torch.long, device=device) + table.index_put_((sorted_source, slot), target[order]) + return table + + +def _khop_sequence_layout( + neighbor_table: Int64[Tensor, "num_nodes_plus_one width"], + anchor_indices: Int[Tensor, "num_anchors"], + hop_distance: int, + max_seq_len: int, + ordering: KhopOrdering = "node_id", +) -> KhopLayout: + """Picks the tokens of every anchor's k-hop neighbourhood, with fixed shapes. + + Gathering through the neighbour table, rather than sparse matrix ops, keeps + every tensor's shape independent of the data. The gathers produce the + candidates, hop by hop, and ``select_sequences`` turns them into sequences. + With the default ``ordering``, each sequence is the anchor, then every other + node within ``hop_distance`` hops in ascending node id, truncated to + ``max_seq_len``. The sparse path gets the same order from ``coalesce``, so + this layout equals ``_get_k_hop_neighbors_sparse`` followed by + ``_build_sequence_layout_from_sparse_neighbors`` with the anchor moved first. + With ``ordering="hop"`` the neighbours are ordered by first-reach hop, then + node id, so truncation drops the farthest hops first, as in + ``_build_hop_ordered_sequence_layout_from_sparse_edges``. With + ``ordering="input"`` they are ordered by first reach in the expansion: by hop, + then by the position of the node's parent in the previous hop, then by its + slot in the parent's table row. + + A token's hop is the hop at which it is first reached; the anchor's is 0. + + Args and returns are as in ``khop_layout``, with the neighbour table from + ``_build_neighbor_table`` in place of the edges. + """ + anchor_indices = anchor_indices.long() + hop_bits = hop_distance.bit_length() + device = anchor_indices.device + + # The candidates are laid out hop by hop, so a node's lowest column is its + # first-reach hop, which is the copy ``select_sequences`` keeps. + frontier = anchor_indices.unsqueeze(1) + candidate_parts = [frontier] + hop_parts = [torch.zeros(1, dtype=torch.long, device=device)] + for hop in range(1, hop_distance + 1): + # A -1 id selects the all -1 last row of the table. + frontier = neighbor_table[frontier].flatten(1) + candidate_parts.append(frontier) + hop_parts.append(torch.full((frontier.size(1),), hop, device=device)) + candidates = torch.cat(candidate_parts, dim=1) + hop_of_column = torch.cat(hop_parts) + if ordering == "hop": + priority = hop_of_column.expand_as(candidates) + elif ordering == "input": + priority = torch.arange(candidates.size(1), device=device).expand_as(candidates) + else: + priority = None + nodes, source, lengths = select_sequences( + anchor_indices, candidates, max_seq_len, priority + ) + hops = hop_of_column[source.clamp(min=0)] + layout = torch.where(nodes >= 0, (nodes << hop_bits) | hops, -1) + return KhopLayout(layout, lengths, hop_bits) + + +def khop_layout( + edge_index: Int[Tensor, "2 num_edges"], + num_nodes: int, + anchor_indices: Int[Tensor, "num_anchors"], + hop_distance: int, + max_seq_len: int, + sampling_direction: Literal["in", "out"] = "out", + max_candidates: int = MAX_LAYOUT_CANDIDATES, + ordering: KhopOrdering = "node_id", +) -> KhopLayout: + """Builds the k-hop token layout of every anchor; the entry point of this module. + + Builds the neighbour table, checks the candidate budget, then expands and + deduplicates each anchor's neighbourhood into a padded + ``[num_anchors, max_seq_len]`` layout. Pass ``lengths`` to ``pack_sequences`` + to drop the padding. A host read of the table width is the only sync here. + + Args: + edge_index: ``[2, num_edges]`` int32 or int64 local node ids in + ``[0, num_nodes)``. + num_nodes: Number of nodes. + anchor_indices: ``[num_anchors]`` int32 or int64 local node ids of the + anchors. + hop_distance: Number of hops to expand. + max_seq_len: Tokens per sequence, anchor included. + sampling_direction: ``"in"`` reverses the edges, ``"out"`` keeps them. + max_candidates: Maximum number of candidates plus neighbour-table cells; see + ``MAX_LAYOUT_CANDIDATES``. + ordering: How neighbours are ordered, and so which are cut at + ``max_seq_len``: ``"node_id"`` by ascending node id, ``"hop"`` by + first-reach hop and then node id, ``"input"`` by first reach in the + expansion (hop, then the parent's position, then the slot in the + parent's neighbour-table row, which follows the edge order). + + Returns: + The ``KhopLayout``: ``layout``, ``lengths`` and ``hop_bits``. + + Raises: + ValueError: If ``sampling_direction`` is not ``"in"`` or ``"out"``, or + ``ordering`` is not ``"node_id"``, ``"hop"`` or ``"input"``. + KhopLayoutTooLargeError: If the table cells plus candidates exceed ``max_candidates``. + """ + if sampling_direction not in {"in", "out"}: + raise ValueError( + "sampling_direction must be one of {'in', 'out'}, " + f"got '{sampling_direction}'." + ) + if ordering not in {"node_id", "hop", "input"}: + raise ValueError( + f"ordering must be one of {{'node_id', 'hop', 'input'}}, got '{ordering}'." + ) + neighbor_table = _build_neighbor_table( + edge_index, + num_nodes, + anchor_indices.numel(), + hop_distance, + max_candidates, + sampling_direction, + ) + return _khop_sequence_layout( + neighbor_table, + anchor_indices, + hop_distance, + max_seq_len, + ordering, + ) diff --git a/tests/test_assets/runtime_type_checking.py b/tests/test_assets/runtime_type_checking.py index f5faa2f0d..004eca76b 100644 --- a/tests/test_assets/runtime_type_checking.py +++ b/tests/test_assets/runtime_type_checking.py @@ -13,7 +13,9 @@ from types import FunctionType, UnionType from typing import Any, Final, Optional, Union, cast, get_args, get_origin -from jaxtyping import AbstractArray, install_import_hook +import jaxtyping +import torch +from jaxtyping import AbstractArray, Float, install_import_hook from typeguard import typechecked from gigl.common.logger import Logger @@ -90,6 +92,56 @@ def retain_shape_contract(annotation: object) -> object: return wrapped_function +# The config class is private to Jaxtyping, so it is read from the config object. +_JaxtypingConfig: Any = cast(Any, type(jaxtyping.config)) + + +class _SkipWhileCompilingConfig(_JaxtypingConfig): + """Jaxtyping config whose ``jaxtyping_disable`` reads true while Dynamo traces. + + Jaxtyping's wrapper returns the undecorated function when ``jaxtyping_disable`` + is true, before it binds arguments or pushes a shape memo. Dynamo cannot trace + that bookkeeping or Typeguard, and guards on it cause recompiles. Eager calls + read the configured value, so they are still checked. + """ + + @property + def jaxtyping_disable(self) -> bool: + if torch.compiler.is_compiling(): + return True + return self._eager_jaxtyping_disable + + @jaxtyping_disable.setter + def jaxtyping_disable(self, value: bool) -> None: + self._eager_jaxtyping_disable = value + + +def _skip_shape_checks_while_compiling() -> None: + """Keeps Shape Contract wrappers out of ``torch.compile``. + + Inside a compiled frame, ``jaxtyping.config`` reads ``jaxtyping_disable`` as + true, so the inlined wrapper calls the undecorated function. + + Every wrapped function shares the code object of Jaxtyping's wrapper. Dynamo + caches compiled frames per code object, so a wrapper that starts a frame + after a graph break would count each contracted function as a recompile of + the same code. ``skip_code`` runs that frame eagerly, with its checks, and + compiles the function it calls as its own frame. + """ + config = jaxtyping.config + if isinstance(config, _SkipWhileCompilingConfig): + return + eager_jaxtyping_disable = config.__dict__.pop("jaxtyping_disable") + config.__class__ = _SkipWhileCompilingConfig + config.jaxtyping_disable = eager_jaxtyping_disable + + def probe(x: Float[torch.Tensor, "_"]) -> Float[torch.Tensor, "_"]: + return x + + wrapper = jaxtyping.jaxtyped(typechecker=shape_contract_typechecker)(probe) + torch._dynamo.eval_frame.skip_code(wrapper.__code__) + + def install_runtime_typechecking() -> None: """Enable test-only runtime checks for tensor Shape Contracts. @@ -105,6 +157,8 @@ def install_runtime_typechecking() -> None: cannot be used with ``isinstance``. The custom typechecker filters those out so Typeguard enforces Shape Contracts only. + Checks run in eager calls and are skipped while ``torch.compile`` traces. + Repeated calls are safe. Runtime checking remains scoped to the current process and modules imported after this function runs. A contract violation raises ``jaxtyping.TypeCheckError`` at the call site, so an uncaught @@ -112,6 +166,7 @@ def install_runtime_typechecking() -> None: """ global _import_hook if _import_hook is None: + _skip_shape_checks_while_compiling() _import_hook = install_import_hook( modules=_SHAPE_CONTRACT_PACKAGES, typechecker="tests.test_assets.runtime_type_checking.shape_contract_typechecker", diff --git a/tests/unit/test_assets/runtime_type_checking_test.py b/tests/unit/test_assets/runtime_type_checking_test.py new file mode 100644 index 000000000..346667831 --- /dev/null +++ b/tests/unit/test_assets/runtime_type_checking_test.py @@ -0,0 +1,54 @@ +import jaxtyping +import torch +from absl.testing import absltest +from jaxtyping import Float + +from tests.test_assets.runtime_type_checking import ( + install_runtime_typechecking, + shape_contract_typechecker, +) +from tests.test_assets.test_case import TestCase + + +def _contracted(function): + return jaxtyping.jaxtyped(typechecker=shape_contract_typechecker)(function) + + +@_contracted +def _double(x: Float[torch.Tensor, "rows 2"]) -> Float[torch.Tensor, "rows 2"]: + return x * 2 + + +@_contracted +def _increment(x: Float[torch.Tensor, "rows 2"]) -> Float[torch.Tensor, "rows 2"]: + return x + 1 + + +class RuntimeTypeCheckingTest(TestCase): + def setUp(self) -> None: + super().setUp() + install_runtime_typechecking() + torch._dynamo.reset() + + def test_eager_call_with_wrong_shape_raises(self) -> None: + with self.assertRaises(jaxtyping.TypeCheckError): + _double(torch.ones(3, 4)) + + def test_check_is_skipped_inside_a_compiled_frame(self) -> None: + def run(x: torch.Tensor) -> torch.Tensor: + return _double(x) + 0 + + compiled = torch.compile(run, backend="eager", fullgraph=True) + torch.testing.assert_close(compiled(torch.ones(3, 4)), torch.full((3, 4), 2.0)) + + def test_contracted_functions_as_frame_roots_do_not_recompile(self) -> None: + # Both functions share the code object of Jaxtyping's wrapper. + with torch._dynamo.config.patch( + recompile_limit=1, fail_on_recompile_limit_hit=True + ): + torch.compile(_double, backend="eager")(torch.ones(3, 2)) + torch.compile(_increment, backend="eager")(torch.ones(3, 2)) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/transforms/graph_transformer_test.py b/tests/unit/transforms/graph_transformer_test.py index 14236fa60..0e74a3d6a 100644 --- a/tests/unit/transforms/graph_transformer_test.py +++ b/tests/unit/transforms/graph_transformer_test.py @@ -11,6 +11,7 @@ from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation from gigl.transforms.graph_transformer import ( + _build_sequence_layout_from_ppr_edges, _get_k_hop_neighbors_sparse, heterodata_to_graph_transformer_input, tensors_to_graph_transformer_input, @@ -966,6 +967,74 @@ def test_ppr_sequence_can_return_token_input_and_attention_bias_features(self): ) +class TestBuildSequenceLayoutFromPprEdges(TestCase): + """Anchor 0's edges, in edge order, are 0->5 (0.5), 0->3 (0.9), 0->0 (0.7, a + self edge), 0->6 (0.5), 0->4 (0.9) and 0->7 (0.1), so by descending weight + with ties in edge order it ranks 3, 4, 0, 5, 6, 7. Anchor 1 has the single + edge 1->2. The edge 2->0 starts at a non-anchor and is ignored. The second + edge feature column is the edge position plus one.""" + + def _build(self, max_seq_len: int, include_anchor_first: bool): + edge_index = torch.tensor([[0, 0, 0, 2, 0, 1, 0, 0], [5, 3, 0, 0, 6, 2, 4, 7]]) + weights = torch.tensor([0.5, 0.9, 0.7, 0.99, 0.5, 0.3, 0.9, 0.1]) + edge_attr = torch.stack([weights, torch.arange(8) + 1.0], dim=1) + return _build_sequence_layout_from_ppr_edges( + edge_index=edge_index, + edge_attr=edge_attr, + anchor_indices=torch.tensor([0, 1]), + max_seq_len=max_seq_len, + include_anchor_first=include_anchor_first, + num_nodes=8, + device=torch.device("cpu"), + return_edge_weights=True, + return_relation_features=True, + ) + + def test_anchor_first_drops_self_edge_and_orders_by_weight(self): + nodes, valid_mask, weights, relation = self._build(6, True) + self.assert_tensor_equality( + nodes, torch.tensor([[0, 3, 4, 5, 6, 7], [1, 2, -1, -1, -1, -1]]) + ) + self.assert_tensor_equality( + valid_mask, + torch.tensor([[True] * 6, [True, True] + [False] * 4]), + ) + self.assert_tensor_equality( + weights[:, :, 0], + torch.tensor( + [[0.0, 0.9, 0.9, 0.5, 0.5, 0.1], [0.0, 0.3, 0.0, 0.0, 0.0, 0.0]] + ), + ) + self.assert_tensor_equality( + relation[:, :, 0], + torch.tensor([[0.0, 2.0, 7.0, 1.0, 5.0, 8.0], [0.0, 6.0, 0, 0, 0, 0]]), + ) + + def test_anchor_first_truncates_at_max_seq_len(self): + nodes, valid_mask, _, _ = self._build(4, True) + self.assert_tensor_equality(nodes, torch.tensor([[0, 3, 4, 5], [1, 2, -1, -1]])) + self.assert_tensor_equality( + valid_mask, torch.tensor([[True] * 4, [True, True, False, False]]) + ) + + def test_without_anchor_first_keeps_self_edge(self): + nodes, valid_mask, weights, relation = self._build(5, False) + self.assert_tensor_equality( + nodes, torch.tensor([[3, 4, 0, 5, 6], [2, -1, -1, -1, -1]]) + ) + self.assert_tensor_equality( + valid_mask, torch.tensor([[True] * 5, [True] + [False] * 4]) + ) + self.assert_tensor_equality( + weights[:, :, 0], + torch.tensor([[0.9, 0.9, 0.7, 0.5, 0.5], [0.3, 0.0, 0.0, 0.0, 0.0]]), + ) + self.assert_tensor_equality( + relation[:, :, 0], + torch.tensor([[2.0, 7.0, 3.0, 1.0, 5.0], [6.0, 0, 0, 0, 0]]), + ) + + class TestPyTorchTransformerIntegration(TestCase): """Tests for integration with PyTorch TransformerEncoderLayer.""" diff --git a/tests/unit/transforms/static_khop_test.py b/tests/unit/transforms/static_khop_test.py new file mode 100644 index 000000000..3043326c9 --- /dev/null +++ b/tests/unit/transforms/static_khop_test.py @@ -0,0 +1,350 @@ +"""The fixed-shape k-hop layout equals the sparse k-hop layout.""" + +import doctest +import itertools +from typing import Any, Literal +from unittest import mock + +import torch +from absl.testing import absltest + +from gigl.transforms import sequence_layout, static_khop +from gigl.transforms.graph_transformer import ( + _build_hop_ordered_sequence_layout_from_sparse_edges, + _build_sequence_layout_from_sparse_neighbors, + _get_k_hop_neighbors_sparse, +) +from gigl.transforms.sequence_layout import pack_sequences, select_sequences +from gigl.transforms.static_khop import KhopLayoutTooLargeError +from gigl.transforms.static_khop import khop_layout as _khop_layout +from tests.test_assets.test_case import TestCase + +_CPU = torch.device("cpu") + + +def khop_layout( + *args: Any, **kwargs: Any +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``khop_layout`` with the cells decoded: node ids, first-reach hops (0 at padding), lengths.""" + layout, lengths, hop_bits = _khop_layout(*args, **kwargs) + valid = layout >= 0 + nodes = torch.where(valid, layout >> hop_bits, -1) + hops = torch.where(valid, layout & ((1 << hop_bits) - 1), 0) + return nodes, hops, lengths + + +def _random_graph(num_nodes: int, num_edges: int, seed: int) -> torch.Tensor: + """Random edges plus duplicates, self loops (on a seed and a non-seed) and a 2-cycle.""" + generator = torch.Generator().manual_seed(seed) + edge_index = torch.randint(0, num_nodes, (2, num_edges), generator=generator) + extra = torch.tensor([[0, 7, 1, 2], [0, 7, 2, 1]]) + return torch.cat([edge_index, edge_index[:, :5], extra], dim=1) + + +def _first_reach_hops( + node_index: torch.Tensor, + valid: torch.Tensor, + adjacency: torch.Tensor, + num_nodes: int, + hop_distance: int, +) -> torch.Tensor: + """The smallest hop at which each valid token is reachable; 0 at the anchor and padding.""" + anchors = torch.arange(node_index.size(0)) + hops = torch.zeros_like(node_index) + # Go from the farthest hop to the nearest so the nearest reach wins. + for hop in range(hop_distance, 0, -1): + reach = _get_k_hop_neighbors_sparse(anchors, adjacency, num_nodes, hop, _CPU) + in_reach = reach.to_dense()[anchors.unsqueeze(1), node_index.clamp(min=0)] > 0 + hops = torch.where(in_reach & (node_index >= 0), hop, hops) + hops = torch.where(valid, hops, 0) + hops[:, 0] = 0 + return hops + + +def _sparse_layout( + edge_index: torch.Tensor, + num_nodes: int, + num_anchors: int, + max_seq_len: int, + direction: Literal["in", "out"], + hop_distance: int = 2, + ordering: Literal["node_id", "hop"] = "node_id", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Node ids, validity and first-reach hops from the sparse path.""" + adjacency = edge_index.flip(0) if direction == "in" else edge_index + anchors = torch.arange(num_anchors) + if ordering == "hop": + node_index, valid = _build_hop_ordered_sequence_layout_from_sparse_edges( + anchor_indices=anchors, + edge_index=adjacency, + num_nodes=num_nodes, + hop_distance=hop_distance, + max_seq_len=max_seq_len, + include_anchor_first=True, + device=_CPU, + ) + else: + node_index, valid = _build_sequence_layout_from_sparse_neighbors( + reachable=_get_k_hop_neighbors_sparse( + anchors, adjacency, num_nodes, hop_distance, _CPU + ), + anchor_indices=anchors, + max_seq_len=max_seq_len, + include_anchor_first=True, + device=_CPU, + ) + hops = _first_reach_hops(node_index, valid, adjacency, num_nodes, hop_distance) + return node_index, valid, hops + + +class StaticKhopLayoutTest(TestCase): + def _assert_layouts_equal( + self, + edge_index: torch.Tensor, + num_nodes: int, + num_anchors: int, + max_seq_len: int, + direction: Literal["in", "out"], + hop_distance: int = 2, + ordering: Literal["node_id", "hop"] = "node_id", + ) -> None: + expected_nodes, expected_valid, expected_hops = _sparse_layout( + edge_index, + num_nodes, + num_anchors, + max_seq_len, + direction, + hop_distance, + ordering, + ) + nodes, hops, lengths = khop_layout( + edge_index, + num_nodes, + torch.arange(num_anchors), + hop_distance=hop_distance, + max_seq_len=max_seq_len, + sampling_direction=direction, + ordering=ordering, + ) + valid = torch.arange(max_seq_len).unsqueeze(0) < lengths.unsqueeze(1) + self.assertTrue(torch.equal(valid, expected_valid)) + self.assertTrue(torch.equal(nodes, expected_nodes)) + self.assertTrue(torch.equal(hops, expected_hops)) + + def test_matches_sparse_layout(self) -> None: + for ( + hop_distance, + ordering, + direction, + max_seq_len, + seed, + ) in itertools.product( + (1, 2, 3, 4), ("node_id", "hop"), ("in", "out"), (3, 4, 16, 64), range(2) + ): + with self.subTest( + hop=hop_distance, + ordering=ordering, + direction=direction, + L=max_seq_len, + seed=seed, + ): + self._assert_layouts_equal( + _random_graph(40, 80, seed), + num_nodes=40, + num_anchors=10, + max_seq_len=max_seq_len, + direction=direction, + hop_distance=hop_distance, + ordering=ordering, + ) + + def test_hop_order_truncation_drops_the_farthest_hops_first(self) -> None: + # Anchor 0 has hop-1 neighbours 5 and 6 and hop-2 neighbours 1 and 2. + edge_index = torch.tensor([[0, 0, 5, 6], [5, 6, 1, 2]]) + for ordering, expected in (("node_id", [0, 1, 2]), ("hop", [0, 5, 6])): + nodes, hops, _ = khop_layout( + edge_index, + 7, + torch.tensor([0]), + hop_distance=2, + max_seq_len=3, + sampling_direction="out", + ordering=ordering, + ) + self.assertEqual(nodes[0].tolist(), expected) + + def test_hop_is_first_reach_and_isolated_anchors_have_length_one(self) -> None: + # Node 3 is a hop-1 and a hop-2 neighbour of anchor 0; anchor 1 is isolated + # and node 4 is shared between anchors 0 and 2. + edge_index = torch.tensor([[3, 5, 3, 4, 4], [0, 0, 5, 3, 2]]) + nodes, hops, lengths = khop_layout( + edge_index, + 6, + torch.arange(3), + hop_distance=2, + max_seq_len=5, + sampling_direction="in", + ) + self.assertEqual( + nodes.tolist(), [[0, 3, 4, 5, -1], [1, -1, -1, -1, -1], [2, 4, -1, -1, -1]] + ) + self.assertEqual( + hops.tolist(), [[0, 1, 2, 1, 0], [0, 0, 0, 0, 0], [0, 1, 0, 0, 0]] + ) + self.assertEqual(lengths.tolist(), [4, 1, 2]) + self._assert_layouts_equal(edge_index, 6, 3, 5, direction="in") + + def test_accepts_int32_edge_index_and_anchors(self) -> None: + edge_index = _random_graph(40, 90, 0) + anchors = torch.tensor([5, 0, 17]) + for hop_distance in (0, 1, 2): + with self.subTest(hop_distance=hop_distance): + expected = khop_layout(edge_index, 40, anchors, hop_distance, 8) + actual = khop_layout( + edge_index.int(), 40, anchors.int(), hop_distance, 8 + ) + for expected_tensor, actual_tensor in zip(expected, actual): + self.assertEqual(actual_tensor.dtype, torch.long) + self.assertTrue(torch.equal(actual_tensor, expected_tensor)) + + def test_rejects_an_unknown_sampling_direction(self) -> None: + edge_index = torch.tensor([[0, 1], [1, 2]]) + with self.assertRaises(ValueError): + khop_layout(edge_index, 3, torch.arange(1), 2, 4, sampling_direction="both") + + def test_rejects_a_layout_over_the_candidate_budget(self) -> None: + # Width 3 at two hops is 1 + 3 + 9 = 13 candidates per anchor, plus a + # 5 x 3 neighbour table. + edge_index = torch.tensor([[0, 0, 0], [1, 2, 3]]) + anchors = torch.arange(2) + khop_layout(edge_index, 4, anchors, 2, 8, max_candidates=41) + with self.assertRaises(KhopLayoutTooLargeError): + khop_layout(edge_index, 4, anchors, 2, 8, max_candidates=40) + + def test_rejects_a_huge_hub_table_before_allocating_it(self) -> None: + # One anchor passes the candidate check on its own (1 + 3000 + 3000 ** 2 + # is about 9M), but a 200k-node table of that width would be 4.8 GB. + num_nodes = 200_000 + edge_index = torch.stack( + [torch.zeros(3000, dtype=torch.long), torch.arange(1, 3001)] + ) + with mock.patch.object(torch, "full", side_effect=AssertionError("allocated")): + with self.assertRaises(KhopLayoutTooLargeError): + khop_layout(edge_index, num_nodes, torch.tensor([5]), 2, 8) + + def test_orderings_pick_different_tokens_when_truncated(self) -> None: + # Anchor 0 reaches 5 then 3 in slot order. 5 leads to 8 and 1; 3 leads to + # 7 and 1. Node 1 is first reached through 5, at hop 2. + edge_index = torch.tensor([[0, 0, 5, 5, 3, 3], [5, 3, 8, 1, 7, 1]]) + expected = { + "node_id": ([0, 1, 3, 5, 7, 8], [0, 2, 1, 1, 2, 2]), + "hop": ([0, 3, 5, 1, 7, 8], [0, 1, 1, 2, 2, 2]), + "input": ([0, 5, 3, 8, 1, 7], [0, 1, 1, 2, 2, 2]), + } + for ordering, (expected_nodes, expected_hops) in expected.items(): + for max_seq_len in (6, 4): + with self.subTest(ordering=ordering, L=max_seq_len): + nodes, hops, lengths = khop_layout( + edge_index, + 10, + torch.tensor([0]), + hop_distance=2, + max_seq_len=max_seq_len, + ordering=ordering, + ) + self.assertEqual(nodes[0].tolist(), expected_nodes[:max_seq_len]) + self.assertEqual(hops[0].tolist(), expected_hops[:max_seq_len]) + self.assertEqual(lengths.tolist(), [max_seq_len]) + + def test_rejects_an_unknown_ordering(self) -> None: + edge_index = torch.tensor([[0, 1], [1, 2]]) + with self.assertRaises(ValueError): + khop_layout(edge_index, 3, torch.arange(1), 2, 4, ordering="degree") + + def test_select_sequences_orders_by_priority_and_keeps_the_first_copy(self) -> None: + # Anchor 4 is among its candidates and is dropped from them. Node 9 is in + # columns 0 and 2; the first copy (priority 0.5) is the one kept. Nodes 6 + # and 9 tie on priority and fall back to ascending node id. + anchors = torch.tensor([4, 1]) + candidates = torch.tensor([[9, 2, 9, 4, -1, 6], [-1] * 6]) + priority = torch.tensor([[0.5, 0.9, 0.1, 0.0, 0.0, 0.5], [0.0] * 6]) + nodes, source, lengths = select_sequences(anchors, candidates, 5, priority) + self.assertEqual(nodes.tolist(), [[4, 6, 9, 2, -1], [1, -1, -1, -1, -1]]) + self.assertEqual(source.tolist(), [[-1, 5, 0, 1, -1], [-1] * 5]) + self.assertEqual(lengths.tolist(), [4, 1]) + + nodes, source, lengths = select_sequences(anchors, candidates, 3, priority) + self.assertEqual(nodes.tolist(), [[4, 6, 9], [1, -1, -1]]) + self.assertEqual(lengths.tolist(), [3, 1]) + + nodes, source, lengths = select_sequences(anchors, candidates, 5) + self.assertEqual(nodes[0].tolist(), [4, 2, 6, 9, -1]) + self.assertEqual(source[0].tolist(), [-1, 1, 5, 0, -1]) + + def test_select_sequences_can_tie_by_column_and_leave_out_the_anchor(self) -> None: + anchors = torch.tensor([4, 1]) + candidates = torch.tensor([[9, 2, 9, 4, -1, 6], [-1] * 6]) + priority = torch.tensor([[0.5, 0.9, 0.1, 0.0, 0.0, 0.5], [0.0] * 6]) + # Nodes 9 and 6 tie on priority; 9 comes first because its column does. + nodes, source, lengths = select_sequences( + anchors, candidates, 5, priority, ties_by_column=True + ) + self.assertEqual(nodes.tolist(), [[4, 9, 6, 2, -1], [1, -1, -1, -1, -1]]) + self.assertEqual(source[0].tolist(), [-1, 0, 5, 1, -1]) + self.assertEqual(lengths.tolist(), [4, 1]) + + # Without the anchor slot, anchor 4 is an ordinary candidate (column 3). + nodes, source, lengths = select_sequences( + anchors, candidates, 5, priority, include_anchor=False + ) + self.assertEqual(nodes.tolist(), [[4, 6, 9, 2, -1], [-1] * 5]) + self.assertEqual(source.tolist(), [[3, 5, 0, 1, -1], [-1] * 5]) + self.assertEqual(lengths.tolist(), [4, 0]) + + nodes, _, lengths = select_sequences( + anchors, candidates, 3, priority, include_anchor=False + ) + self.assertEqual(nodes.tolist(), [[4, 6, 9], [-1] * 3]) + self.assertEqual(lengths.tolist(), [3, 0]) + + def test_select_sequences_does_not_recompile_across_batch_shapes(self) -> None: + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + self.addCleanup(torch._dynamo.reset) + compiled = torch.compile(select_sequences, backend="eager", dynamic=True) + generator = torch.Generator().manual_seed(0) + for num_anchors, num_candidates in ((6, 9), (11, 30), (4, 17)): + candidates = torch.randint( + -1, 50, (num_anchors, num_candidates), generator=generator + ) + priority = torch.rand(num_anchors, num_candidates, generator=generator) + anchors = torch.arange(num_anchors) + for args in ((anchors, candidates, 8, priority), (anchors, candidates, 8)): + expected = select_sequences(*args) + for expected_tensor, actual_tensor in zip(expected, compiled(*args)): + self.assertTrue(torch.equal(expected_tensor, actual_tensor)) + # One graph for each of the two call signatures. + self.assertEqual(torch._dynamo.utils.counters["stats"]["unique_graphs"], 2) + + def test_docstring_examples(self) -> None: + for module in (sequence_layout, static_khop): + with self.subTest(module=module.__name__): + result = doctest.testmod(module) + self.assertGreater(result.attempted, 0) + self.assertEqual(result.failed, 0) + + def test_pack_sequences_layout(self) -> None: + flat_index, packed = pack_sequences(torch.tensor([2, 1, 3]), 3) + self.assertEqual(flat_index.tolist(), [0, 1, 3, 6, 7, 8]) + self.assertEqual(packed.cu_seqlens.tolist(), [0, 2, 3, 6]) + self.assertEqual(packed.cu_seqlens.dtype, torch.int32) + self.assertEqual(packed.max_seq_len, 3) + + # A length above the grid width is clamped, so the offsets match the indices. + flat_index, packed = pack_sequences(torch.tensor([5, 1]), 3) + self.assertEqual(flat_index.tolist(), [0, 1, 2, 3]) + self.assertEqual(packed.cu_seqlens.tolist(), [0, 3, 4]) + + +if __name__ == "__main__": + absltest.main()