gigl.transforms.static_khop#

Builds each anchor’s k-hop token sequence with fixed-shape tensor ops.

This is the layout half of the packed backend of GraphTransformerEncoder. For every anchor it picks which nodes become its tokens (the anchor, then its neighbours within hop_distance hops), and pack_sequences (in gigl.transforms.sequence_layout) turns the padded result into the offsets that gigl.nn.packed_attention needs to run the transformer on real tokens only.

Why this exists:

  • Sequences are short and uneven (about 90 real tokens out of a 256-token budget in our models), so a padded [num_sequences, max_seq_len] batch spends most of its compute and memory on padding. Packing the real tokens into one buffer, with cu_seqlens marking the boundaries, avoids that.

  • The sparse k-hop path in gigl.transforms.graph_transformer (sparse matmul, coalesce, unique, nonzero, bool-mask indexing) produces tensors whose sizes depend on the data. It needs many host syncs and blocks torch.compile. This module instead uses a dense neighbour table and only sort, gather and scatter, whose output shapes are known up front. It produces the same layout as that path’s default: anchor first, then neighbours in ascending node id.

  • A dense table has one cost: the candidates before deduplication grow as num_anchors * width ** hop_distance, and one hub node sets width for the whole batch. khop_layout counts the candidates before it builds anything and raises KhopLayoutTooLargeError past max_candidates. GraphTransformerEncoder catches that and runs the batch on the padded path.

  • The whole path makes exactly two host reads (each is a GPU sync): the table width (the largest per-node degree, all edge types combined) and, in pack_sequences, the total token count.

How a layout is built:

edge_index [2, num_edges]
    |  sort edges by source; an edge's slot is its rank in its source's group
    v
neighbour table [num_nodes + 1, width]    width = largest degree (host read 1)
    |                                     row i = nodes one hop from i, -1 pads;
    |                                     the extra last row is all -1
    |  hop h: gather the rows of the hop h - 1 frontier -> [num_anchors, width ** h]
    v
candidates [num_anchors, 1 + width + ... + width ** hop_distance]
    |  select_sequences: the first copy of a node is its first-reach hop, so
    |  drop repeats, padding and the anchor itself
    v
anchor first, then neighbours in ascending node id (or by hop, then node id,
or in expansion order, per ``ordering``), cut at max_seq_len
    |
    v
layout [num_anchors, max_seq_len] of (node << hop_bits) | hop (-1 at padding), lengths
    |
    |  pack_sequences(lengths, max_seq_len)       (host read 2: total tokens)
    v
flat_index + PackedLayout(cu_seqlens, max_seq_len)  for gigl.nn.packed_attention
Constraints:

  • Memory is O(num_anchors * width ** hop_distance); see MAX_LAYOUT_CANDIDATES for the limit.

  • cu_seqlens is int32, so the total token count must fit in int32.

Example

>>> from gigl.transforms.sequence_layout import pack_sequences
>>> edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]])
>>> layout, lengths, hop_bits = khop_layout(
...     edge_index,
...     num_nodes=5,
...     anchor_indices=torch.arange(2),
...     hop_distance=2,
...     max_seq_len=4,
...     sampling_direction="out",
... )
>>> hop_bits
2
>>> (layout >> hop_bits).tolist()
[[0, 1, 2, -1], [1, 2, 3, -1]]
>>> (layout & 3).tolist()
[[0, 1, 2, 3], [0, 1, 2, 3]]
>>> lengths.tolist()
[3, 3]
>>> flat_index, packed = pack_sequences(lengths, max_seq_len=4)
>>> (layout.flatten()[flat_index] >> hop_bits).tolist()
[0, 1, 2, 1, 2, 3]
>>> packed.cu_seqlens.tolist()
[0, 3, 6]

Attributes#

KhopOrdering

MAX_LAYOUT_CANDIDATES

Exceptions#

KhopLayoutTooLargeError

The k-hop layout would have more candidates than the allowed maximum.

Classes#

KhopLayout

The k-hop token layout of every anchor, as built by khop_layout.

Functions#

khop_layout(edge_index, num_nodes, anchor_indices, ...)

Builds the k-hop token layout of every anchor; the entry point of this module.

Module Contents#

exception gigl.transforms.static_khop.KhopLayoutTooLargeError[source]#

Bases: ValueError

The k-hop layout would have more candidates than the allowed maximum.

Initialize self. See help(type(self)) for accurate signature.

class gigl.transforms.static_khop.KhopLayout[source]#

Bases: NamedTuple

The k-hop token layout of every anchor, as built by khop_layout.

layout[source]#

[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[source]#

[num_anchors] long count of valid tokens per anchor, anchor included, so at least 1 and at most max_seq_len.

hop_bits[source]#

Python int, hop_distance.bit_length().

hop_bits: int[source]#
layout: jaxtyping.Int64[torch.Tensor, num_anchors max_seq_len][source]#
lengths: jaxtyping.Int64[torch.Tensor, num_anchors][source]#
gigl.transforms.static_khop.khop_layout(edge_index, num_nodes, anchor_indices, hop_distance, max_seq_len, sampling_direction='out', max_candidates=MAX_LAYOUT_CANDIDATES, ordering='node_id')[source]#

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.

Parameters:

  • edge_index (jaxtyping.Int[torch.Tensor, 2 num_edges]) – [2, num_edges] int32 or int64 local node ids in [0, num_nodes).

  • num_nodes (int) – Number of nodes.

  • anchor_indices (jaxtyping.Int[torch.Tensor, num_anchors]) – [num_anchors] int32 or int64 local node ids of the anchors.

  • hop_distance (int) – Number of hops to expand.

  • max_seq_len (int) – Tokens per sequence, anchor included.

  • sampling_direction (Literal['in', 'out']) – "in" reverses the edges, "out" keeps them.

  • max_candidates (int) – Maximum number of candidates plus neighbour-table cells; see MAX_LAYOUT_CANDIDATES.

  • ordering (KhopOrdering) – 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, lengths and hop_bits.

Return type:

The KhopLayout

Raises:

  • ValueError – If sampling_direction is not "in" or "out", or ordering is not "node_id", "hop" or "input".

  • KhopLayoutTooLargeError – If the table cells plus candidates exceed max_candidates.

gigl.transforms.static_khop.KhopOrdering[source]#
gigl.transforms.static_khop.MAX_LAYOUT_CANDIDATES = 33554432[source]#