From e11308b517ee84d4dc38c66f59308d5db1328a05 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 23 Oct 2026 16:44:46 +0000 Subject: [PATCH 1/9] Add a fixed-shape k-hop sequence layout for packed graph transformer tokens Adds gigl/transforms/static_khop.py (khop_layout, pack_sequences, PackedSequences, KhopLayoutTooLargeError) and its tests. Nothing uses it yet; the attention module and the encoder integration follow in later PRs. Co-Authored-By: Claude Opus 5.5 --- gigl/transforms/static_khop.py | 347 ++++++++++++++++++++++ tests/unit/transforms/static_khop_test.py | 241 +++++++++++++++ 2 files changed, 588 insertions(+) create mode 100644 gigl/transforms/static_khop.py create mode 100644 tests/unit/transforms/static_khop_test.py diff --git a/gigl/transforms/static_khop.py b/gigl/transforms/static_khop.py new file mode 100644 index 000000000..a23003a25 --- /dev/null +++ b/gigl/transforms/static_khop.py @@ -0,0 +1,347 @@ +"""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`` 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] + | sort each row by (node id, hop); 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, + with ``prioritize_hop_order``), 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 + PackedSequences(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: + >>> 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 +import torch.nn.functional as F +from jaxtyping import Int, Int32, Int64 +from torch import Tensor + +# 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 + + +class KhopLayoutTooLargeError(ValueError): + """The k-hop layout would have more candidates than the allowed maximum.""" + + +class PackedSequences(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 + + +def pack_sequences( + lengths: Int64[Tensor, "num_sequences"], max_seq_len: int +) -> tuple[Int64[Tensor, "num_tokens"], PackedSequences]: + """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. + + 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, PackedSequences(cu_seqlens.to(torch.int32), max_seq_len) + + +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. + order = torch.argsort(source, stable=True) + sorted_source = source[order] + 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, + prioritize_hop_order: bool = False, +) -> tuple[Int64[Tensor, "num_anchors max_seq_len"], Int64[Tensor, "num_anchors"], int]: + """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. 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 ``prioritize_hop_order`` 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``. + + 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() + padding_key = torch.iinfo(torch.long).max + + # A candidate's sort key is ``(node << hop_bits) | hop``. It groups each node's + # copies with the smallest hop first, so the first copy carries the + # first-reach hop. Padding gets the largest key and sorts last. + frontier = anchor_indices.unsqueeze(1) + key_parts = [frontier << hop_bits] + 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) + key_parts.append( + torch.where(frontier >= 0, (frontier << hop_bits) | hop, padding_key) + ) + keys, _ = torch.sort(torch.cat(key_parts, dim=1), dim=1) + nodes = keys >> hop_bits + hops = keys & ((1 << hop_bits) - 1) + keep = (keys != padding_key) & F.pad(nodes.diff(dim=1) != 0, (1, 0), value=True) + if prioritize_hop_order: + # Re-sort the kept candidates by (hop, node id) so truncation below drops + # the farthest hops first. + node_radix = neighbor_table.size(0) + order_keys, _ = torch.sort( + torch.where(keep, hops * node_radix + nodes, padding_key), dim=1 + ) + keep = order_keys != padding_key + nodes = order_keys % node_radix + keys = (nodes << hop_bits) | (order_keys // node_radix) + # The anchor goes to column 0 below; the other kept candidates are numbered + # from 1 and cut at ``max_seq_len``. + keep = keep & (nodes != anchor_indices.unsqueeze(1)) + positions = torch.cumsum(keep, dim=1) + keep = keep & (positions < max_seq_len) + # Dropped candidates scatter into column 0, which the anchor overwrites. + layout = torch.full( + (anchor_indices.numel(), max_seq_len), -1, dtype=torch.long, device=keys.device + ) + layout.scatter_(1, torch.where(keep, positions, 0), keys) + layout[:, 0] = anchor_indices << hop_bits + return layout, 1 + keep.sum(dim=1), 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, + prioritize_hop_order: bool = False, +) -> tuple[Int64[Tensor, "num_anchors max_seq_len"], Int64[Tensor, "num_anchors"], int]: + """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``. + prioritize_hop_order: Order neighbours by hop before node id. + + Returns: + - ``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 ascending + node id (or by hop, then node id, with ``prioritize_hop_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()``. + + Raises: + KhopLayoutTooLargeError: If the table cells plus candidates exceed ``max_candidates``. + """ + 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, + prioritize_hop_order, + ) diff --git a/tests/unit/transforms/static_khop_test.py b/tests/unit/transforms/static_khop_test.py new file mode 100644 index 000000000..c9721fcd4 --- /dev/null +++ b/tests/unit/transforms/static_khop_test.py @@ -0,0 +1,241 @@ +"""The fixed-shape k-hop layout equals the sparse k-hop layout.""" + +import itertools +from typing import Any, Literal +from unittest import mock + +import torch +from absl.testing import absltest + +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.static_khop import KhopLayoutTooLargeError, pack_sequences +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, + prioritize_hop_order: bool = False, +) -> 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 prioritize_hop_order: + 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, + prioritize_hop_order: bool = False, + ) -> None: + expected_nodes, expected_valid, expected_hops = _sparse_layout( + edge_index, + num_nodes, + num_anchors, + max_seq_len, + direction, + hop_distance, + prioritize_hop_order, + ) + 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, + prioritize_hop_order=prioritize_hop_order, + ) + 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, + prioritize_hop_order, + direction, + max_seq_len, + seed, + ) in itertools.product( + (1, 2, 3, 4), (False, True), ("in", "out"), (3, 4, 16, 64), range(2) + ): + with self.subTest( + hop=hop_distance, + hop_order=prioritize_hop_order, + 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, + prioritize_hop_order=prioritize_hop_order, + ) + + 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 prioritize_hop_order, expected in ((False, [0, 1, 2]), (True, [0, 5, 6])): + nodes, hops, _ = khop_layout( + edge_index, + 7, + torch.tensor([0]), + hop_distance=2, + max_seq_len=3, + sampling_direction="out", + prioritize_hop_order=prioritize_hop_order, + ) + 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_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_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() From 8b06c18d8d1dbeb976fbe77d35ce695757bfd086 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 23 Oct 2026 17:33:16 +0000 Subject: [PATCH 2/9] Static k-hop: get the sorted sources and their order from one sort Co-Authored-By: Claude Sonnet 5.5 --- gigl/transforms/static_khop.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/gigl/transforms/static_khop.py b/gigl/transforms/static_khop.py index a23003a25..2364f87e7 100644 --- a/gigl/transforms/static_khop.py +++ b/gigl/transforms/static_khop.py @@ -196,8 +196,7 @@ def _build_neighbor_table( # 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. - order = torch.argsort(source, stable=True) - sorted_source = source[order] + 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) @@ -256,7 +255,7 @@ def _khop_sequence_layout( key_parts.append( torch.where(frontier >= 0, (frontier << hop_bits) | hop, padding_key) ) - keys, _ = torch.sort(torch.cat(key_parts, dim=1), dim=1) + keys = torch.sort(torch.cat(key_parts, dim=1), dim=1).values nodes = keys >> hop_bits hops = keys & ((1 << hop_bits) - 1) keep = (keys != padding_key) & F.pad(nodes.diff(dim=1) != 0, (1, 0), value=True) @@ -264,9 +263,9 @@ def _khop_sequence_layout( # Re-sort the kept candidates by (hop, node id) so truncation below drops # the farthest hops first. node_radix = neighbor_table.size(0) - order_keys, _ = torch.sort( + order_keys = torch.sort( torch.where(keep, hops * node_radix + nodes, padding_key), dim=1 - ) + ).values keep = order_keys != padding_key nodes = order_keys % node_radix keys = (nodes << hop_bits) | (order_keys // node_radix) From 2c8c2d2395a89dd755c2bc7cbd964bcae2a805b0 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 29 Oct 2026 16:56:57 +0000 Subject: [PATCH 3/9] Static k-hop: validate sampling_direction; move the packed layout type to packed_layout.py - khop_layout raises ValueError for a sampling_direction other than 'in' or 'out', instead of treating every other value as 'out'. - PackedSequences is renamed PackedLayout and, with pack_sequences, moves to gigl/transforms/packed_layout.py, since neither depends on how the sequences were built. static_khop.py keeps only the k-hop code. Co-Authored-By: Claude Opus 5.5 --- gigl/transforms/packed_layout.py | 59 ++++++++++++++++++++++ gigl/transforms/static_khop.py | 61 ++++------------------- tests/unit/transforms/static_khop_test.py | 8 ++- 3 files changed, 77 insertions(+), 51 deletions(-) create mode 100644 gigl/transforms/packed_layout.py diff --git a/gigl/transforms/packed_layout.py b/gigl/transforms/packed_layout.py new file mode 100644 index 000000000..83f1d623f --- /dev/null +++ b/gigl/transforms/packed_layout.py @@ -0,0 +1,59 @@ +"""Packed token layout: where each sequence lives in a buffer without padding. + +``pack_sequences`` turns a padded ``[num_sequences, max_seq_len]`` grid of +sequences into the offsets (``PackedLayout``) that ``gigl.nn.packed_attention`` +needs to run on real tokens only. It does not depend on how the sequences were +built. +""" + +from typing import NamedTuple + +import torch +from jaxtyping import 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 + + +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. + + 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) diff --git a/gigl/transforms/static_khop.py b/gigl/transforms/static_khop.py index 2364f87e7..c9eb91839 100644 --- a/gigl/transforms/static_khop.py +++ b/gigl/transforms/static_khop.py @@ -2,7 +2,7 @@ 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`` turns the padded +neighbours within ``hop_distance`` hops), and ``pack_sequences`` (in ``gigl.transforms.packed_layout``) turns the padded result into the offsets that ``gigl.nn.packed_attention`` needs to run the transformer on real tokens only. @@ -50,7 +50,7 @@ | | pack_sequences(lengths, max_seq_len) (host read 2: total tokens) v - flat_index + PackedSequences(cu_seqlens, max_seq_len) for gigl.nn.packed_attention + flat_index + PackedLayout(cu_seqlens, max_seq_len) for gigl.nn.packed_attention Constraints: - Memory is ``O(num_anchors * width ** hop_distance)``; see @@ -58,6 +58,7 @@ - ``cu_seqlens`` is int32, so the total token count must fit in int32. Example: + >>> from gigl.transforms.packed_layout import pack_sequences >>> edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]]) >>> layout, lengths, hop_bits = khop_layout( ... edge_index, @@ -82,11 +83,11 @@ [0, 3, 6] """ -from typing import Literal, NamedTuple +from typing import Literal import torch import torch.nn.functional as F -from jaxtyping import Int, Int32, Int64 +from jaxtyping import Int, Int64 from torch import Tensor # Maximum number of int64 cells ``khop_layout`` builds: the candidates, @@ -102,52 +103,6 @@ class KhopLayoutTooLargeError(ValueError): """The k-hop layout would have more candidates than the allowed maximum.""" -class PackedSequences(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 - - -def pack_sequences( - lengths: Int64[Tensor, "num_sequences"], max_seq_len: int -) -> tuple[Int64[Tensor, "num_tokens"], PackedSequences]: - """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. - - 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, PackedSequences(cu_seqlens.to(torch.int32), max_seq_len) - - def _build_neighbor_table( edge_index: Int[Tensor, "2 num_edges"], num_nodes: int, @@ -327,8 +282,14 @@ def khop_layout( - ``hop_bits``: Python int, ``hop_distance.bit_length()``. Raises: + ValueError: If ``sampling_direction`` is not ``"in"`` or ``"out"``. 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}'." + ) neighbor_table = _build_neighbor_table( edge_index, num_nodes, diff --git a/tests/unit/transforms/static_khop_test.py b/tests/unit/transforms/static_khop_test.py index c9721fcd4..72403da1b 100644 --- a/tests/unit/transforms/static_khop_test.py +++ b/tests/unit/transforms/static_khop_test.py @@ -12,7 +12,8 @@ _build_sequence_layout_from_sparse_neighbors, _get_k_hop_neighbors_sparse, ) -from gigl.transforms.static_khop import KhopLayoutTooLargeError, pack_sequences +from gigl.transforms.packed_layout import pack_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 @@ -204,6 +205,11 @@ def test_accepts_int32_edge_index_and_anchors(self) -> None: 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. From b3fec1496cfacb64f1720ca75270402d5d04ae2d Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 29 Oct 2026 17:41:46 +0000 Subject: [PATCH 4/9] Static k-hop: share the select-and-pack step as select_sequences; replace prioritize_hop_order with an ordering option select_sequences dedupes each anchor's candidates (the lowest column of a node wins), orders them by an optional priority and cuts them at max_seq_len, returning the nodes, the source column of every token and the lengths. The k-hop layout generates candidates through the neighbour table and calls it. khop_layout takes ordering in {node_id, hop, input} instead of prioritize_hop_order. Co-Authored-By: Claude Opus 5.5 --- gigl/transforms/packed_layout.py | 103 ++++++++++++++++++++- gigl/transforms/static_khop.py | 106 +++++++++++----------- tests/unit/transforms/static_khop_test.py | 91 ++++++++++++++++--- 3 files changed, 234 insertions(+), 66 deletions(-) diff --git a/gigl/transforms/packed_layout.py b/gigl/transforms/packed_layout.py index 83f1d623f..c0c59f574 100644 --- a/gigl/transforms/packed_layout.py +++ b/gigl/transforms/packed_layout.py @@ -3,13 +3,15 @@ ``pack_sequences`` turns a padded ``[num_sequences, max_seq_len]`` grid of sequences into the offsets (``PackedLayout``) that ``gigl.nn.packed_attention`` needs to run on real tokens only. It does not depend on how the sequences were -built. +built. ``select_sequences`` is the shared step that builds the padded grid from +per-anchor candidates. """ -from typing import NamedTuple +from typing import NamedTuple, Optional import torch -from jaxtyping import Int32, Int64 +import torch.nn.functional as F +from jaxtyping import Int, Int32, Int64 from torch import Tensor @@ -57,3 +59,98 @@ def pack_sequences( 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, +) -> tuple[ + Int64[Tensor, "num_anchors max_seq_len"], + Int64[Tensor, "num_anchors max_seq_len"], + Int64[Tensor, "num_anchors"], +]: + """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) 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. + + 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. + + Returns: + - ``nodes``: ``[num_anchors, max_seq_len]`` long node ids, the anchor in + column 0, -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, so at least 1 and at most ``max_seq_len``. + """ + 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) + & (nodes != anchor_indices.unsqueeze(1)) + ) + if priority is not None: + # 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) + # The anchor goes to column 0 below; the other kept candidates are numbered + # from 1 and cut at ``max_seq_len``. + positions = torch.cumsum(keep, dim=1) + keep = keep & (positions < max_seq_len) + # Dropped candidates scatter into column 0, which the anchor overwrites. + packed = torch.full( + (anchor_indices.numel(), max_seq_len), -1, dtype=torch.long, device=keys.device + ) + packed.scatter_(1, torch.where(keep, positions, 0), keys) + packed[:, 0] = anchor_indices << _COLUMN_BITS + source = torch.where(packed >= 0, packed & _COLUMN_MASK, -1) + source[:, 0] = -1 + return packed >> _COLUMN_BITS, source, 1 + keep.sum(dim=1) diff --git a/gigl/transforms/static_khop.py b/gigl/transforms/static_khop.py index c9eb91839..c194c3f9a 100644 --- a/gigl/transforms/static_khop.py +++ b/gigl/transforms/static_khop.py @@ -39,11 +39,11 @@ | hop h: gather the rows of the hop h - 1 frontier -> [num_anchors, width ** h] v candidates [num_anchors, 1 + width + ... + width ** hop_distance] - | sort each row by (node id, hop); the first copy of a node is its - | first-reach hop, so drop repeats, padding and the anchor itself + | 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, - with ``prioritize_hop_order``), cut at max_seq_len + 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 @@ -86,10 +86,11 @@ from typing import Literal import torch -import torch.nn.functional as F from jaxtyping import Int, Int64 from torch import Tensor +from gigl.transforms.packed_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 @@ -98,6 +99,8 @@ # over the limit and send the batch to the padded path. MAX_LAYOUT_CANDIDATES = 2**25 +KhopOrdering = Literal["node_id", "hop", "input"] + class KhopLayoutTooLargeError(ValueError): """The k-hop layout would have more candidates than the allowed maximum.""" @@ -176,19 +179,24 @@ def _khop_sequence_layout( anchor_indices: Int[Tensor, "num_anchors"], hop_distance: int, max_seq_len: int, - prioritize_hop_order: bool = False, + ordering: KhopOrdering = "node_id", ) -> tuple[Int64[Tensor, "num_anchors max_seq_len"], Int64[Tensor, "num_anchors"], int]: """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. 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 + 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 ``prioritize_hop_order`` 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="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. @@ -197,45 +205,32 @@ def _khop_sequence_layout( """ anchor_indices = anchor_indices.long() hop_bits = hop_distance.bit_length() - padding_key = torch.iinfo(torch.long).max + device = anchor_indices.device - # A candidate's sort key is ``(node << hop_bits) | hop``. It groups each node's - # copies with the smallest hop first, so the first copy carries the - # first-reach hop. Padding gets the largest key and sorts last. + # 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) - key_parts = [frontier << hop_bits] + 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) - key_parts.append( - torch.where(frontier >= 0, (frontier << hop_bits) | hop, padding_key) - ) - keys = torch.sort(torch.cat(key_parts, dim=1), dim=1).values - nodes = keys >> hop_bits - hops = keys & ((1 << hop_bits) - 1) - keep = (keys != padding_key) & F.pad(nodes.diff(dim=1) != 0, (1, 0), value=True) - if prioritize_hop_order: - # Re-sort the kept candidates by (hop, node id) so truncation below drops - # the farthest hops first. - node_radix = neighbor_table.size(0) - order_keys = torch.sort( - torch.where(keep, hops * node_radix + nodes, padding_key), dim=1 - ).values - keep = order_keys != padding_key - nodes = order_keys % node_radix - keys = (nodes << hop_bits) | (order_keys // node_radix) - # The anchor goes to column 0 below; the other kept candidates are numbered - # from 1 and cut at ``max_seq_len``. - keep = keep & (nodes != anchor_indices.unsqueeze(1)) - positions = torch.cumsum(keep, dim=1) - keep = keep & (positions < max_seq_len) - # Dropped candidates scatter into column 0, which the anchor overwrites. - layout = torch.full( - (anchor_indices.numel(), max_seq_len), -1, dtype=torch.long, device=keys.device + 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 ) - layout.scatter_(1, torch.where(keep, positions, 0), keys) - layout[:, 0] = anchor_indices << hop_bits - return layout, 1 + keep.sum(dim=1), hop_bits + hops = hop_of_column[source.clamp(min=0)] + layout = torch.where(nodes >= 0, (nodes << hop_bits) | hops, -1) + return layout, lengths, hop_bits def khop_layout( @@ -246,7 +241,7 @@ def khop_layout( max_seq_len: int, sampling_direction: Literal["in", "out"] = "out", max_candidates: int = MAX_LAYOUT_CANDIDATES, - prioritize_hop_order: bool = False, + ordering: KhopOrdering = "node_id", ) -> tuple[Int64[Tensor, "num_anchors max_seq_len"], Int64[Tensor, "num_anchors"], int]: """Builds the k-hop token layout of every anchor; the entry point of this module. @@ -266,15 +261,19 @@ def khop_layout( sampling_direction: ``"in"`` reverses the edges, ``"out"`` keeps them. max_candidates: Maximum number of candidates plus neighbour-table cells; see ``MAX_LAYOUT_CANDIDATES``. - prioritize_hop_order: Order neighbours by hop before node id. + 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: - ``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 ascending - node id (or by hop, then node id, with ``prioritize_hop_order``). The + 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, @@ -282,7 +281,8 @@ def khop_layout( - ``hop_bits``: Python int, ``hop_distance.bit_length()``. Raises: - ValueError: If ``sampling_direction`` is not ``"in"`` or ``"out"``. + 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"}: @@ -290,6 +290,10 @@ def khop_layout( "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, @@ -303,5 +307,5 @@ def khop_layout( anchor_indices, hop_distance, max_seq_len, - prioritize_hop_order, + ordering, ) diff --git a/tests/unit/transforms/static_khop_test.py b/tests/unit/transforms/static_khop_test.py index 72403da1b..1fa5dcff7 100644 --- a/tests/unit/transforms/static_khop_test.py +++ b/tests/unit/transforms/static_khop_test.py @@ -12,7 +12,7 @@ _build_sequence_layout_from_sparse_neighbors, _get_k_hop_neighbors_sparse, ) -from gigl.transforms.packed_layout import pack_sequences +from gigl.transforms.packed_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 @@ -66,12 +66,12 @@ def _sparse_layout( max_seq_len: int, direction: Literal["in", "out"], hop_distance: int = 2, - prioritize_hop_order: bool = False, + 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 prioritize_hop_order: + if ordering == "hop": node_index, valid = _build_hop_ordered_sequence_layout_from_sparse_edges( anchor_indices=anchors, edge_index=adjacency, @@ -104,7 +104,7 @@ def _assert_layouts_equal( max_seq_len: int, direction: Literal["in", "out"], hop_distance: int = 2, - prioritize_hop_order: bool = False, + ordering: Literal["node_id", "hop"] = "node_id", ) -> None: expected_nodes, expected_valid, expected_hops = _sparse_layout( edge_index, @@ -113,7 +113,7 @@ def _assert_layouts_equal( max_seq_len, direction, hop_distance, - prioritize_hop_order, + ordering, ) nodes, hops, lengths = khop_layout( edge_index, @@ -122,7 +122,7 @@ def _assert_layouts_equal( hop_distance=hop_distance, max_seq_len=max_seq_len, sampling_direction=direction, - prioritize_hop_order=prioritize_hop_order, + ordering=ordering, ) valid = torch.arange(max_seq_len).unsqueeze(0) < lengths.unsqueeze(1) self.assertTrue(torch.equal(valid, expected_valid)) @@ -132,16 +132,16 @@ def _assert_layouts_equal( def test_matches_sparse_layout(self) -> None: for ( hop_distance, - prioritize_hop_order, + ordering, direction, max_seq_len, seed, ) in itertools.product( - (1, 2, 3, 4), (False, True), ("in", "out"), (3, 4, 16, 64), range(2) + (1, 2, 3, 4), ("node_id", "hop"), ("in", "out"), (3, 4, 16, 64), range(2) ): with self.subTest( hop=hop_distance, - hop_order=prioritize_hop_order, + ordering=ordering, direction=direction, L=max_seq_len, seed=seed, @@ -153,13 +153,13 @@ def test_matches_sparse_layout(self) -> None: max_seq_len=max_seq_len, direction=direction, hop_distance=hop_distance, - prioritize_hop_order=prioritize_hop_order, + 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 prioritize_hop_order, expected in ((False, [0, 1, 2]), (True, [0, 5, 6])): + for ordering, expected in (("node_id", [0, 1, 2]), ("hop", [0, 5, 6])): nodes, hops, _ = khop_layout( edge_index, 7, @@ -167,7 +167,7 @@ def test_hop_order_truncation_drops_the_farthest_hops_first(self) -> None: hop_distance=2, max_seq_len=3, sampling_direction="out", - prioritize_hop_order=prioritize_hop_order, + ordering=ordering, ) self.assertEqual(nodes[0].tolist(), expected) @@ -230,6 +230,73 @@ def test_rejects_a_huge_hub_table_before_allocating_it(self) -> None: 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_does_not_recompile_across_batch_shapes(self) -> None: + torch._dynamo.reset() + 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_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]) From b7be346a0bc777b0ea0fa37a3177803321fb1aac Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 29 Oct 2026 20:16:12 +0000 Subject: [PATCH 5/9] Rename gigl.transforms.packed_layout to sequence_layout The module holds the pieces shared by the per-anchor sequence builders (selection, the packed type and packing), not only the packed layout. Co-Authored-By: Claude Opus 5.5 --- .../{packed_layout.py => sequence_layout.py} | 12 ++++++------ gigl/transforms/static_khop.py | 6 +++--- tests/unit/transforms/static_khop_test.py | 2 +- 3 files changed, 10 insertions(+), 10 deletions(-) rename gigl/transforms/{packed_layout.py => sequence_layout.py} (94%) diff --git a/gigl/transforms/packed_layout.py b/gigl/transforms/sequence_layout.py similarity index 94% rename from gigl/transforms/packed_layout.py rename to gigl/transforms/sequence_layout.py index c0c59f574..d44b06c3a 100644 --- a/gigl/transforms/packed_layout.py +++ b/gigl/transforms/sequence_layout.py @@ -1,10 +1,10 @@ -"""Packed token layout: where each sequence lives in a buffer without padding. +"""Shared pieces for building per-anchor token sequences. -``pack_sequences`` turns a padded ``[num_sequences, max_seq_len]`` grid of -sequences into the offsets (``PackedLayout``) that ``gigl.nn.packed_attention`` -needs to run on real tokens only. It does not depend on how the sequences were -built. ``select_sequences`` is the shared step that builds the padded grid from -per-anchor candidates. +``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 diff --git a/gigl/transforms/static_khop.py b/gigl/transforms/static_khop.py index c194c3f9a..0dceb749c 100644 --- a/gigl/transforms/static_khop.py +++ b/gigl/transforms/static_khop.py @@ -2,7 +2,7 @@ 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.packed_layout``) turns the padded +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. @@ -58,7 +58,7 @@ - ``cu_seqlens`` is int32, so the total token count must fit in int32. Example: - >>> from gigl.transforms.packed_layout import pack_sequences + >>> 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, @@ -89,7 +89,7 @@ from jaxtyping import Int, Int64 from torch import Tensor -from gigl.transforms.packed_layout import select_sequences +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 diff --git a/tests/unit/transforms/static_khop_test.py b/tests/unit/transforms/static_khop_test.py index 1fa5dcff7..0545cd07a 100644 --- a/tests/unit/transforms/static_khop_test.py +++ b/tests/unit/transforms/static_khop_test.py @@ -12,7 +12,7 @@ _build_sequence_layout_from_sparse_neighbors, _get_k_hop_neighbors_sparse, ) -from gigl.transforms.packed_layout import pack_sequences, select_sequences +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 From 74dae181fc616828e64b132f11fedc2f5253a81b Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 29 Oct 2026 20:21:13 +0000 Subject: [PATCH 6/9] Build the padded PPR sequences with select_sequences The builder pads each anchor's ragged PPR edges into a [anchors, edges] table and lets select_sequences order them by descending weight, then gathers the weights and relation features by the returned source column. select_sequences gains include_anchor (no anchor slot; the anchor stays if it is a candidate) and ties_by_column (priority ties keep column order instead of ascending node id). Output matches the previous builder on random PPR batches, and it is faster on CPU. Co-Authored-By: Claude Opus 5.5 --- gigl/transforms/graph_transformer.py | 156 ++++++------------ gigl/transforms/sequence_layout.py | 65 +++++--- .../unit/transforms/graph_transformer_test.py | 69 ++++++++ tests/unit/transforms/static_khop_test.py | 26 +++ 4 files changed, 192 insertions(+), 124 deletions(-) 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 index d44b06c3a..a0a08297a 100644 --- a/gigl/transforms/sequence_layout.py +++ b/gigl/transforms/sequence_layout.py @@ -73,6 +73,8 @@ def select_sequences( candidates: Int64[Tensor, "num_anchors num_candidates"], max_seq_len: int, priority: Optional[Tensor] = None, + ties_by_column: bool = False, + include_anchor: bool = True, ) -> tuple[ Int64[Tensor, "num_anchors max_seq_len"], Int64[Tensor, "num_anchors max_seq_len"], @@ -81,11 +83,12 @@ def select_sequences( """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) 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. + 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. Args: anchor_indices: ``[num_anchors]`` node ids of the anchors. @@ -97,16 +100,23 @@ def select_sequences( 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: - ``nodes``: ``[num_anchors, max_seq_len]`` long node ids, the anchor in - column 0, -1 from the sequence length on. + 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, so at least 1 and at most ``max_seq_len``. + included. It is at most ``max_seq_len``, except that it counts the + anchor even when ``max_seq_len`` is 0. """ anchor_indices = anchor_indices.long() num_candidates = candidates.size(1) @@ -120,12 +130,18 @@ def select_sequences( ) 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) - & (nodes != anchor_indices.unsqueeze(1)) - ) + 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. @@ -141,16 +157,23 @@ def select_sequences( ).indices keys = keys.gather(1, order) keep = keep.gather(1, order) - # The anchor goes to column 0 below; the other kept candidates are numbered - # from 1 and cut at ``max_seq_len``. - positions = torch.cumsum(keep, dim=1) + # 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 column 0, which the anchor overwrites. + # Dropped candidates scatter into an extra last column, which is sliced off. packed = torch.full( - (anchor_indices.numel(), max_seq_len), -1, dtype=torch.long, device=keys.device + (anchor_indices.numel(), max_seq_len + 1), + -1, + dtype=torch.long, + device=keys.device, ) - packed.scatter_(1, torch.where(keep, positions, 0), keys) - packed[:, 0] = anchor_indices << _COLUMN_BITS + 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) - source[:, 0] = -1 - return packed >> _COLUMN_BITS, source, 1 + keep.sum(dim=1) + if include_anchor and max_seq_len > 0: + source[:, 0] = -1 + return packed >> _COLUMN_BITS, source, first_position + keep.sum(dim=1) 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 index 0545cd07a..01f4b9b01 100644 --- a/tests/unit/transforms/static_khop_test.py +++ b/tests/unit/transforms/static_khop_test.py @@ -279,6 +279,32 @@ def test_select_sequences_orders_by_priority_and_keeps_the_first_copy(self) -> N 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() self.addCleanup(torch._dynamo.reset) From 80f7f5d7b79aec3209dadf437374bee5321e0c5c Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 30 Oct 2026 18:12:57 +0000 Subject: [PATCH 7/9] Reset dynamo counters in the no-recompile test The test counts unique graphs, which includes compiles from earlier tests in the same process. Co-Authored-By: Claude Opus 5.5 --- tests/unit/transforms/static_khop_test.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/unit/transforms/static_khop_test.py b/tests/unit/transforms/static_khop_test.py index 01f4b9b01..3043326c9 100644 --- a/tests/unit/transforms/static_khop_test.py +++ b/tests/unit/transforms/static_khop_test.py @@ -1,5 +1,6 @@ """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 @@ -7,6 +8,7 @@ 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, @@ -307,6 +309,7 @@ def test_select_sequences_can_tie_by_column_and_leave_out_the_anchor(self) -> No 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) @@ -323,6 +326,13 @@ def test_select_sequences_does_not_recompile_across_batch_shapes(self) -> None: # 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]) From 9a7f03c075677ab658d231f20364f4fa9a8744bf Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 30 Oct 2026 18:13:02 +0000 Subject: [PATCH 8/9] Return named tuples from select_sequences and khop_layout, with doctests SelectedSequences(nodes, source, lengths) and KhopLayout(layout, lengths, hop_bits) replace the bare tuples. select_sequences and pack_sequences get function-level examples; the test that runs the docstring examples came in with the previous commit. Co-Authored-By: Claude Opus 5.5 --- gigl/transforms/sequence_layout.py | 83 ++++++++++++++++++++++++------ gigl/transforms/static_khop.py | 42 +++++++++------ 2 files changed, 95 insertions(+), 30 deletions(-) diff --git a/gigl/transforms/sequence_layout.py b/gigl/transforms/sequence_layout.py index a0a08297a..ca6e90d43 100644 --- a/gigl/transforms/sequence_layout.py +++ b/gigl/transforms/sequence_layout.py @@ -32,6 +32,26 @@ class PackedLayout(NamedTuple): 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]: @@ -46,6 +66,13 @@ def pack_sequences( 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. @@ -75,11 +102,7 @@ def select_sequences( priority: Optional[Tensor] = None, ties_by_column: bool = False, include_anchor: bool = True, -) -> tuple[ - Int64[Tensor, "num_anchors max_seq_len"], - Int64[Tensor, "num_anchors max_seq_len"], - Int64[Tensor, "num_anchors"], -]: +) -> 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 @@ -90,6 +113,42 @@ def select_sequences( 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, @@ -108,15 +167,7 @@ def select_sequences( the anchor stays in it if it is one of them. Returns: - - ``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. + The ``SelectedSequences``: ``nodes``, ``source`` and ``lengths``. """ anchor_indices = anchor_indices.long() num_candidates = candidates.size(1) @@ -176,4 +227,6 @@ def select_sequences( source = torch.where(packed >= 0, packed & _COLUMN_MASK, -1) if include_anchor and max_seq_len > 0: source[:, 0] = -1 - return packed >> _COLUMN_BITS, source, first_position + keep.sum(dim=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 index 0dceb749c..9023cf157 100644 --- a/gigl/transforms/static_khop.py +++ b/gigl/transforms/static_khop.py @@ -83,7 +83,7 @@ [0, 3, 6] """ -from typing import Literal +from typing import Literal, NamedTuple import torch from jaxtyping import Int, Int64 @@ -102,6 +102,28 @@ 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.""" @@ -180,7 +202,7 @@ def _khop_sequence_layout( hop_distance: int, max_seq_len: int, ordering: KhopOrdering = "node_id", -) -> tuple[Int64[Tensor, "num_anchors max_seq_len"], Int64[Tensor, "num_anchors"], int]: +) -> 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 @@ -230,7 +252,7 @@ def _khop_sequence_layout( ) hops = hop_of_column[source.clamp(min=0)] layout = torch.where(nodes >= 0, (nodes << hop_bits) | hops, -1) - return layout, lengths, hop_bits + return KhopLayout(layout, lengths, hop_bits) def khop_layout( @@ -242,7 +264,7 @@ def khop_layout( sampling_direction: Literal["in", "out"] = "out", max_candidates: int = MAX_LAYOUT_CANDIDATES, ordering: KhopOrdering = "node_id", -) -> tuple[Int64[Tensor, "num_anchors max_seq_len"], Int64[Tensor, "num_anchors"], int]: +) -> 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 @@ -268,17 +290,7 @@ def khop_layout( parent's neighbour-table row, which follows the edge order). Returns: - - ``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()``. + The ``KhopLayout``: ``layout``, ``lengths`` and ``hop_bits``. Raises: ValueError: If ``sampling_direction`` is not ``"in"`` or ``"out"``, or From 97331fe41fb8a829dba4c3e94b8294f2e2add6a0 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Tue, 3 Nov 2026 23:26:06 +0000 Subject: [PATCH 9/9] Keep Jaxtyping's wrapper out of torch.compile in Shape Contract tests The Shape Contract wrapper shares one code object across every contracted function, and its typeguard checks ran inside compiled frames. Dynamo now skips that code object, and the checks are off while Dynamo traces. Co-Authored-By: Claude Opus 5.5 --- tests/test_assets/runtime_type_checking.py | 57 ++++++++++++++++++- .../test_assets/runtime_type_checking_test.py | 54 ++++++++++++++++++ 2 files changed, 110 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_assets/runtime_type_checking_test.py 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()