torchcell.nn.masked_gin_conv.MaskedGINConv

class MaskedGINConv(nn: Callable, eps: float = 0.0, train_eps: bool = False, **kwargs)[source]

Bases: GINConv

Graph Isomorphism Network with masked message passing.

Key Innovation: Masks messages, NOT edges. This avoids expensive edge filtering (edge_index[:, mask]) that would negate the speedup from LazySubgraphRepresentation.

How It Works: 1. Process ALL edges (no filtering of edge_index) 2. Zero out messages from masked edges (element-wise multiply) 3. Aggregate using PyG’s optimized scatter operations 4. No tensor copying or allocation

This approach is fast because: - No edge_index[:, mask] filtering (expensive) - Element-wise mask multiplication is GPU-accelerated - Same edge_index reused every forward pass (cache-friendly) - Memory-efficient (no new tensor allocation)

Parameters:

  • nn (torch.nn.Module) – MLP that transforms aggregated features

  • eps (float, optional) – Initial epsilon value for self-loop weighting

  • train_eps (bool, optional) – Whether epsilon should be trainable

  • **kwargs – Additional arguments for MessagePassing

Example

>>> import torch.nn as nn
>>> from torch_geometric.data import Data
>>>
>>> # Create MLP for GIN
>>> mlp = nn.Sequential(
...     nn.Linear(64, 64),
...     nn.ReLU(),
...     nn.Linear(64, 64)
... )
>>>
>>> # Create masked GIN layer
>>> conv = MaskedGINConv(mlp, train_eps=True)
>>>
>>> # Example data
>>> x = torch.randn(100, 64)  # 100 nodes, 64 features
>>> edge_index = torch.randint(0, 100, (2, 500))  # 500 edges
>>> edge_mask = torch.rand(500) > 0.3  # Mask 30% of edges
>>>
>>> # Forward with masking - edge_index is NOT copied!
>>> out = conv(x, edge_index, edge_mask=edge_mask)
>>> out.shape
torch.Size([100, 64])
forward(x: Tensor | Tuple[Tensor, Tensor | None], edge_index: Tensor | SparseTensor, edge_mask: Tensor | None = None, size: Tuple[int, int] | None = None) → Tensor[source]

Forward pass with masked message passing.

IMPORTANT: This does NOT filter edge_index. Instead, it masks messages in the message() method, preserving lazy speedup.

Parameters:

  • x – Node features [num_nodes, in_channels] or pair

  • edge_index – Edge connectivity [2, num_edges] - NOT filtered!

  • edge_mask – Optional boolean mask [num_edges] - True=keep, False=zero

  • size – Size of source/destination nodes (bipartite graphs)

Returns:

Updated node features [num_nodes, out_channels]

Return type:

out

message(x_j: Tensor, edge_mask: Tensor | None = None) → Tensor[source]

Compute messages with optional masking.

This is where masking happens - messages are zeroed out for masked edges, but edge_index itself is never modified.

Parameters:

  • x_j – Source node features for each edge [num_edges, F]

  • edge_mask – Optional mask [num_edges] - True=keep, False=zero

Returns:

Masked messages [num_edges, F]