From f56ce027b39b3238f99679526eb743119dbc1c74 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Mon, 3 Aug 2026 19:02:55 +0000 Subject: [PATCH 01/13] Optimize anchor-only graph transformer inference --- gigl/nn/graph_transformer.py | 161 +++++++++++++++++++++++- tests/unit/nn/graph_transformer_test.py | 114 +++++++++++++++++ 2 files changed, 274 insertions(+), 1 deletion(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 6270b407f..71c6d9f80 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -533,6 +533,133 @@ def forward( return x + def forward_anchor_only( + self, + x: Tensor, + attn_bias: Optional[Tensor] = None, + valid_mask: Optional[Tensor] = None, + pairwise_relation_indices: Optional[Tensor] = None, + ) -> Tensor: + """Compute the final layer output for the anchor token only. + + Keys and values still cover the complete sequence because every token + can contribute to the anchor. Query-side attention, relation messages, + output projection, and feed-forward work are restricted to position + zero because later token outputs cannot affect the anchor. + + Args: + x: Input tensor of shape ``(batch, seq, model_dim)``. + attn_bias: Optional attention bias broadcastable to + ``(batch, num_heads, seq, seq)``. + valid_mask: Optional boolean tensor of shape ``(batch, seq)``. + pairwise_relation_indices: Optional sparse relation coordinates + shaped ``(num_relation_edges, 4)``. + + Returns: + Anchor output of shape ``(batch, 1, model_dim)``. + + Raises: + ValueError: If relation-aware attention is enabled. Its square + query/key bias construction requires the full layer path. + """ + if self._relation_attention_mode != "none": + raise ValueError( + "Anchor-only final-layer execution does not support " + "relation-aware attention." + ) + + batch_size, seq_len, model_dim = x.shape + anchor_valid_mask = valid_mask[:, :1] if valid_mask is not None else None + residual_anchor = x[:, :1, :] + x_norm = self._attention_norm(x) + + query = self._query_projection(x_norm[:, :1, :]) + key = self._key_projection(x_norm) + value = self._value_projection(x_norm) + + query = query.view(batch_size, 1, self._num_heads, self._head_dim).transpose( + 1, 2 + ) + key = key.view(batch_size, seq_len, self._num_heads, self._head_dim).transpose( + 1, 2 + ) + value = value.view( + batch_size, seq_len, self._num_heads, self._head_dim + ).transpose(1, 2) + + anchor_attn_bias = attn_bias + if anchor_attn_bias is not None and anchor_attn_bias.size(-2) == seq_len: + anchor_attn_bias = anchor_attn_bias[..., :1, :] + attention_output = self._run_attention( + query=query, + key=key, + value=value, + attn_bias=anchor_attn_bias, + pairwise_relation_indices=pairwise_relation_indices, + ) + attention_output = attention_output.transpose(1, 2).reshape( + batch_size, 1, model_dim + ) + attention_output = self._dropout(self._output_projection(attention_output)) + anchor = residual_anchor + attention_output + + anchor_relation_indices = pairwise_relation_indices + if self._relation_message_mode != "none": + if pairwise_relation_indices is None: + raise ValueError( + "pairwise_relation_indices is required when " + "relation_message_mode is relation-aware." + ) + if ( + pairwise_relation_indices.dim() != 2 + or pairwise_relation_indices.size(-1) != 4 + ): + raise ValueError( + "pairwise_relation_indices must have shape (num_relation_edges, 4)." + ) + if pairwise_relation_indices.numel() > 0: + relation_indices = pairwise_relation_indices[:, 3] + if ( + relation_indices.min().item() < 0 + or relation_indices.max().item() >= self._num_relations + ): + raise ValueError( + "pairwise_relation_indices contains relation ids outside " + f"[0, {self._num_relations})." + ) + anchor_relation_indices = pairwise_relation_indices[ + pairwise_relation_indices[:, 1] == 0 + ] + if self._relation_message_mode == "edge_type_attention": + anchor = anchor + self._dropout( + self._compute_relation_attention_messages( + x_norm=x_norm, + query=query, + key=key, + pairwise_relation_indices=anchor_relation_indices, + batch_size=batch_size, + seq_len=1, + ) + ) + elif self._relation_message_mode != "none": + anchor = anchor + self._dropout( + self._compute_relation_messages( + x_norm=x_norm, + pairwise_relation_indices=anchor_relation_indices, + batch_size=batch_size, + seq_len=1, + ) + ) + if anchor_valid_mask is not None: + anchor = anchor * anchor_valid_mask.unsqueeze(-1).to(anchor.dtype) + + residual_anchor = anchor + anchor = residual_anchor + self._ffn(self._ffn_norm(anchor)) + if anchor_valid_mask is not None: + anchor = anchor * anchor_valid_mask.unsqueeze(-1).to(anchor.dtype) + + return anchor + def _run_attention( self, query: Tensor, @@ -1849,7 +1976,22 @@ def _encode_and_readout( """ x = sequences * valid_mask.unsqueeze(-1).to(sequences.dtype) - for encoder_layer in self._encoder_layers: + encoder_layers = self._encoder_layers + use_anchor_only_final_layer = ( + self._readout_mode == "anchor_only" + and not self.training + and len(encoder_layers) > 0 + and encoder_layers[-1]._relation_attention_mode == "none" + ) + num_full_sequence_layers = len(encoder_layers) - int( + use_anchor_only_final_layer + ) + for layer_index in range(num_full_sequence_layers): + encoder_layer = encoder_layers[layer_index] + if not isinstance(encoder_layer, GraphTransformerEncoderLayer): + raise TypeError( + "Graph transformer encoder contains an unexpected layer type." + ) x = encoder_layer( x, attn_bias=attn_bias, @@ -1857,6 +1999,23 @@ def _encode_and_readout( valid_mask=valid_mask, ) + if use_anchor_only_final_layer: + final_encoder_layer = encoder_layers[-1] + if not isinstance(final_encoder_layer, GraphTransformerEncoderLayer): + raise TypeError( + "Graph transformer encoder contains an unexpected layer type." + ) + x = final_encoder_layer.forward_anchor_only( + x, + attn_bias=attn_bias, + pairwise_relation_indices=pairwise_relation_indices, + valid_mask=valid_mask, + ) + anchor_valid_mask = valid_mask[:, :1] + x = self._final_norm(x) + x = x * anchor_valid_mask.unsqueeze(-1).to(x.dtype) + return x.squeeze(1) + x = self._final_norm(x) x = x * valid_mask.unsqueeze(-1).to(x.dtype) diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index 49d1b4fa6..be334c6be 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -437,6 +437,27 @@ def _pairwise_relation_indices(coords: list[tuple[int, int, int, int]]) -> Tenso return torch.tensor(coords, dtype=torch.long) +def _full_sequence_anchor_reference( + encoder: GraphTransformerEncoder, + sequences: Tensor, + valid_mask: Tensor, + attn_bias: Tensor | None = None, + pairwise_relation_indices: Tensor | None = None, +) -> Tensor: + """Evaluate the full sequence before selecting the anchor token.""" + x = sequences * valid_mask.unsqueeze(-1).to(sequences.dtype) + for encoder_layer in encoder._encoder_layers: + x = encoder_layer( + x, + attn_bias=attn_bias, + valid_mask=valid_mask, + pairwise_relation_indices=pairwise_relation_indices, + ) + x = encoder._final_norm(x) + x = x * valid_mask.unsqueeze(-1).to(x.dtype) + return x[:, 0, :] + + class TestGraphTransformerEncoderPEModes(TestCase): def setUp(self) -> None: self._node_type = NodeType("user") @@ -463,6 +484,99 @@ def _create_encoder(self, **kwargs: object) -> GraphTransformerEncoder: defaults.update(kwargs) return GraphTransformerEncoder(**defaults) + def test_anchor_only_final_layer_matches_full_sequence_relation_messages( + self, + ) -> None: + """Anchor specialization preserves every relation-message mode.""" + sequences = torch.randn(2, 4, 8) + valid_mask = torch.tensor( + [[True, True, True, False], [True, True, False, False]] + ) + attn_bias = 0.1 * torch.randn(2, 2, 4, 4) + relation_indices = _pairwise_relation_indices( + [ + (0, 0, 1, 0), + (0, 0, 2, 0), + (0, 1, 2, 0), + (1, 0, 1, 0), + (1, 1, 0, 0), + ] + ) + + for relation_message_mode in [ + "none", + "edge_type_linear", + "edge_type_attention", + ]: + with self.subTest(relation_message_mode=relation_message_mode): + torch.manual_seed(0) + encoder = self._create_encoder( + num_layers=2, + readout_mode="anchor_only", + relation_message_mode=relation_message_mode, + ) + encoder.eval() + with torch.no_grad(): + for encoder_layer in encoder._encoder_layers: + if encoder_layer._relation_message_matrices is not None: + relation_message_matrices = ( + encoder_layer._relation_message_matrices + ) + assert isinstance(relation_message_matrices, Tensor) + relation_message_matrices.normal_() + expected = _full_sequence_anchor_reference( + encoder=encoder, + sequences=sequences, + valid_mask=valid_mask, + attn_bias=attn_bias, + pairwise_relation_indices=relation_indices, + ) + actual = encoder._encode_and_readout( + sequences=sequences, + valid_mask=valid_mask, + attn_bias=attn_bias, + pairwise_relation_indices=relation_indices, + ) + + self.assertTrue( + torch.allclose(actual, expected, atol=1e-6, rtol=1e-5), + f"relation_message_mode={relation_message_mode}", + ) + + def test_anchor_only_relation_attention_uses_full_sequence_path(self) -> None: + """Relation-aware attention retains its square full-sequence path.""" + sequences = torch.randn(2, 4, 8) + valid_mask = torch.tensor( + [[True, True, True, False], [True, True, False, False]] + ) + relation_indices = _pairwise_relation_indices( + [(0, 0, 1, 0), (0, 1, 2, 0), (1, 0, 1, 0)] + ) + + for relation_attention_mode in ["edge_type_bilinear", "edge_type_hgt"]: + with self.subTest(relation_attention_mode=relation_attention_mode): + torch.manual_seed(0) + encoder = self._create_encoder( + num_layers=2, + readout_mode="anchor_only", + relation_attention_mode=relation_attention_mode, + ) + encoder.eval() + with torch.no_grad(): + expected = _full_sequence_anchor_reference( + encoder=encoder, + sequences=sequences, + valid_mask=valid_mask, + pairwise_relation_indices=relation_indices, + ) + actual = encoder._encode_and_readout( + sequences=sequences, + valid_mask=valid_mask, + pairwise_relation_indices=relation_indices, + ) + + self.assertTrue(torch.equal(actual, expected)) + def test_additive_mode_matches_base_encoder_when_node_pe_projection_is_zero( self, ) -> None: From 6989558f7d01325e7d0d46fc0ae1b14479a432f6 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Mon, 3 Aug 2026 20:16:50 +0000 Subject: [PATCH 02/13] Support anchor-only relation attention --- gigl/nn/graph_transformer.py | 87 ++++++++++++------------- tests/unit/nn/graph_transformer_test.py | 32 +++++++-- 2 files changed, 69 insertions(+), 50 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 71c6d9f80..5a6c71145 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -558,21 +558,43 @@ def forward_anchor_only( Returns: Anchor output of shape ``(batch, 1, model_dim)``. - Raises: - ValueError: If relation-aware attention is enabled. Its square - query/key bias construction requires the full layer path. """ - if self._relation_attention_mode != "none": - raise ValueError( - "Anchor-only final-layer execution does not support " - "relation-aware attention." - ) - batch_size, seq_len, model_dim = x.shape anchor_valid_mask = valid_mask[:, :1] if valid_mask is not None else None residual_anchor = x[:, :1, :] x_norm = self._attention_norm(x) + anchor_relation_indices = pairwise_relation_indices + if ( + self._relation_attention_mode != "none" + or self._relation_message_mode != "none" + ): + if pairwise_relation_indices is None: + raise ValueError( + "pairwise_relation_indices is required when relation-aware " + "attention or messages are enabled." + ) + if ( + pairwise_relation_indices.dim() != 2 + or pairwise_relation_indices.size(-1) != 4 + ): + raise ValueError( + "pairwise_relation_indices must have shape (num_relation_edges, 4)." + ) + if pairwise_relation_indices.numel() > 0: + relation_indices = pairwise_relation_indices[:, 3] + if ( + relation_indices.min().item() < 0 + or relation_indices.max().item() >= self._num_relations + ): + raise ValueError( + "pairwise_relation_indices contains relation ids outside " + f"[0, {self._num_relations})." + ) + anchor_relation_indices = pairwise_relation_indices[ + pairwise_relation_indices[:, 1] == 0 + ] + query = self._query_projection(x_norm[:, :1, :]) key = self._key_projection(x_norm) value = self._value_projection(x_norm) @@ -595,7 +617,7 @@ def forward_anchor_only( key=key, value=value, attn_bias=anchor_attn_bias, - pairwise_relation_indices=pairwise_relation_indices, + pairwise_relation_indices=anchor_relation_indices, ) attention_output = attention_output.transpose(1, 2).reshape( batch_size, 1, model_dim @@ -603,33 +625,6 @@ def forward_anchor_only( attention_output = self._dropout(self._output_projection(attention_output)) anchor = residual_anchor + attention_output - anchor_relation_indices = pairwise_relation_indices - if self._relation_message_mode != "none": - if pairwise_relation_indices is None: - raise ValueError( - "pairwise_relation_indices is required when " - "relation_message_mode is relation-aware." - ) - if ( - pairwise_relation_indices.dim() != 2 - or pairwise_relation_indices.size(-1) != 4 - ): - raise ValueError( - "pairwise_relation_indices must have shape (num_relation_edges, 4)." - ) - if pairwise_relation_indices.numel() > 0: - relation_indices = pairwise_relation_indices[:, 3] - if ( - relation_indices.min().item() < 0 - or relation_indices.max().item() >= self._num_relations - ): - raise ValueError( - "pairwise_relation_indices contains relation ids outside " - f"[0, {self._num_relations})." - ) - anchor_relation_indices = pairwise_relation_indices[ - pairwise_relation_indices[:, 1] == 0 - ] if self._relation_message_mode == "edge_type_attention": anchor = anchor + self._dropout( self._compute_relation_attention_messages( @@ -697,8 +692,9 @@ def _build_relation_attention_bias( ): return None - batch_size, _, seq_len, _ = query.shape - empty_bias = query.new_zeros((batch_size, self._num_heads, seq_len, seq_len)) + batch_size, _, query_len, _ = query.shape + key_len = key.size(2) + empty_bias = query.new_zeros((batch_size, self._num_heads, query_len, key_len)) return self._add_relation_attention_bias( attn_bias=empty_bias, query=query, @@ -745,15 +741,15 @@ def _add_relation_attention_bias( f"[0, {self._num_relations})." ) - batch_size, _, seq_len, _ = query.shape + batch_size, _, query_len, _ = query.shape + key_len = key.size(2) + bias_shape = (self._num_heads, query_len, key_len) if attn_bias is None: - attn_bias = query.new_zeros((batch_size, self._num_heads, seq_len, seq_len)) - elif attn_bias.shape[1:] != (self._num_heads, seq_len, seq_len): + attn_bias = query.new_zeros((batch_size, *bias_shape)) + elif attn_bias.shape[1:] != bias_shape: attn_bias = attn_bias.expand( batch_size, - self._num_heads, - seq_len, - seq_len, + *bias_shape, ).clone() else: attn_bias = attn_bias.clone() @@ -1981,7 +1977,6 @@ def _encode_and_readout( self._readout_mode == "anchor_only" and not self.training and len(encoder_layers) > 0 - and encoder_layers[-1]._relation_attention_mode == "none" ) num_full_sequence_layers = len(encoder_layers) - int( use_anchor_only_final_layer diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index be334c6be..6e825fc8b 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -543,14 +543,21 @@ def test_anchor_only_final_layer_matches_full_sequence_relation_messages( f"relation_message_mode={relation_message_mode}", ) - def test_anchor_only_relation_attention_uses_full_sequence_path(self) -> None: - """Relation-aware attention retains its square full-sequence path.""" + def test_anchor_only_final_layer_matches_relation_attention(self) -> None: + """Anchor specialization preserves rectangular relation attention.""" sequences = torch.randn(2, 4, 8) valid_mask = torch.tensor( [[True, True, True, False], [True, True, False, False]] ) + attn_bias = 0.1 * torch.randn(2, 2, 4, 4) relation_indices = _pairwise_relation_indices( - [(0, 0, 1, 0), (0, 1, 2, 0), (1, 0, 1, 0)] + [ + (0, 0, 1, 0), + (0, 0, 2, 0), + (0, 1, 2, 0), + (1, 0, 1, 0), + (1, 1, 0, 0), + ] ) for relation_attention_mode in ["edge_type_bilinear", "edge_type_hgt"]: @@ -560,22 +567,39 @@ def test_anchor_only_relation_attention_uses_full_sequence_path(self) -> None: num_layers=2, readout_mode="anchor_only", relation_attention_mode=relation_attention_mode, + relation_message_mode="edge_type_attention", ) encoder.eval() with torch.no_grad(): + for encoder_layer in encoder._encoder_layers: + relation_parameters = [ + encoder_layer._relation_attention_matrices, + encoder_layer._relation_hgt_attention_matrices, + encoder_layer._relation_hgt_attention_priors, + encoder_layer._relation_message_matrices, + ] + for relation_parameter in relation_parameters: + if relation_parameter is not None: + assert isinstance(relation_parameter, Tensor) + relation_parameter.normal_() expected = _full_sequence_anchor_reference( encoder=encoder, sequences=sequences, valid_mask=valid_mask, + attn_bias=attn_bias, pairwise_relation_indices=relation_indices, ) actual = encoder._encode_and_readout( sequences=sequences, valid_mask=valid_mask, + attn_bias=attn_bias, pairwise_relation_indices=relation_indices, ) - self.assertTrue(torch.equal(actual, expected)) + self.assertTrue( + torch.allclose(actual, expected, atol=1e-6, rtol=1e-5), + f"relation_attention_mode={relation_attention_mode}", + ) def test_additive_mode_matches_base_encoder_when_node_pe_projection_is_zero( self, From 8925668414855696427b9bcc27c901d2481ff57f Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Mon, 3 Aug 2026 23:17:40 +0000 Subject: [PATCH 03/13] Enable anchor-only final layer during training --- gigl/nn/graph_transformer.py | 31 +++- tests/unit/nn/graph_transformer_test.py | 182 ++++++++++++++++++++++++ 2 files changed, 210 insertions(+), 3 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 5a6c71145..020f792b5 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -648,6 +648,15 @@ def forward_anchor_only( if anchor_valid_mask is not None: anchor = anchor * anchor_valid_mask.unsqueeze(-1).to(anchor.dtype) + if ( + self.training + and pairwise_relation_indices is not None + and pairwise_relation_indices.numel() > 0 + and anchor_relation_indices is not None + and anchor_relation_indices.numel() == 0 + ): + anchor = anchor + self._zero_relation_parameter_dependency(anchor) + residual_anchor = anchor anchor = residual_anchor + self._ffn(self._ffn_norm(anchor)) if anchor_valid_mask is not None: @@ -655,6 +664,24 @@ def forward_anchor_only( return anchor + def _zero_relation_parameter_dependency(self, reference: Tensor) -> Tensor: + """Keep relation parameters in the training graph after anchor filtering.""" + zero_dependency = reference.new_zeros(()) + relation_parameters = [ + self._relation_attention_matrices, + self._relation_hgt_attention_matrices, + self._relation_hgt_attention_priors, + self._relation_message_matrices, + self._relation_message_attention_matrices, + self._relation_message_attention_priors, + ] + for parameter in relation_parameters: + if parameter is not None: + zero_dependency = zero_dependency + ( + parameter.reshape(-1)[0].to(dtype=reference.dtype) * 0.0 + ) + return zero_dependency + def _run_attention( self, query: Tensor, @@ -1974,9 +2001,7 @@ def _encode_and_readout( encoder_layers = self._encoder_layers use_anchor_only_final_layer = ( - self._readout_mode == "anchor_only" - and not self.training - and len(encoder_layers) > 0 + self._readout_mode == "anchor_only" and len(encoder_layers) > 0 ) num_full_sequence_layers = len(encoder_layers) - int( use_anchor_only_final_layer diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index 6e825fc8b..4f0c10acc 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -1,5 +1,6 @@ """Tests for GraphTransformerEncoder.""" +import copy from typing import Literal, cast import torch @@ -601,6 +602,187 @@ def test_anchor_only_final_layer_matches_relation_attention(self) -> None: f"relation_attention_mode={relation_attention_mode}", ) + def test_anchor_only_training_matches_full_sequence_gradients(self) -> None: + """Anchor specialization preserves training outputs and gradients.""" + valid_mask = torch.tensor( + [[True, True, True, False], [True, True, False, False]] + ) + relation_indices = _pairwise_relation_indices( + [ + (0, 0, 1, 0), + (0, 0, 2, 0), + (0, 1, 2, 0), + (1, 0, 1, 0), + (1, 1, 0, 0), + ] + ) + relation_modes = [ + ("none", "none"), + ("none", "edge_type_linear"), + ("none", "edge_type_attention"), + ("edge_type_bilinear", "none"), + ("edge_type_hgt", "edge_type_attention"), + ] + + for num_layers in [1, 2]: + for relation_attention_mode, relation_message_mode in relation_modes: + with self.subTest( + num_layers=num_layers, + relation_attention_mode=relation_attention_mode, + relation_message_mode=relation_message_mode, + ): + torch.manual_seed(0) + encoder = self._create_encoder( + num_layers=num_layers, + readout_mode="anchor_only", + relation_attention_mode=relation_attention_mode, + relation_message_mode=relation_message_mode, + ) + with torch.no_grad(): + for encoder_layer in encoder._encoder_layers: + relation_parameters = [ + encoder_layer._relation_attention_matrices, + encoder_layer._relation_hgt_attention_matrices, + encoder_layer._relation_hgt_attention_priors, + encoder_layer._relation_message_matrices, + ] + for relation_parameter in relation_parameters: + if relation_parameter is not None: + assert isinstance(relation_parameter, Tensor) + relation_parameter.normal_() + + full_encoder = copy.deepcopy(encoder).train() + optimized_encoder = copy.deepcopy(encoder).train() + torch.manual_seed(1) + full_sequences = torch.randn(2, 4, 8, requires_grad=True) + optimized_sequences = ( + full_sequences.detach().clone().requires_grad_() + ) + full_attn_bias = (0.1 * torch.randn(2, 2, 4, 4)).requires_grad_() + optimized_attn_bias = ( + full_attn_bias.detach().clone().requires_grad_() + ) + + expected = _full_sequence_anchor_reference( + encoder=full_encoder, + sequences=full_sequences, + valid_mask=valid_mask, + attn_bias=full_attn_bias, + pairwise_relation_indices=relation_indices, + ) + actual = optimized_encoder._encode_and_readout( + sequences=optimized_sequences, + valid_mask=valid_mask, + attn_bias=optimized_attn_bias, + pairwise_relation_indices=relation_indices, + ) + upstream_gradient = torch.randn_like(expected) + expected.backward(upstream_gradient) + actual.backward(upstream_gradient) + + torch.testing.assert_close(actual, expected) + torch.testing.assert_close( + optimized_sequences.grad, + full_sequences.grad, + ) + torch.testing.assert_close( + optimized_attn_bias.grad, + full_attn_bias.grad, + ) + assert optimized_sequences.grad is not None + self.assertGreater( + optimized_sequences.grad[:, 1:, :].abs().sum().item(), + 0.0, + ) + + full_parameters = dict(full_encoder.named_parameters()) + optimized_parameters = dict(optimized_encoder.named_parameters()) + self.assertEqual( + full_parameters.keys(), optimized_parameters.keys() + ) + for name, full_parameter in full_parameters.items(): + optimized_parameter = optimized_parameters[name] + self.assertEqual( + optimized_parameter.grad is None, + full_parameter.grad is None, + name, + ) + if full_parameter.grad is not None: + torch.testing.assert_close( + optimized_parameter.grad, + full_parameter.grad, + msg=lambda message, name=name: f"{name}: {message}", + ) + + def test_anchor_only_training_preserves_zero_relation_gradients(self) -> None: + """Filtered non-anchor relations remain visible to DDP and optimizers.""" + valid_mask = torch.tensor( + [[True, True, True, False], [True, True, False, False]] + ) + relation_indices = _pairwise_relation_indices( + [ + (0, 1, 2, 0), + (0, 2, 1, 0), + (1, 1, 0, 0), + ] + ) + relation_modes = [ + ("none", "edge_type_linear"), + ("none", "edge_type_attention"), + ("edge_type_bilinear", "none"), + ("edge_type_hgt", "edge_type_attention"), + ] + + for relation_attention_mode, relation_message_mode in relation_modes: + with self.subTest( + relation_attention_mode=relation_attention_mode, + relation_message_mode=relation_message_mode, + ): + torch.manual_seed(0) + encoder = self._create_encoder( + num_layers=1, + readout_mode="anchor_only", + relation_attention_mode=relation_attention_mode, + relation_message_mode=relation_message_mode, + ) + full_encoder = copy.deepcopy(encoder).train() + optimized_encoder = copy.deepcopy(encoder).train() + torch.manual_seed(1) + full_sequences = torch.randn(2, 4, 8, requires_grad=True) + optimized_sequences = full_sequences.detach().clone().requires_grad_() + + expected = _full_sequence_anchor_reference( + encoder=full_encoder, + sequences=full_sequences, + valid_mask=valid_mask, + pairwise_relation_indices=relation_indices, + ) + actual = optimized_encoder._encode_and_readout( + sequences=optimized_sequences, + valid_mask=valid_mask, + pairwise_relation_indices=relation_indices, + ) + upstream_gradient = torch.randn_like(expected) + expected.backward(upstream_gradient) + actual.backward(upstream_gradient) + + torch.testing.assert_close(actual, expected) + full_parameters = dict(full_encoder.named_parameters()) + optimized_parameters = dict(optimized_encoder.named_parameters()) + for name, full_parameter in full_parameters.items(): + optimized_parameter = optimized_parameters[name] + self.assertEqual( + optimized_parameter.grad is None, + full_parameter.grad is None, + name, + ) + if full_parameter.grad is not None: + torch.testing.assert_close( + optimized_parameter.grad, + full_parameter.grad, + msg=lambda message, name=name: f"{name}: {message}", + ) + def test_additive_mode_matches_base_encoder_when_node_pe_projection_is_zero( self, ) -> None: From cf484f4a18e98e318028e300e7808c893bc1da39 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 5 Aug 2026 19:01:57 +0000 Subject: [PATCH 04/13] Simplify anchor-only transformer execution --- gigl/nn/graph_transformer.py | 412 +++++++++++------------- tests/unit/nn/graph_transformer_test.py | 79 ++++- 2 files changed, 264 insertions(+), 227 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 020f792b5..5a4bfc393 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -463,77 +463,15 @@ def forward( Returns: Output tensor of shape ``(batch, seq, model_dim)``. """ - batch_size, seq_len, model_dim = x.shape - - # Self-attention block (pre-norm) - residual = x - x_norm = self._attention_norm(x) - - query = self._query_projection(x_norm) - key = self._key_projection(x_norm) - value = self._value_projection(x_norm) - - # Reshape to (batch, num_heads, seq, head_dim) - query = query.view( - batch_size, seq_len, self._num_heads, self._head_dim - ).transpose(1, 2) - key = key.view(batch_size, seq_len, self._num_heads, self._head_dim).transpose( - 1, 2 - ) - value = value.view( - batch_size, seq_len, self._num_heads, self._head_dim - ).transpose(1, 2) - - attention_output = self._run_attention( - query=query, - key=key, - value=value, + return self._forward_query_prefix( + x=x, + query_length=x.size(1), attn_bias=attn_bias, + valid_mask=valid_mask, pairwise_relation_indices=pairwise_relation_indices, ) - # Reshape back to (batch, seq, model_dim) - attention_output = attention_output.transpose(1, 2).reshape( - batch_size, seq_len, model_dim - ) - attention_output = self._output_projection(attention_output) - attention_output = self._dropout(attention_output) - - x = residual + attention_output - if self._relation_message_mode == "edge_type_attention": - x = x + self._dropout( - self._compute_relation_attention_messages( - x_norm=x_norm, - query=query, - key=key, - pairwise_relation_indices=pairwise_relation_indices, - batch_size=batch_size, - seq_len=seq_len, - ) - ) - elif self._relation_message_mode != "none": - x = x + self._dropout( - self._compute_relation_messages( - x_norm=x_norm, - pairwise_relation_indices=pairwise_relation_indices, - batch_size=batch_size, - seq_len=seq_len, - ) - ) - if valid_mask is not None: - x = x * valid_mask.unsqueeze(-1).to(x.dtype) - - # Feed-forward block (pre-norm) - residual = x - x_norm = self._ffn_norm(x) - ffn_output = self._ffn(x_norm) - x = residual + ffn_output - if valid_mask is not None: - x = x * valid_mask.unsqueeze(-1).to(x.dtype) - - return x - - def forward_anchor_only( + def _forward_anchor_only( self, x: Tensor, attn_bias: Optional[Tensor] = None, @@ -557,51 +495,80 @@ def forward_anchor_only( Returns: Anchor output of shape ``(batch, 1, model_dim)``. + """ + return self._forward_query_prefix( + x=x, + query_length=1, + attn_bias=attn_bias, + valid_mask=valid_mask, + pairwise_relation_indices=pairwise_relation_indices, + ) + def _forward_query_prefix( + self, + x: Tensor, + query_length: int, + attn_bias: Optional[Tensor], + valid_mask: Optional[Tensor], + pairwise_relation_indices: Optional[Tensor], + ) -> Tensor: + """Compute layer outputs for a prefix of query positions. + + Keys and values retain the full sequence so every token can contribute + to each requested query output. + + Args: + x: Input tensor of shape ``(batch, seq, model_dim)``. + query_length: Number of leading query positions to compute. + attn_bias: Optional attention bias broadcastable to + ``(batch, num_heads, query_length, seq)``. + valid_mask: Optional boolean tensor of shape ``(batch, seq)``. + pairwise_relation_indices: Optional sparse relation coordinates + shaped ``(num_relation_edges, 4)``. + + Returns: + Output tensor of shape ``(batch, query_length, model_dim)``. + + Raises: + ValueError: If ``query_length`` is outside ``[1, seq]`` or relation + coordinates are invalid. """ batch_size, seq_len, model_dim = x.shape - anchor_valid_mask = valid_mask[:, :1] if valid_mask is not None else None - residual_anchor = x[:, :1, :] - x_norm = self._attention_norm(x) + if query_length < 1 or query_length > seq_len: + raise ValueError( + f"query_length must be in [1, {seq_len}], got {query_length}." + ) - anchor_relation_indices = pairwise_relation_indices + output_valid_mask = ( + valid_mask[:, :query_length] if valid_mask is not None else None + ) + output = x[:, :query_length, :] + + selected_relation_indices = pairwise_relation_indices if ( self._relation_attention_mode != "none" or self._relation_message_mode != "none" ): - if pairwise_relation_indices is None: - raise ValueError( - "pairwise_relation_indices is required when relation-aware " - "attention or messages are enabled." - ) - if ( - pairwise_relation_indices.dim() != 2 - or pairwise_relation_indices.size(-1) != 4 - ): - raise ValueError( - "pairwise_relation_indices must have shape (num_relation_edges, 4)." - ) - if pairwise_relation_indices.numel() > 0: - relation_indices = pairwise_relation_indices[:, 3] - if ( - relation_indices.min().item() < 0 - or relation_indices.max().item() >= self._num_relations - ): - raise ValueError( - "pairwise_relation_indices contains relation ids outside " - f"[0, {self._num_relations})." - ) - anchor_relation_indices = pairwise_relation_indices[ - pairwise_relation_indices[:, 1] == 0 - ] + selected_relation_indices = self._validate_pairwise_relation_indices( + pairwise_relation_indices=pairwise_relation_indices, + device=x.device, + batch_size=batch_size, + query_length=seq_len, + key_length=seq_len, + ) + if query_length < seq_len: + selected_relation_indices = selected_relation_indices[ + selected_relation_indices[:, 1] < query_length + ] - query = self._query_projection(x_norm[:, :1, :]) + x_norm = self._attention_norm(x) + query = self._query_projection(x_norm[:, :query_length, :]) key = self._key_projection(x_norm) value = self._value_projection(x_norm) - query = query.view(batch_size, 1, self._num_heads, self._head_dim).transpose( - 1, 2 - ) + query = query.view( + batch_size, query_length, self._num_heads, self._head_dim + ).transpose(1, 2) key = key.view(batch_size, seq_len, self._num_heads, self._head_dim).transpose( 1, 2 ) @@ -609,60 +576,107 @@ def forward_anchor_only( batch_size, seq_len, self._num_heads, self._head_dim ).transpose(1, 2) - anchor_attn_bias = attn_bias - if anchor_attn_bias is not None and anchor_attn_bias.size(-2) == seq_len: - anchor_attn_bias = anchor_attn_bias[..., :1, :] + if attn_bias is not None and attn_bias.dim() >= 2: + attn_bias = attn_bias[..., :query_length, :] attention_output = self._run_attention( query=query, key=key, value=value, - attn_bias=anchor_attn_bias, - pairwise_relation_indices=anchor_relation_indices, + attn_bias=attn_bias, + pairwise_relation_indices=selected_relation_indices, ) attention_output = attention_output.transpose(1, 2).reshape( - batch_size, 1, model_dim + batch_size, query_length, model_dim ) attention_output = self._dropout(self._output_projection(attention_output)) - anchor = residual_anchor + attention_output + output = output + attention_output if self._relation_message_mode == "edge_type_attention": - anchor = anchor + self._dropout( + output = output + self._dropout( self._compute_relation_attention_messages( x_norm=x_norm, query=query, key=key, - pairwise_relation_indices=anchor_relation_indices, + pairwise_relation_indices=selected_relation_indices, batch_size=batch_size, - seq_len=1, + seq_len=query_length, ) ) elif self._relation_message_mode != "none": - anchor = anchor + self._dropout( + output = output + self._dropout( self._compute_relation_messages( x_norm=x_norm, - pairwise_relation_indices=anchor_relation_indices, + pairwise_relation_indices=selected_relation_indices, batch_size=batch_size, - seq_len=1, + seq_len=query_length, ) ) - if anchor_valid_mask is not None: - anchor = anchor * anchor_valid_mask.unsqueeze(-1).to(anchor.dtype) + if output_valid_mask is not None: + output = output * output_valid_mask.unsqueeze(-1).to(output.dtype) if ( self.training and pairwise_relation_indices is not None and pairwise_relation_indices.numel() > 0 - and anchor_relation_indices is not None - and anchor_relation_indices.numel() == 0 + and selected_relation_indices is not None + and selected_relation_indices.numel() == 0 ): - anchor = anchor + self._zero_relation_parameter_dependency(anchor) + output = output + self._zero_relation_parameter_dependency(output) - residual_anchor = anchor - anchor = residual_anchor + self._ffn(self._ffn_norm(anchor)) - if anchor_valid_mask is not None: - anchor = anchor * anchor_valid_mask.unsqueeze(-1).to(anchor.dtype) + output = output + self._ffn(self._ffn_norm(output)) + if output_valid_mask is not None: + output = output * output_valid_mask.unsqueeze(-1).to(output.dtype) - return anchor + return output + + def _validate_pairwise_relation_indices( + self, + pairwise_relation_indices: Optional[Tensor], + device: torch.device, + batch_size: int, + query_length: int, + key_length: int, + ) -> Tensor: + """Normalize and validate sparse relation coordinates.""" + if pairwise_relation_indices is None: + raise ValueError( + "pairwise_relation_indices is required for relation-aware " + "attention or messages." + ) + if ( + pairwise_relation_indices.dim() != 2 + or pairwise_relation_indices.size(-1) != 4 + ): + raise ValueError( + "pairwise_relation_indices must have shape (num_relation_edges, 4)." + ) + + pairwise_relation_indices = pairwise_relation_indices.to( + device=device, + dtype=torch.long, + ) + if pairwise_relation_indices.numel() == 0: + return pairwise_relation_indices + + coordinate_names_and_limits = [ + ("batch", pairwise_relation_indices[:, 0], batch_size), + ("query", pairwise_relation_indices[:, 1], query_length), + ("key", pairwise_relation_indices[:, 2], key_length), + ("relation", pairwise_relation_indices[:, 3], self._num_relations), + ] + for coordinate_name, coordinates, upper_bound in coordinate_names_and_limits: + if coordinates.min().item() < 0 or coordinates.max().item() >= upper_bound: + if coordinate_name == "relation": + raise ValueError( + "pairwise_relation_indices contains relation ids outside " + f"[0, {upper_bound})." + ) + raise ValueError( + "pairwise_relation_indices contains " + f"{coordinate_name} indices outside [0, {upper_bound})." + ) + + return pairwise_relation_indices def _zero_relation_parameter_dependency(self, reference: Tensor) -> Tensor: """Keep relation parameters in the training graph after anchor filtering.""" @@ -713,14 +727,18 @@ def _build_relation_attention_bias( key: Tensor, pairwise_relation_indices: Optional[Tensor], ) -> Optional[Tensor]: - if ( - pairwise_relation_indices is not None - and pairwise_relation_indices.numel() == 0 - ): - return None - batch_size, _, query_len, _ = query.shape key_len = key.size(2) + pairwise_relation_indices = self._validate_pairwise_relation_indices( + pairwise_relation_indices=pairwise_relation_indices, + device=query.device, + batch_size=batch_size, + query_length=query_len, + key_length=key_len, + ) + if pairwise_relation_indices.numel() == 0: + return None + empty_bias = query.new_zeros((batch_size, self._num_heads, query_len, key_len)) return self._add_relation_attention_bias( attn_bias=empty_bias, @@ -736,40 +754,22 @@ def _add_relation_attention_bias( key: Tensor, pairwise_relation_indices: Optional[Tensor], ) -> Optional[Tensor]: - if pairwise_relation_indices is None: - raise ValueError( - "pairwise_relation_indices is required when " - "relation_attention_mode is relation-aware." - ) - if pairwise_relation_indices.numel() == 0: - return attn_bias - if ( - pairwise_relation_indices.dim() != 2 - or pairwise_relation_indices.size(-1) != 4 - ): - raise ValueError( - "pairwise_relation_indices must have shape (num_relation_edges, 4)." - ) - - pairwise_relation_indices = pairwise_relation_indices.to( + batch_size, _, query_len, _ = query.shape + key_len = key.size(2) + pairwise_relation_indices = self._validate_pairwise_relation_indices( + pairwise_relation_indices=pairwise_relation_indices, device=query.device, - dtype=torch.long, + batch_size=batch_size, + query_length=query_len, + key_length=key_len, ) + if pairwise_relation_indices.numel() == 0: + return attn_bias batch_indices = pairwise_relation_indices[:, 0] query_indices = pairwise_relation_indices[:, 1] key_indices = pairwise_relation_indices[:, 2] relation_indices = pairwise_relation_indices[:, 3] - if ( - relation_indices.min().item() < 0 - or relation_indices.max().item() >= self._num_relations - ): - raise ValueError( - "pairwise_relation_indices contains relation ids outside " - f"[0, {self._num_relations})." - ) - batch_size, _, query_len, _ = query.shape - key_len = key.size(2) bias_shape = (self._num_heads, query_len, key_len) if attn_bias is None: attn_bias = query.new_zeros((batch_size, *bias_shape)) @@ -893,11 +893,13 @@ def _compute_relation_messages( """ if self._relation_message_matrices is None: raise ValueError("Relation message matrices are not initialized.") - if pairwise_relation_indices is None: - raise ValueError( - "pairwise_relation_indices is required when " - "relation_message_mode is relation-aware." - ) + pairwise_relation_indices = self._validate_pairwise_relation_indices( + pairwise_relation_indices=pairwise_relation_indices, + device=x_norm.device, + batch_size=batch_size, + query_length=seq_len, + key_length=x_norm.size(1), + ) messages = torch.zeros( (batch_size, seq_len, x_norm.size(-1)), dtype=x_norm.dtype, @@ -905,30 +907,10 @@ def _compute_relation_messages( ) if pairwise_relation_indices.numel() == 0: return messages - if ( - pairwise_relation_indices.dim() != 2 - or pairwise_relation_indices.size(-1) != 4 - ): - raise ValueError( - "pairwise_relation_indices must have shape (num_relation_edges, 4)." - ) - - pairwise_relation_indices = pairwise_relation_indices.to( - device=x_norm.device, - dtype=torch.long, - ) batch_indices = pairwise_relation_indices[:, 0] target_indices = pairwise_relation_indices[:, 1] source_indices = pairwise_relation_indices[:, 2] relation_indices = pairwise_relation_indices[:, 3] - if ( - relation_indices.min().item() < 0 - or relation_indices.max().item() >= self._num_relations - ): - raise ValueError( - "pairwise_relation_indices contains relation ids outside " - f"[0, {self._num_relations})." - ) # Per-(relation, batch, target) in-degree for mean aggregation, computed # globally up front so grouping order of the entries does not matter. @@ -1017,11 +999,13 @@ def _compute_relation_attention_messages( raise ValueError("Relation message attention matrices are not initialized.") if self._relation_message_attention_priors is None: raise ValueError("Relation message attention priors are not initialized.") - if pairwise_relation_indices is None: - raise ValueError( - "pairwise_relation_indices is required when " - "relation_message_mode is relation-aware." - ) + pairwise_relation_indices = self._validate_pairwise_relation_indices( + pairwise_relation_indices=pairwise_relation_indices, + device=x_norm.device, + batch_size=batch_size, + query_length=seq_len, + key_length=x_norm.size(1), + ) messages = torch.zeros( (batch_size, seq_len, x_norm.size(-1)), dtype=x_norm.dtype, @@ -1029,30 +1013,10 @@ def _compute_relation_attention_messages( ) if pairwise_relation_indices.numel() == 0: return messages - if ( - pairwise_relation_indices.dim() != 2 - or pairwise_relation_indices.size(-1) != 4 - ): - raise ValueError( - "pairwise_relation_indices must have shape (num_relation_edges, 4)." - ) - - pairwise_relation_indices = pairwise_relation_indices.to( - device=x_norm.device, - dtype=torch.long, - ) batch_indices = pairwise_relation_indices[:, 0] target_indices = pairwise_relation_indices[:, 1] source_indices = pairwise_relation_indices[:, 2] relation_indices = pairwise_relation_indices[:, 3] - if ( - relation_indices.min().item() < 0 - or relation_indices.max().item() >= self._num_relations - ): - raise ValueError( - "pairwise_relation_indices contains relation ids outside " - f"[0, {self._num_relations})." - ) # Gather per-edge queries/keys: (num_edges, num_heads, head_dim). edge_queries = query[batch_indices, :, target_indices] @@ -2003,15 +1967,18 @@ def _encode_and_readout( use_anchor_only_final_layer = ( self._readout_mode == "anchor_only" and len(encoder_layers) > 0 ) - num_full_sequence_layers = len(encoder_layers) - int( - use_anchor_only_final_layer + final_encoder_layer = ( + cast(GraphTransformerEncoderLayer, encoder_layers[-1]) + if use_anchor_only_final_layer + else None ) - for layer_index in range(num_full_sequence_layers): - encoder_layer = encoder_layers[layer_index] - if not isinstance(encoder_layer, GraphTransformerEncoderLayer): - raise TypeError( - "Graph transformer encoder contains an unexpected layer type." - ) + for encoder_layer_module in encoder_layers: + encoder_layer = cast( + GraphTransformerEncoderLayer, + encoder_layer_module, + ) + if encoder_layer is final_encoder_layer: + break x = encoder_layer( x, attn_bias=attn_bias, @@ -2019,25 +1986,18 @@ def _encode_and_readout( valid_mask=valid_mask, ) - if use_anchor_only_final_layer: - final_encoder_layer = encoder_layers[-1] - if not isinstance(final_encoder_layer, GraphTransformerEncoderLayer): - raise TypeError( - "Graph transformer encoder contains an unexpected layer type." - ) - x = final_encoder_layer.forward_anchor_only( + output_valid_mask = valid_mask + if final_encoder_layer is not None: + x = final_encoder_layer._forward_anchor_only( x, attn_bias=attn_bias, pairwise_relation_indices=pairwise_relation_indices, valid_mask=valid_mask, ) - anchor_valid_mask = valid_mask[:, :1] - x = self._final_norm(x) - x = x * anchor_valid_mask.unsqueeze(-1).to(x.dtype) - return x.squeeze(1) + output_valid_mask = valid_mask[:, :1] x = self._final_norm(x) - x = x * valid_mask.unsqueeze(-1).to(x.dtype) + x = x * output_valid_mask.unsqueeze(-1).to(x.dtype) # Readout: anchor (position 0) + attention-weighted neighbor aggregation anchor = x[:, 0, :].unsqueeze(1) # (batch, 1, hid_dim) diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index 4f0c10acc..cd7f158f6 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -1,7 +1,7 @@ """Tests for GraphTransformerEncoder.""" import copy -from typing import Literal, cast +from typing import Callable, Literal, cast import torch import torch.nn as nn @@ -485,6 +485,83 @@ def _create_encoder(self, **kwargs: object) -> GraphTransformerEncoder: defaults.update(kwargs) return GraphTransformerEncoder(**defaults) + def test_anchor_only_final_layer_restricts_query_side_work(self) -> None: + """Anchor specialization keeps full keys and values but one query.""" + encoder = self._create_encoder( + num_layers=2, + readout_mode="anchor_only", + ) + first_layer = cast( + GraphTransformerEncoderLayer, + encoder._encoder_layers[0], + ) + final_layer = cast( + GraphTransformerEncoderLayer, + encoder._encoder_layers[-1], + ) + input_shapes: dict[str, tuple[int, ...]] = {} + + def record_input_shape( + name: str, + ) -> Callable[[nn.Module, tuple[Tensor, ...]], None]: + def hook(_module: nn.Module, inputs: tuple[Tensor, ...]) -> None: + input_shapes[name] = tuple(inputs[0].shape) + + return hook + + hooks = [ + first_layer._query_projection.register_forward_pre_hook( + record_input_shape("first_query") + ), + final_layer._query_projection.register_forward_pre_hook( + record_input_shape("final_query") + ), + final_layer._key_projection.register_forward_pre_hook( + record_input_shape("final_key") + ), + final_layer._value_projection.register_forward_pre_hook( + record_input_shape("final_value") + ), + final_layer._output_projection.register_forward_pre_hook( + record_input_shape("final_output") + ), + final_layer._ffn.register_forward_pre_hook(record_input_shape("final_ffn")), + ] + try: + output = encoder._encode_and_readout( + sequences=torch.randn(2, 4, 8), + valid_mask=torch.tensor( + [[True, True, True, False], [True, True, False, False]] + ), + ) + finally: + for hook in hooks: + hook.remove() + + self.assertEqual(output.shape, (2, 8)) + self.assertEqual(input_shapes["first_query"], (2, 4, 8)) + self.assertEqual(input_shapes["final_query"], (2, 1, 8)) + self.assertEqual(input_shapes["final_key"], (2, 4, 8)) + self.assertEqual(input_shapes["final_value"], (2, 4, 8)) + self.assertEqual(input_shapes["final_output"], (2, 1, 8)) + self.assertEqual(input_shapes["final_ffn"], (2, 1, 8)) + + def test_anchor_only_rejects_invalid_non_anchor_relation_coordinates( + self, + ) -> None: + """Validate all relation coordinates before filtering anchor rows.""" + encoder = self._create_encoder( + readout_mode="anchor_only", + relation_attention_mode="edge_type_bilinear", + ) + + with self.assertRaisesRegex(ValueError, "query indices outside"): + encoder._encode_and_readout( + sequences=torch.randn(1, 2, 8), + valid_mask=torch.ones((1, 2), dtype=torch.bool), + pairwise_relation_indices=_pairwise_relation_indices([(0, 2, 1, 0)]), + ) + def test_anchor_only_final_layer_matches_full_sequence_relation_messages( self, ) -> None: From db1a1dd9c06c0dd7799ee4198f8634f408280e67 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 5 Aug 2026 19:14:46 +0000 Subject: [PATCH 05/13] Simplify shared relation validation --- gigl/nn/graph_transformer.py | 135 ++++++++---------------- tests/unit/nn/graph_transformer_test.py | 26 +---- 2 files changed, 47 insertions(+), 114 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 5a4bfc393..2a070b76b 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -549,13 +549,32 @@ def _forward_query_prefix( self._relation_attention_mode != "none" or self._relation_message_mode != "none" ): - selected_relation_indices = self._validate_pairwise_relation_indices( - pairwise_relation_indices=pairwise_relation_indices, + if pairwise_relation_indices is None: + raise ValueError( + "pairwise_relation_indices is required for relation-aware " + "attention or messages." + ) + if ( + pairwise_relation_indices.dim() != 2 + or pairwise_relation_indices.size(-1) != 4 + ): + raise ValueError( + "pairwise_relation_indices must have shape (num_relation_edges, 4)." + ) + selected_relation_indices = pairwise_relation_indices.to( device=x.device, - batch_size=batch_size, - query_length=seq_len, - key_length=seq_len, - ) + dtype=torch.long, + ) + if selected_relation_indices.numel() > 0: + relation_indices = selected_relation_indices[:, 3] + if ( + relation_indices.min().item() < 0 + or relation_indices.max().item() >= self._num_relations + ): + raise ValueError( + "pairwise_relation_indices contains relation ids outside " + f"[0, {self._num_relations})." + ) if query_length < seq_len: selected_relation_indices = selected_relation_indices[ selected_relation_indices[:, 1] < query_length @@ -597,7 +616,10 @@ def _forward_query_prefix( x_norm=x_norm, query=query, key=key, - pairwise_relation_indices=selected_relation_indices, + pairwise_relation_indices=cast( + Tensor, + selected_relation_indices, + ), batch_size=batch_size, seq_len=query_length, ) @@ -606,7 +628,10 @@ def _forward_query_prefix( output = output + self._dropout( self._compute_relation_messages( x_norm=x_norm, - pairwise_relation_indices=selected_relation_indices, + pairwise_relation_indices=cast( + Tensor, + selected_relation_indices, + ), batch_size=batch_size, seq_len=query_length, ) @@ -629,55 +654,6 @@ def _forward_query_prefix( return output - def _validate_pairwise_relation_indices( - self, - pairwise_relation_indices: Optional[Tensor], - device: torch.device, - batch_size: int, - query_length: int, - key_length: int, - ) -> Tensor: - """Normalize and validate sparse relation coordinates.""" - if pairwise_relation_indices is None: - raise ValueError( - "pairwise_relation_indices is required for relation-aware " - "attention or messages." - ) - if ( - pairwise_relation_indices.dim() != 2 - or pairwise_relation_indices.size(-1) != 4 - ): - raise ValueError( - "pairwise_relation_indices must have shape (num_relation_edges, 4)." - ) - - pairwise_relation_indices = pairwise_relation_indices.to( - device=device, - dtype=torch.long, - ) - if pairwise_relation_indices.numel() == 0: - return pairwise_relation_indices - - coordinate_names_and_limits = [ - ("batch", pairwise_relation_indices[:, 0], batch_size), - ("query", pairwise_relation_indices[:, 1], query_length), - ("key", pairwise_relation_indices[:, 2], key_length), - ("relation", pairwise_relation_indices[:, 3], self._num_relations), - ] - for coordinate_name, coordinates, upper_bound in coordinate_names_and_limits: - if coordinates.min().item() < 0 or coordinates.max().item() >= upper_bound: - if coordinate_name == "relation": - raise ValueError( - "pairwise_relation_indices contains relation ids outside " - f"[0, {upper_bound})." - ) - raise ValueError( - "pairwise_relation_indices contains " - f"{coordinate_name} indices outside [0, {upper_bound})." - ) - - return pairwise_relation_indices - def _zero_relation_parameter_dependency(self, reference: Tensor) -> Tensor: """Keep relation parameters in the training graph after anchor filtering.""" zero_dependency = reference.new_zeros(()) @@ -709,7 +685,10 @@ def _run_attention( attn_bias=attn_bias, query=query, key=key, - pairwise_relation_indices=pairwise_relation_indices, + pairwise_relation_indices=cast( + Tensor, + pairwise_relation_indices, + ), ) return F.scaled_dot_product_attention( @@ -725,20 +704,13 @@ def _build_relation_attention_bias( self, query: Tensor, key: Tensor, - pairwise_relation_indices: Optional[Tensor], + pairwise_relation_indices: Tensor, ) -> Optional[Tensor]: - batch_size, _, query_len, _ = query.shape - key_len = key.size(2) - pairwise_relation_indices = self._validate_pairwise_relation_indices( - pairwise_relation_indices=pairwise_relation_indices, - device=query.device, - batch_size=batch_size, - query_length=query_len, - key_length=key_len, - ) if pairwise_relation_indices.numel() == 0: return None + batch_size, _, query_len, _ = query.shape + key_len = key.size(2) empty_bias = query.new_zeros((batch_size, self._num_heads, query_len, key_len)) return self._add_relation_attention_bias( attn_bias=empty_bias, @@ -752,17 +724,10 @@ def _add_relation_attention_bias( attn_bias: Optional[Tensor], query: Tensor, key: Tensor, - pairwise_relation_indices: Optional[Tensor], + pairwise_relation_indices: Tensor, ) -> Optional[Tensor]: batch_size, _, query_len, _ = query.shape key_len = key.size(2) - pairwise_relation_indices = self._validate_pairwise_relation_indices( - pairwise_relation_indices=pairwise_relation_indices, - device=query.device, - batch_size=batch_size, - query_length=query_len, - key_length=key_len, - ) if pairwise_relation_indices.numel() == 0: return attn_bias batch_indices = pairwise_relation_indices[:, 0] @@ -873,7 +838,7 @@ def _compute_relation_attention_scores( def _compute_relation_messages( self, x_norm: Tensor, - pairwise_relation_indices: Optional[Tensor], + pairwise_relation_indices: Tensor, batch_size: int, seq_len: int, ) -> Tensor: @@ -893,13 +858,6 @@ def _compute_relation_messages( """ if self._relation_message_matrices is None: raise ValueError("Relation message matrices are not initialized.") - pairwise_relation_indices = self._validate_pairwise_relation_indices( - pairwise_relation_indices=pairwise_relation_indices, - device=x_norm.device, - batch_size=batch_size, - query_length=seq_len, - key_length=x_norm.size(1), - ) messages = torch.zeros( (batch_size, seq_len, x_norm.size(-1)), dtype=x_norm.dtype, @@ -963,7 +921,7 @@ def _compute_relation_attention_messages( x_norm: Tensor, query: Tensor, key: Tensor, - pairwise_relation_indices: Optional[Tensor], + pairwise_relation_indices: Tensor, batch_size: int, seq_len: int, ) -> Tensor: @@ -999,13 +957,6 @@ def _compute_relation_attention_messages( raise ValueError("Relation message attention matrices are not initialized.") if self._relation_message_attention_priors is None: raise ValueError("Relation message attention priors are not initialized.") - pairwise_relation_indices = self._validate_pairwise_relation_indices( - pairwise_relation_indices=pairwise_relation_indices, - device=x_norm.device, - batch_size=batch_size, - query_length=seq_len, - key_length=x_norm.size(1), - ) messages = torch.zeros( (batch_size, seq_len, x_norm.size(-1)), dtype=x_norm.dtype, diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index cd7f158f6..3684ae4e4 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -546,22 +546,6 @@ def hook(_module: nn.Module, inputs: tuple[Tensor, ...]) -> None: self.assertEqual(input_shapes["final_output"], (2, 1, 8)) self.assertEqual(input_shapes["final_ffn"], (2, 1, 8)) - def test_anchor_only_rejects_invalid_non_anchor_relation_coordinates( - self, - ) -> None: - """Validate all relation coordinates before filtering anchor rows.""" - encoder = self._create_encoder( - readout_mode="anchor_only", - relation_attention_mode="edge_type_bilinear", - ) - - with self.assertRaisesRegex(ValueError, "query indices outside"): - encoder._encode_and_readout( - sequences=torch.randn(1, 2, 8), - valid_mask=torch.ones((1, 2), dtype=torch.bool), - pairwise_relation_indices=_pairwise_relation_indices([(0, 2, 1, 0)]), - ) - def test_anchor_only_final_layer_matches_full_sequence_relation_messages( self, ) -> None: @@ -1776,9 +1760,8 @@ def test_relation_attention_rejects_invalid_relation_ids(self) -> None: ) with self.assertRaisesRegex(ValueError, "relation ids outside"): - layer._build_relation_attention_bias( - query=torch.zeros((1, 1, 2, 2)), - key=torch.zeros((1, 1, 2, 2)), + layer( + x=torch.zeros((1, 2, 2)), pairwise_relation_indices=_pairwise_relation_indices([(0, 1, 0, 1)]), ) @@ -1794,9 +1777,8 @@ def test_relation_attention_hgt_rejects_invalid_relation_ids(self) -> None: ) with self.assertRaisesRegex(ValueError, "relation ids outside"): - layer._build_relation_attention_bias( - query=torch.zeros((1, 1, 2, 2)), - key=torch.zeros((1, 1, 2, 2)), + layer( + x=torch.zeros((1, 2, 2)), pairwise_relation_indices=_pairwise_relation_indices([(0, 1, 0, 1)]), ) From 88b5e4254800f8a5bfd0981c3024c05a90fb6d4e Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 5 Aug 2026 20:31:04 +0000 Subject: [PATCH 06/13] Clarify query prefix tensor dimensions --- gigl/nn/graph_transformer.py | 146 ++++++++++++++---------- tests/unit/nn/graph_transformer_test.py | 6 +- 2 files changed, 88 insertions(+), 64 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 2a070b76b..8840daaa8 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -518,11 +518,12 @@ def _forward_query_prefix( to each requested query output. Args: - x: Input tensor of shape ``(batch, seq, model_dim)``. + x: Input tensor of shape ``(batch, sequence_length, model_dim)``. query_length: Number of leading query positions to compute. attn_bias: Optional attention bias broadcastable to - ``(batch, num_heads, query_length, seq)``. - valid_mask: Optional boolean tensor of shape ``(batch, seq)``. + ``(batch, num_heads, query_length, sequence_length)``. + valid_mask: Optional boolean tensor of shape + ``(batch, sequence_length)``. pairwise_relation_indices: Optional sparse relation coordinates shaped ``(num_relation_edges, 4)``. @@ -530,21 +531,22 @@ def _forward_query_prefix( Output tensor of shape ``(batch, query_length, model_dim)``. Raises: - ValueError: If ``query_length`` is outside ``[1, seq]`` or relation - coordinates are invalid. + ValueError: If ``query_length`` is outside + ``[1, sequence_length]`` or relation coordinates are invalid. """ - batch_size, seq_len, model_dim = x.shape - if query_length < 1 or query_length > seq_len: + batch_size, sequence_length, model_dim = x.shape + if query_length < 1 or query_length > sequence_length: raise ValueError( - f"query_length must be in [1, {seq_len}], got {query_length}." + f"query_length must be in [1, {sequence_length}], got {query_length}." ) output_valid_mask = ( valid_mask[:, :query_length] if valid_mask is not None else None - ) - output = x[:, :query_length, :] + ) # [batch, query_length] or None + output = x[:, :query_length, :] # [batch, query_length, model_dim] selected_relation_indices = pairwise_relation_indices + all_relation_edges_filtered = False if ( self._relation_attention_mode != "none" or self._relation_message_mode != "none" @@ -564,7 +566,7 @@ def _forward_query_prefix( selected_relation_indices = pairwise_relation_indices.to( device=x.device, dtype=torch.long, - ) + ) # [num_relation_edges, 4] if selected_relation_indices.numel() > 0: relation_indices = selected_relation_indices[:, 3] if ( @@ -575,40 +577,56 @@ def _forward_query_prefix( "pairwise_relation_indices contains relation ids outside " f"[0, {self._num_relations})." ) - if query_length < seq_len: + if query_length < sequence_length: selected_relation_indices = selected_relation_indices[ selected_relation_indices[:, 1] < query_length - ] + ] # [num_selected_relation_edges, 4] + all_relation_edges_filtered = ( + pairwise_relation_indices.numel() > 0 + and selected_relation_indices.numel() == 0 + ) - x_norm = self._attention_norm(x) - query = self._query_projection(x_norm[:, :query_length, :]) - key = self._key_projection(x_norm) - value = self._value_projection(x_norm) + x_norm = self._attention_norm(x) # [batch, sequence_length, model_dim] + query = self._query_projection( + x_norm[:, :query_length, :] + ) # [batch, query_length, model_dim] + key = self._key_projection(x_norm) # [batch, sequence_length, model_dim] + value = self._value_projection(x_norm) # [batch, sequence_length, model_dim] query = query.view( batch_size, query_length, self._num_heads, self._head_dim - ).transpose(1, 2) - key = key.view(batch_size, seq_len, self._num_heads, self._head_dim).transpose( - 1, 2 - ) + ).transpose(1, 2) # [batch, num_heads, query_length, head_dim] + key = key.view( + batch_size, + sequence_length, + self._num_heads, + self._head_dim, + ).transpose(1, 2) # [batch, num_heads, sequence_length, head_dim] value = value.view( - batch_size, seq_len, self._num_heads, self._head_dim - ).transpose(1, 2) + batch_size, + sequence_length, + self._num_heads, + self._head_dim, + ).transpose(1, 2) # [batch, num_heads, sequence_length, head_dim] if attn_bias is not None and attn_bias.dim() >= 2: - attn_bias = attn_bias[..., :query_length, :] + attn_bias = attn_bias[ + ..., :query_length, : + ] # [..., query_length, sequence_length] attention_output = self._run_attention( query=query, key=key, value=value, attn_bias=attn_bias, pairwise_relation_indices=selected_relation_indices, - ) + ) # [batch, num_heads, query_length, head_dim] attention_output = attention_output.transpose(1, 2).reshape( batch_size, query_length, model_dim - ) - attention_output = self._dropout(self._output_projection(attention_output)) - output = output + attention_output + ) # [batch, query_length, model_dim] + attention_output = self._dropout( + self._output_projection(attention_output) + ) # [batch, query_length, model_dim] + output = output + attention_output # [batch, query_length, model_dim] if self._relation_message_mode == "edge_type_attention": output = output + self._dropout( @@ -621,7 +639,7 @@ def _forward_query_prefix( selected_relation_indices, ), batch_size=batch_size, - seq_len=query_length, + query_length=query_length, ) ) elif self._relation_message_mode != "none": @@ -633,24 +651,26 @@ def _forward_query_prefix( selected_relation_indices, ), batch_size=batch_size, - seq_len=query_length, + query_length=query_length, ) ) if output_valid_mask is not None: - output = output * output_valid_mask.unsqueeze(-1).to(output.dtype) + output = output * output_valid_mask.unsqueeze(-1).to( + output.dtype + ) # [batch, query_length, model_dim] - if ( - self.training - and pairwise_relation_indices is not None - and pairwise_relation_indices.numel() > 0 - and selected_relation_indices is not None - and selected_relation_indices.numel() == 0 - ): + # Match the full layer's autograd participation when prefix filtering + # removes every relation edge from an otherwise nonempty relation set. + if self.training and all_relation_edges_filtered: output = output + self._zero_relation_parameter_dependency(output) - output = output + self._ffn(self._ffn_norm(output)) + output = output + self._ffn( + self._ffn_norm(output) + ) # [batch, query_length, model_dim] if output_valid_mask is not None: - output = output * output_valid_mask.unsqueeze(-1).to(output.dtype) + output = output * output_valid_mask.unsqueeze(-1).to( + output.dtype + ) # [batch, query_length, model_dim] return output @@ -709,9 +729,11 @@ def _build_relation_attention_bias( if pairwise_relation_indices.numel() == 0: return None - batch_size, _, query_len, _ = query.shape - key_len = key.size(2) - empty_bias = query.new_zeros((batch_size, self._num_heads, query_len, key_len)) + batch_size, _, query_length, _ = query.shape + key_length = key.size(2) + empty_bias = query.new_zeros( + (batch_size, self._num_heads, query_length, key_length) + ) return self._add_relation_attention_bias( attn_bias=empty_bias, query=query, @@ -726,8 +748,8 @@ def _add_relation_attention_bias( key: Tensor, pairwise_relation_indices: Tensor, ) -> Optional[Tensor]: - batch_size, _, query_len, _ = query.shape - key_len = key.size(2) + batch_size, _, query_length, _ = query.shape + key_length = key.size(2) if pairwise_relation_indices.numel() == 0: return attn_bias batch_indices = pairwise_relation_indices[:, 0] @@ -735,7 +757,7 @@ def _add_relation_attention_bias( key_indices = pairwise_relation_indices[:, 2] relation_indices = pairwise_relation_indices[:, 3] - bias_shape = (self._num_heads, query_len, key_len) + bias_shape = (self._num_heads, query_length, key_length) if attn_bias is None: attn_bias = query.new_zeros((batch_size, *bias_shape)) elif attn_bias.shape[1:] != bias_shape: @@ -840,7 +862,7 @@ def _compute_relation_messages( x_norm: Tensor, pairwise_relation_indices: Tensor, batch_size: int, - seq_len: int, + query_length: int, ) -> Tensor: """Aggregate per-relation linear messages over sampled directed edges. @@ -854,12 +876,12 @@ def _compute_relation_messages( path and its memory profile are unaffected. Returns: - Message tensor of shape ``(batch_size, seq_len, model_dim)``. + Message tensor of shape ``(batch_size, query_length, model_dim)``. """ if self._relation_message_matrices is None: raise ValueError("Relation message matrices are not initialized.") messages = torch.zeros( - (batch_size, seq_len, x_norm.size(-1)), + (batch_size, query_length, x_norm.size(-1)), dtype=x_norm.dtype, device=x_norm.device, ) @@ -873,7 +895,7 @@ def _compute_relation_messages( # Per-(relation, batch, target) in-degree for mean aggregation, computed # globally up front so grouping order of the entries does not matter. target_degrees = torch.zeros( - (self._num_relations, batch_size, seq_len), + (self._num_relations, batch_size, query_length), dtype=x_norm.dtype, device=x_norm.device, ) @@ -923,7 +945,7 @@ def _compute_relation_attention_messages( key: Tensor, pairwise_relation_indices: Tensor, batch_size: int, - seq_len: int, + query_length: int, ) -> Tensor: """Aggregate per-relation messages weighted by renormalized attention. @@ -938,18 +960,19 @@ def _compute_relation_attention_messages( Args: x_norm: Pre-norm token states of shape - ``(batch_size, seq_len, model_dim)``. + ``(batch_size, key_length, model_dim)``. query: Projected queries of shape - ``(batch_size, num_heads, seq_len, head_dim)``. - key: Projected keys of the same shape as ``query``. + ``(batch_size, num_heads, query_length, head_dim)``. + key: Projected keys of shape + ``(batch_size, num_heads, key_length, head_dim)``. pairwise_relation_indices: Sparse relation coordinates shaped ``(num_relation_edges, 4)`` storing ``(batch_idx, query_pos=target, key_pos=source, relation_idx)``. batch_size: Number of sequences in the batch. - seq_len: Sequence length. + query_length: Number of query outputs. Returns: - Message tensor of shape ``(batch_size, seq_len, model_dim)``. + Message tensor of shape ``(batch_size, query_length, model_dim)``. """ if self._relation_message_matrices is None: raise ValueError("Relation message matrices are not initialized.") @@ -958,7 +981,7 @@ def _compute_relation_attention_messages( if self._relation_message_attention_priors is None: raise ValueError("Relation message attention priors are not initialized.") messages = torch.zeros( - (batch_size, seq_len, x_norm.size(-1)), + (batch_size, query_length, x_norm.size(-1)), dtype=x_norm.dtype, device=x_norm.device, ) @@ -1008,13 +1031,14 @@ def _compute_relation_attention_messages( scores = torch.cat(score_parts, dim=0) # (num_edges, num_heads) # Numerically stable scatter-softmax per (relation, batch, target) group - # and head. Group buffers are (num_relations * batch * seq, num_heads) — - # tiny compared to any (seq, seq) tensor. + # and head. Group buffers are + # (num_relations * batch * query_length, num_heads), much smaller than + # any (sequence_length, sequence_length) tensor. group_indices = ( relation_indices * batch_size + batch_indices - ) * seq_len + target_indices + ) * query_length + target_indices head_group_indices = group_indices.unsqueeze(-1).expand(-1, self._num_heads) - num_groups = self._num_relations * batch_size * seq_len + num_groups = self._num_relations * batch_size * query_length group_score_max = scores.new_full( (num_groups, self._num_heads), torch.finfo(scores.dtype).min ) diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index 3684ae4e4..32959faf4 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -1412,7 +1412,7 @@ def test_relation_message_mean_aggregates_only_indexed_targets(self) -> None: [(0, 1, 0, 0), (0, 1, 2, 0), (0, 2, 0, 1)] ), batch_size=1, - seq_len=3, + query_length=3, ) expected_target_1 = (2.0 * x_norm[0, 0] + 2.0 * x_norm[0, 2]) / 2.0 @@ -1530,7 +1530,7 @@ def test_relation_message_attention_uniform_scores_match_mean_mode(self) -> None [(0, 1, 0, 0), (0, 1, 2, 0), (0, 2, 0, 1)] ), batch_size=1, - seq_len=3, + query_length=3, ) expected_target_1 = (2.0 * x_norm[0, 0] + 2.0 * x_norm[0, 2]) / 2.0 @@ -1595,7 +1595,7 @@ def test_relation_message_attention_concentrates_on_high_score_source( [(0, 1, 0, 0), (0, 1, 2, 0)] ), batch_size=1, - seq_len=3, + query_length=3, ) self.assertTrue(torch.allclose(messages[0, 1], 2.0 * x_norm[0, 0], atol=1e-3)) From 1af0b27c68181d6aff8dd506757d014f105c5a21 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 5 Aug 2026 20:45:24 +0000 Subject: [PATCH 07/13] Preserve transformer length variable names --- gigl/nn/graph_transformer.py | 87 ++++++++++++------------- tests/unit/nn/graph_transformer_test.py | 6 +- 2 files changed, 43 insertions(+), 50 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 8840daaa8..d408befec 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -518,12 +518,11 @@ def _forward_query_prefix( to each requested query output. Args: - x: Input tensor of shape ``(batch, sequence_length, model_dim)``. + x: Input tensor of shape ``(batch, seq, model_dim)``. query_length: Number of leading query positions to compute. attn_bias: Optional attention bias broadcastable to - ``(batch, num_heads, query_length, sequence_length)``. - valid_mask: Optional boolean tensor of shape - ``(batch, sequence_length)``. + ``(batch, num_heads, query_length, seq)``. + valid_mask: Optional boolean tensor of shape ``(batch, seq)``. pairwise_relation_indices: Optional sparse relation coordinates shaped ``(num_relation_edges, 4)``. @@ -531,13 +530,13 @@ def _forward_query_prefix( Output tensor of shape ``(batch, query_length, model_dim)``. Raises: - ValueError: If ``query_length`` is outside - ``[1, sequence_length]`` or relation coordinates are invalid. + ValueError: If ``query_length`` is outside ``[1, seq]`` or relation + coordinates are invalid. """ - batch_size, sequence_length, model_dim = x.shape - if query_length < 1 or query_length > sequence_length: + batch_size, seq_len, model_dim = x.shape + if query_length < 1 or query_length > seq_len: raise ValueError( - f"query_length must be in [1, {sequence_length}], got {query_length}." + f"query_length must be in [1, {seq_len}], got {query_length}." ) output_valid_mask = ( @@ -577,7 +576,7 @@ def _forward_query_prefix( "pairwise_relation_indices contains relation ids outside " f"[0, {self._num_relations})." ) - if query_length < sequence_length: + if query_length < seq_len: selected_relation_indices = selected_relation_indices[ selected_relation_indices[:, 1] < query_length ] # [num_selected_relation_edges, 4] @@ -586,33 +585,31 @@ def _forward_query_prefix( and selected_relation_indices.numel() == 0 ) - x_norm = self._attention_norm(x) # [batch, sequence_length, model_dim] + x_norm = self._attention_norm(x) # [batch, seq_len, model_dim] query = self._query_projection( x_norm[:, :query_length, :] ) # [batch, query_length, model_dim] - key = self._key_projection(x_norm) # [batch, sequence_length, model_dim] - value = self._value_projection(x_norm) # [batch, sequence_length, model_dim] + key = self._key_projection(x_norm) # [batch, seq_len, model_dim] + value = self._value_projection(x_norm) # [batch, seq_len, model_dim] query = query.view( batch_size, query_length, self._num_heads, self._head_dim ).transpose(1, 2) # [batch, num_heads, query_length, head_dim] key = key.view( batch_size, - sequence_length, + seq_len, self._num_heads, self._head_dim, - ).transpose(1, 2) # [batch, num_heads, sequence_length, head_dim] + ).transpose(1, 2) # [batch, num_heads, seq_len, head_dim] value = value.view( batch_size, - sequence_length, + seq_len, self._num_heads, self._head_dim, - ).transpose(1, 2) # [batch, num_heads, sequence_length, head_dim] + ).transpose(1, 2) # [batch, num_heads, seq_len, head_dim] if attn_bias is not None and attn_bias.dim() >= 2: - attn_bias = attn_bias[ - ..., :query_length, : - ] # [..., query_length, sequence_length] + attn_bias = attn_bias[..., :query_length, :] # [..., query_length, seq_len] attention_output = self._run_attention( query=query, key=key, @@ -639,7 +636,7 @@ def _forward_query_prefix( selected_relation_indices, ), batch_size=batch_size, - query_length=query_length, + seq_len=query_length, ) ) elif self._relation_message_mode != "none": @@ -651,7 +648,7 @@ def _forward_query_prefix( selected_relation_indices, ), batch_size=batch_size, - query_length=query_length, + seq_len=query_length, ) ) if output_valid_mask is not None: @@ -729,11 +726,9 @@ def _build_relation_attention_bias( if pairwise_relation_indices.numel() == 0: return None - batch_size, _, query_length, _ = query.shape - key_length = key.size(2) - empty_bias = query.new_zeros( - (batch_size, self._num_heads, query_length, key_length) - ) + batch_size, _, query_len, _ = query.shape + key_len = key.size(2) + empty_bias = query.new_zeros((batch_size, self._num_heads, query_len, key_len)) return self._add_relation_attention_bias( attn_bias=empty_bias, query=query, @@ -748,8 +743,8 @@ def _add_relation_attention_bias( key: Tensor, pairwise_relation_indices: Tensor, ) -> Optional[Tensor]: - batch_size, _, query_length, _ = query.shape - key_length = key.size(2) + batch_size, _, query_len, _ = query.shape + key_len = key.size(2) if pairwise_relation_indices.numel() == 0: return attn_bias batch_indices = pairwise_relation_indices[:, 0] @@ -757,7 +752,7 @@ def _add_relation_attention_bias( key_indices = pairwise_relation_indices[:, 2] relation_indices = pairwise_relation_indices[:, 3] - bias_shape = (self._num_heads, query_length, key_length) + bias_shape = (self._num_heads, query_len, key_len) if attn_bias is None: attn_bias = query.new_zeros((batch_size, *bias_shape)) elif attn_bias.shape[1:] != bias_shape: @@ -862,7 +857,7 @@ def _compute_relation_messages( x_norm: Tensor, pairwise_relation_indices: Tensor, batch_size: int, - query_length: int, + seq_len: int, ) -> Tensor: """Aggregate per-relation linear messages over sampled directed edges. @@ -876,12 +871,12 @@ def _compute_relation_messages( path and its memory profile are unaffected. Returns: - Message tensor of shape ``(batch_size, query_length, model_dim)``. + Message tensor of shape ``(batch_size, seq_len, model_dim)``. """ if self._relation_message_matrices is None: raise ValueError("Relation message matrices are not initialized.") messages = torch.zeros( - (batch_size, query_length, x_norm.size(-1)), + (batch_size, seq_len, x_norm.size(-1)), dtype=x_norm.dtype, device=x_norm.device, ) @@ -895,7 +890,7 @@ def _compute_relation_messages( # Per-(relation, batch, target) in-degree for mean aggregation, computed # globally up front so grouping order of the entries does not matter. target_degrees = torch.zeros( - (self._num_relations, batch_size, query_length), + (self._num_relations, batch_size, seq_len), dtype=x_norm.dtype, device=x_norm.device, ) @@ -945,7 +940,7 @@ def _compute_relation_attention_messages( key: Tensor, pairwise_relation_indices: Tensor, batch_size: int, - query_length: int, + seq_len: int, ) -> Tensor: """Aggregate per-relation messages weighted by renormalized attention. @@ -960,19 +955,18 @@ def _compute_relation_attention_messages( Args: x_norm: Pre-norm token states of shape - ``(batch_size, key_length, model_dim)``. + ``(batch_size, seq_len, model_dim)``. query: Projected queries of shape - ``(batch_size, num_heads, query_length, head_dim)``. - key: Projected keys of shape - ``(batch_size, num_heads, key_length, head_dim)``. + ``(batch_size, num_heads, seq_len, head_dim)``. + key: Projected keys of the same shape as ``query``. pairwise_relation_indices: Sparse relation coordinates shaped ``(num_relation_edges, 4)`` storing ``(batch_idx, query_pos=target, key_pos=source, relation_idx)``. batch_size: Number of sequences in the batch. - query_length: Number of query outputs. + seq_len: Sequence length. Returns: - Message tensor of shape ``(batch_size, query_length, model_dim)``. + Message tensor of shape ``(batch_size, seq_len, model_dim)``. """ if self._relation_message_matrices is None: raise ValueError("Relation message matrices are not initialized.") @@ -981,7 +975,7 @@ def _compute_relation_attention_messages( if self._relation_message_attention_priors is None: raise ValueError("Relation message attention priors are not initialized.") messages = torch.zeros( - (batch_size, query_length, x_norm.size(-1)), + (batch_size, seq_len, x_norm.size(-1)), dtype=x_norm.dtype, device=x_norm.device, ) @@ -1031,14 +1025,13 @@ def _compute_relation_attention_messages( scores = torch.cat(score_parts, dim=0) # (num_edges, num_heads) # Numerically stable scatter-softmax per (relation, batch, target) group - # and head. Group buffers are - # (num_relations * batch * query_length, num_heads), much smaller than - # any (sequence_length, sequence_length) tensor. + # and head. Group buffers are (num_relations * batch * seq, num_heads) — + # tiny compared to any (seq, seq) tensor. group_indices = ( relation_indices * batch_size + batch_indices - ) * query_length + target_indices + ) * seq_len + target_indices head_group_indices = group_indices.unsqueeze(-1).expand(-1, self._num_heads) - num_groups = self._num_relations * batch_size * query_length + num_groups = self._num_relations * batch_size * seq_len group_score_max = scores.new_full( (num_groups, self._num_heads), torch.finfo(scores.dtype).min ) diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index 32959faf4..3684ae4e4 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -1412,7 +1412,7 @@ def test_relation_message_mean_aggregates_only_indexed_targets(self) -> None: [(0, 1, 0, 0), (0, 1, 2, 0), (0, 2, 0, 1)] ), batch_size=1, - query_length=3, + seq_len=3, ) expected_target_1 = (2.0 * x_norm[0, 0] + 2.0 * x_norm[0, 2]) / 2.0 @@ -1530,7 +1530,7 @@ def test_relation_message_attention_uniform_scores_match_mean_mode(self) -> None [(0, 1, 0, 0), (0, 1, 2, 0), (0, 2, 0, 1)] ), batch_size=1, - query_length=3, + seq_len=3, ) expected_target_1 = (2.0 * x_norm[0, 0] + 2.0 * x_norm[0, 2]) / 2.0 @@ -1595,7 +1595,7 @@ def test_relation_message_attention_concentrates_on_high_score_source( [(0, 1, 0, 0), (0, 1, 2, 0)] ), batch_size=1, - query_length=3, + seq_len=3, ) self.assertTrue(torch.allclose(messages[0, 1], 2.0 * x_norm[0, 0], atol=1e-3)) From f6275a9d93843b2919ff84e84f8eca558aad74b0 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 5 Aug 2026 22:13:46 +0000 Subject: [PATCH 08/13] Restore transformer state naming --- gigl/nn/graph_transformer.py | 97 +++++++++++++++++------------------- 1 file changed, 45 insertions(+), 52 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index d408befec..e330d7a32 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -465,7 +465,7 @@ def forward( """ return self._forward_query_prefix( x=x, - query_length=x.size(1), + seq_len=x.size(1), attn_bias=attn_bias, valid_mask=valid_mask, pairwise_relation_indices=pairwise_relation_indices, @@ -498,7 +498,7 @@ def _forward_anchor_only( """ return self._forward_query_prefix( x=x, - query_length=1, + seq_len=1, attn_bias=attn_bias, valid_mask=valid_mask, pairwise_relation_indices=pairwise_relation_indices, @@ -507,7 +507,7 @@ def _forward_anchor_only( def _forward_query_prefix( self, x: Tensor, - query_length: int, + seq_len: int, attn_bias: Optional[Tensor], valid_mask: Optional[Tensor], pairwise_relation_indices: Optional[Tensor], @@ -519,30 +519,28 @@ def _forward_query_prefix( Args: x: Input tensor of shape ``(batch, seq, model_dim)``. - query_length: Number of leading query positions to compute. + seq_len: Number of leading sequence positions to compute. attn_bias: Optional attention bias broadcastable to - ``(batch, num_heads, query_length, seq)``. + ``(batch, num_heads, seq_len, input_seq_len)``. valid_mask: Optional boolean tensor of shape ``(batch, seq)``. pairwise_relation_indices: Optional sparse relation coordinates shaped ``(num_relation_edges, 4)``. Returns: - Output tensor of shape ``(batch, query_length, model_dim)``. + Output tensor of shape ``(batch, seq_len, model_dim)``. Raises: - ValueError: If ``query_length`` is outside ``[1, seq]`` or relation + ValueError: If ``seq_len`` is outside ``[1, input_seq_len]`` or relation coordinates are invalid. """ - batch_size, seq_len, model_dim = x.shape - if query_length < 1 or query_length > seq_len: - raise ValueError( - f"query_length must be in [1, {seq_len}], got {query_length}." - ) + batch_size, input_seq_len, model_dim = x.shape + if seq_len < 1 or seq_len > input_seq_len: + raise ValueError(f"seq_len must be in [1, {input_seq_len}], got {seq_len}.") - output_valid_mask = ( - valid_mask[:, :query_length] if valid_mask is not None else None - ) # [batch, query_length] or None - output = x[:, :query_length, :] # [batch, query_length, model_dim] + valid_mask = ( + valid_mask[:, :seq_len] if valid_mask is not None else None + ) # [batch, seq_len] or None + residual = x[:, :seq_len, :] # [batch, seq_len, model_dim] selected_relation_indices = pairwise_relation_indices all_relation_edges_filtered = False @@ -576,57 +574,57 @@ def _forward_query_prefix( "pairwise_relation_indices contains relation ids outside " f"[0, {self._num_relations})." ) - if query_length < seq_len: + if seq_len < input_seq_len: selected_relation_indices = selected_relation_indices[ - selected_relation_indices[:, 1] < query_length + selected_relation_indices[:, 1] < seq_len ] # [num_selected_relation_edges, 4] all_relation_edges_filtered = ( pairwise_relation_indices.numel() > 0 and selected_relation_indices.numel() == 0 ) - x_norm = self._attention_norm(x) # [batch, seq_len, model_dim] + x_norm = self._attention_norm(x) # [batch, input_seq_len, model_dim] query = self._query_projection( - x_norm[:, :query_length, :] - ) # [batch, query_length, model_dim] - key = self._key_projection(x_norm) # [batch, seq_len, model_dim] - value = self._value_projection(x_norm) # [batch, seq_len, model_dim] + x_norm[:, :seq_len, :] + ) # [batch, seq_len, model_dim] + key = self._key_projection(x_norm) # [batch, input_seq_len, model_dim] + value = self._value_projection(x_norm) # [batch, input_seq_len, model_dim] query = query.view( - batch_size, query_length, self._num_heads, self._head_dim - ).transpose(1, 2) # [batch, num_heads, query_length, head_dim] + batch_size, seq_len, self._num_heads, self._head_dim + ).transpose(1, 2) # [batch, num_heads, seq_len, head_dim] key = key.view( batch_size, - seq_len, + input_seq_len, self._num_heads, self._head_dim, - ).transpose(1, 2) # [batch, num_heads, seq_len, head_dim] + ).transpose(1, 2) # [batch, num_heads, input_seq_len, head_dim] value = value.view( batch_size, - seq_len, + input_seq_len, self._num_heads, self._head_dim, - ).transpose(1, 2) # [batch, num_heads, seq_len, head_dim] + ).transpose(1, 2) # [batch, num_heads, input_seq_len, head_dim] if attn_bias is not None and attn_bias.dim() >= 2: - attn_bias = attn_bias[..., :query_length, :] # [..., query_length, seq_len] + attn_bias = attn_bias[..., :seq_len, :] # [..., seq_len, input_seq_len] attention_output = self._run_attention( query=query, key=key, value=value, attn_bias=attn_bias, pairwise_relation_indices=selected_relation_indices, - ) # [batch, num_heads, query_length, head_dim] + ) # [batch, num_heads, seq_len, head_dim] attention_output = attention_output.transpose(1, 2).reshape( - batch_size, query_length, model_dim - ) # [batch, query_length, model_dim] + batch_size, seq_len, model_dim + ) # [batch, seq_len, model_dim] attention_output = self._dropout( self._output_projection(attention_output) - ) # [batch, query_length, model_dim] - output = output + attention_output # [batch, query_length, model_dim] + ) # [batch, seq_len, model_dim] + x = residual + attention_output # [batch, seq_len, model_dim] if self._relation_message_mode == "edge_type_attention": - output = output + self._dropout( + x = x + self._dropout( self._compute_relation_attention_messages( x_norm=x_norm, query=query, @@ -636,11 +634,11 @@ def _forward_query_prefix( selected_relation_indices, ), batch_size=batch_size, - seq_len=query_length, + seq_len=seq_len, ) ) elif self._relation_message_mode != "none": - output = output + self._dropout( + x = x + self._dropout( self._compute_relation_messages( x_norm=x_norm, pairwise_relation_indices=cast( @@ -648,28 +646,23 @@ def _forward_query_prefix( selected_relation_indices, ), batch_size=batch_size, - seq_len=query_length, + seq_len=seq_len, ) ) - if output_valid_mask is not None: - output = output * output_valid_mask.unsqueeze(-1).to( - output.dtype - ) # [batch, query_length, model_dim] + if valid_mask is not None: + x = x * valid_mask.unsqueeze(-1).to(x.dtype) # [batch, seq_len, model_dim] # Match the full layer's autograd participation when prefix filtering # removes every relation edge from an otherwise nonempty relation set. if self.training and all_relation_edges_filtered: - output = output + self._zero_relation_parameter_dependency(output) + x = x + self._zero_relation_parameter_dependency(x) - output = output + self._ffn( - self._ffn_norm(output) - ) # [batch, query_length, model_dim] - if output_valid_mask is not None: - output = output * output_valid_mask.unsqueeze(-1).to( - output.dtype - ) # [batch, query_length, model_dim] + residual = x + x = residual + self._ffn(self._ffn_norm(x)) # [batch, seq_len, model_dim] + if valid_mask is not None: + x = x * valid_mask.unsqueeze(-1).to(x.dtype) # [batch, seq_len, model_dim] - return output + return x def _zero_relation_parameter_dependency(self, reference: Tensor) -> Tensor: """Keep relation parameters in the training graph after anchor filtering.""" From 1725c227307b58dc4a1214142a5a557547a4a399 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 5 Aug 2026 22:15:59 +0000 Subject: [PATCH 09/13] Explain zero relation gradient dependency --- gigl/nn/graph_transformer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index e330d7a32..dd3e5a733 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -652,8 +652,10 @@ def _forward_query_prefix( if valid_mask is not None: x = x * valid_mask.unsqueeze(-1).to(x.dtype) # [batch, seq_len, model_dim] - # Match the full layer's autograd participation when prefix filtering - # removes every relation edge from an otherwise nonempty relation set. + # The full-sequence path still touches relation parameters for edges outside + # this prefix, producing zero gradients. Filtering all such edges would leave + # those parameters with grad=None instead, changing DDP/optimizer behavior; + # retain a zero-valued dependency to preserve the original training semantics. if self.training and all_relation_edges_filtered: x = x + self._zero_relation_parameter_dependency(x) From 96583e7efe5718cec76758dd3632b2c5212f50f9 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Mon, 10 Aug 2026 16:45:05 +0000 Subject: [PATCH 10/13] Explain anchor-only final layer selection --- gigl/nn/graph_transformer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index dd3e5a733..22f64af01 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -1949,6 +1949,10 @@ def _encode_and_readout( valid_mask=valid_mask, ) + # ``anchor_only`` readout returns only position zero. In the final layer, + # non-anchor outputs cannot affect that result, but the anchor still + # attends over the full sequence as keys and values. Compute only the + # anchor query/output and shorten the mask to match that one-token result. output_valid_mask = valid_mask if final_encoder_layer is not None: x = final_encoder_layer._forward_anchor_only( From 3960125822f4611cb5d5305ea0d3527d627d204e Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Mon, 10 Aug 2026 22:59:48 +0000 Subject: [PATCH 11/13] Route anchor-only layer through forward --- gigl/nn/graph_transformer.py | 48 ++++++++---------------------------- 1 file changed, 10 insertions(+), 38 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 22f64af01..cf50af761 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -446,6 +446,7 @@ def forward( attn_bias: Optional[Tensor] = None, valid_mask: Optional[Tensor] = None, pairwise_relation_indices: Optional[Tensor] = None, + query_seq_len: Optional[int] = None, ) -> Tensor: """Forward pass. @@ -459,46 +460,15 @@ def forward( pairwise_relation_indices: Optional long tensor of shape ``(num_relation_edges, 4)`` with sparse ``(batch_idx, query_pos, key_pos, relation_idx)`` coordinates. + query_seq_len: Number of leading query positions to compute. Defaults + to the full input sequence length. Returns: - Output tensor of shape ``(batch, seq, model_dim)``. - """ - return self._forward_query_prefix( - x=x, - seq_len=x.size(1), - attn_bias=attn_bias, - valid_mask=valid_mask, - pairwise_relation_indices=pairwise_relation_indices, - ) - - def _forward_anchor_only( - self, - x: Tensor, - attn_bias: Optional[Tensor] = None, - valid_mask: Optional[Tensor] = None, - pairwise_relation_indices: Optional[Tensor] = None, - ) -> Tensor: - """Compute the final layer output for the anchor token only. - - Keys and values still cover the complete sequence because every token - can contribute to the anchor. Query-side attention, relation messages, - output projection, and feed-forward work are restricted to position - zero because later token outputs cannot affect the anchor. - - Args: - x: Input tensor of shape ``(batch, seq, model_dim)``. - attn_bias: Optional attention bias broadcastable to - ``(batch, num_heads, seq, seq)``. - valid_mask: Optional boolean tensor of shape ``(batch, seq)``. - pairwise_relation_indices: Optional sparse relation coordinates - shaped ``(num_relation_edges, 4)``. - - Returns: - Anchor output of shape ``(batch, 1, model_dim)``. + Output tensor of shape ``(batch, query_seq_len, model_dim)``. """ return self._forward_query_prefix( x=x, - seq_len=1, + seq_len=x.size(1) if query_seq_len is None else query_seq_len, attn_bias=attn_bias, valid_mask=valid_mask, pairwise_relation_indices=pairwise_relation_indices, @@ -1935,13 +1905,14 @@ def _encode_and_readout( if use_anchor_only_final_layer else None ) + if final_encoder_layer is not None: + encoder_layers = encoder_layers[:-1] + for encoder_layer_module in encoder_layers: encoder_layer = cast( GraphTransformerEncoderLayer, encoder_layer_module, ) - if encoder_layer is final_encoder_layer: - break x = encoder_layer( x, attn_bias=attn_bias, @@ -1955,11 +1926,12 @@ def _encode_and_readout( # anchor query/output and shorten the mask to match that one-token result. output_valid_mask = valid_mask if final_encoder_layer is not None: - x = final_encoder_layer._forward_anchor_only( + x = final_encoder_layer( x, attn_bias=attn_bias, pairwise_relation_indices=pairwise_relation_indices, valid_mask=valid_mask, + query_seq_len=1, ) output_valid_mask = valid_mask[:, :1] From d91ec9728177b01e4291002ba1c9f34995dae684 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Mon, 10 Aug 2026 22:59:48 +0000 Subject: [PATCH 12/13] Route anchor-only layer through forward --- gigl/nn/graph_transformer.py | 52 +++++++++--------------------------- 1 file changed, 12 insertions(+), 40 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 22f64af01..a43c3f5e1 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -446,6 +446,7 @@ def forward( attn_bias: Optional[Tensor] = None, valid_mask: Optional[Tensor] = None, pairwise_relation_indices: Optional[Tensor] = None, + query_seq_len: Optional[int] = None, ) -> Tensor: """Forward pass. @@ -459,46 +460,15 @@ def forward( pairwise_relation_indices: Optional long tensor of shape ``(num_relation_edges, 4)`` with sparse ``(batch_idx, query_pos, key_pos, relation_idx)`` coordinates. + query_seq_len: Number of leading query positions to compute. Defaults + to the full input sequence length. Returns: - Output tensor of shape ``(batch, seq, model_dim)``. + Output tensor of shape ``(batch, query_seq_len, model_dim)``. """ return self._forward_query_prefix( x=x, - seq_len=x.size(1), - attn_bias=attn_bias, - valid_mask=valid_mask, - pairwise_relation_indices=pairwise_relation_indices, - ) - - def _forward_anchor_only( - self, - x: Tensor, - attn_bias: Optional[Tensor] = None, - valid_mask: Optional[Tensor] = None, - pairwise_relation_indices: Optional[Tensor] = None, - ) -> Tensor: - """Compute the final layer output for the anchor token only. - - Keys and values still cover the complete sequence because every token - can contribute to the anchor. Query-side attention, relation messages, - output projection, and feed-forward work are restricted to position - zero because later token outputs cannot affect the anchor. - - Args: - x: Input tensor of shape ``(batch, seq, model_dim)``. - attn_bias: Optional attention bias broadcastable to - ``(batch, num_heads, seq, seq)``. - valid_mask: Optional boolean tensor of shape ``(batch, seq)``. - pairwise_relation_indices: Optional sparse relation coordinates - shaped ``(num_relation_edges, 4)``. - - Returns: - Anchor output of shape ``(batch, 1, model_dim)``. - """ - return self._forward_query_prefix( - x=x, - seq_len=1, + seq_len=x.size(1) if query_seq_len is None else query_seq_len, attn_bias=attn_bias, valid_mask=valid_mask, pairwise_relation_indices=pairwise_relation_indices, @@ -1935,13 +1905,14 @@ def _encode_and_readout( if use_anchor_only_final_layer else None ) - for encoder_layer_module in encoder_layers: + normal_encoder_layer_count = len(encoder_layers) - int( + final_encoder_layer is not None + ) + for encoder_layer_index in range(normal_encoder_layer_count): encoder_layer = cast( GraphTransformerEncoderLayer, - encoder_layer_module, + encoder_layers[encoder_layer_index], ) - if encoder_layer is final_encoder_layer: - break x = encoder_layer( x, attn_bias=attn_bias, @@ -1955,11 +1926,12 @@ def _encode_and_readout( # anchor query/output and shorten the mask to match that one-token result. output_valid_mask = valid_mask if final_encoder_layer is not None: - x = final_encoder_layer._forward_anchor_only( + x = final_encoder_layer( x, attn_bias=attn_bias, pairwise_relation_indices=pairwise_relation_indices, valid_mask=valid_mask, + query_seq_len=1, ) output_valid_mask = valid_mask[:, :1] From 8efe38157982ea737ca11252adccd02d59f0123d Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 12 Aug 2026 00:12:31 +0000 Subject: [PATCH 13/13] Fix encoder layer iteration typing --- gigl/nn/graph_transformer.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index cf50af761..a43c3f5e1 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -1905,13 +1905,13 @@ def _encode_and_readout( if use_anchor_only_final_layer else None ) - if final_encoder_layer is not None: - encoder_layers = encoder_layers[:-1] - - for encoder_layer_module in encoder_layers: + normal_encoder_layer_count = len(encoder_layers) - int( + final_encoder_layer is not None + ) + for encoder_layer_index in range(normal_encoder_layer_count): encoder_layer = cast( GraphTransformerEncoderLayer, - encoder_layer_module, + encoder_layers[encoder_layer_index], ) x = encoder_layer( x,