gigl.transforms.sequence_layout

Shared pieces for building per-anchor token sequences.

select_sequences picks each anchor’s tokens from its candidates into a padded [num_sequences, max_seq_len] grid, whatever generated the candidates (k-hop expansion, PPR edges). PackedLayout holds the offsets gigl.nn.packed_attention needs to run on real tokens only, and pack_sequences computes it from the grid’s sequence lengths.

Classes

PackedLayout	Where each sequence lives in the packed token buffer.
SelectedSequences	The tokens select_sequences picked for each anchor.

Functions

pack_sequences(lengths, max_seq_len)	Computes what is needed to drop the padding from a padded grid of sequences.
select_sequences(anchor_indices, candidates, max_seq_len)	Picks each anchor's tokens from its candidates, with fixed shapes.

Module Contents

class gigl.transforms.sequence_layout.PackedLayout

Bases: NamedTuple

Where each sequence lives in the packed token buffer.

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: jaxtyping.Int32[torch.Tensor, num_sequences_plus_one]

max_seq_len: int

class gigl.transforms.sequence_layout.SelectedSequences

Bases: NamedTuple

The tokens select_sequences picked for each anchor.

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.

lengths: jaxtyping.Int64[torch.Tensor, num_anchors]

nodes: jaxtyping.Int64[torch.Tensor, num_anchors max_seq_len]

source: jaxtyping.Int64[torch.Tensor, num_anchors max_seq_len]

gigl.transforms.sequence_layout.pack_sequences(lengths, max_seq_len)

Computes what is needed to drop the padding from a padded grid of sequences.

Sequence i is valid in its first lengths[i] positions. Any [num_sequences, max_seq_len] per-position tensor gathers into packed order with grid.flatten()[flat_index], and cu_seqlens marks where each packed sequence starts. This costs one host read, of the total token count; it is why the packed buffer holds exactly the real tokens.

A length above max_seq_len is clamped to it, because the grid has no positions beyond that. Checking for it would cost another host read.

Example

>>> flat_index, packed = pack_sequences(torch.tensor([2, 1, 3]), max_seq_len=3)
>>> flat_index.tolist()
[0, 1, 3, 6, 7, 8]
>>> packed.cu_seqlens.tolist()
[0, 2, 3, 6]

Parameters:

• max_seq_len (int) – Width of the padded grid.

• lengths (jaxtyping.Int64[torch.Tensor, num_sequences])

Returns:

• [num_tokens] long indices of the valid positions in the flattened grid, sequence by sequence.

• The packed layout.

Return type:

tuple[jaxtyping.Int64[torch.Tensor, num_tokens], PackedLayout]

gigl.transforms.sequence_layout.select_sequences(anchor_indices, candidates, max_seq_len, priority=None, ties_by_column=False, include_anchor=True)

Picks each anchor’s tokens from its candidates, with fixed shapes.

Every anchor’s sequence is the anchor, then its distinct candidate nodes other than the anchor, ordered by priority (ties by ascending node id, or by column with ties_by_column) and cut at max_seq_len. A node that appears in several columns keeps the lowest column, so put the preferred copy of a node first. This is the selection step shared by the k-hop and PPR sequence builders: they differ in how they generate the candidates and in what priority ranks.

Example

Two anchors, 0 and 1. Node 3 is a duplicate candidate of anchor 0, anchor 0 is among its own candidates, and -1 is padding:

>>> anchors = torch.tensor([0, 1])
>>> candidates = torch.tensor([[5, 3, 0, 3, -1], [4, 1, 6, 4, 2]])
>>> nodes, source, lengths = select_sequences(anchors, candidates, 4)
>>> nodes.tolist()
[[0, 3, 5, -1], [1, 2, 4, 6]]
>>> lengths.tolist()
[3, 4]

Gather a per-candidate payload with source, zeroing the anchor and padding tokens:

>>> payload = torch.arange(10, 20).reshape(2, 5)
>>> (payload.gather(1, source.clamp(min=0)) * (source >= 0)).tolist()
[[0, 11, 10, 0], [0, 19, 15, 17]]

A priority ranks the candidates (smaller first); equal priorities go by node id, or by column with ties_by_column:

>>> priority = torch.tensor([[0, 1, 0, 1, 0], [2, 9, 0, 5, 0]])
>>> select_sequences(anchors, candidates, 4, priority).nodes.tolist()
[[0, 5, 3, -1], [1, 2, 6, 4]]
>>> select_sequences(
...     anchors, candidates, 4, priority, ties_by_column=True
... ).nodes.tolist()
[[0, 5, 3, -1], [1, 6, 2, 4]]

Without include_anchor the candidates are the whole sequence, so anchor 0 stays in its sequence as the candidate it is:

>>> select_sequences(anchors, candidates, 4, include_anchor=False).source.tolist()
[[2, 1, 0, -1], [1, 4, 0, 2]]

Parameters:

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

• candidates (jaxtyping.Int64[torch.Tensor, num_anchors num_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 (int) – Tokens per sequence, anchor included.

• priority (Optional[torch.Tensor]) – [num_anchors, num_candidates] integer or float rank of each candidate; smaller comes first. If None, candidates are ordered by ascending node id, which costs one sort instead of two. Pass the column index to keep the candidates’ own order.

• ties_by_column (bool) – Whether candidates with equal priority stay in column order instead of ascending node id. It has no effect without priority.

• include_anchor (bool) – 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, source and lengths.

Return type:

The SelectedSequences