HeteroData.subgraph() crashes for edge attributes with non-zero concatenation dimensions
Describe the bug
HeteroData.subgraph() raises an IndexError when an edge attribute is concatenated along a dimension other than 0.
For example, an additional edge attribute whose name contains index is assigned a concatenation dimension of -1 by HeteroData.__cat_dim__. EdgeStorage.is_edge_attr() therefore correctly recognizes a tensor shaped [2, num_edges] as an edge attribute. However, HeteroData.subgraph() applies the one-dimensional edge mask as value[edge_mask], which always indexes dimension 0.
This prevents node-induced subgraph extraction for heterogeneous graphs containing such edge attributes.
Minimal reproduction
import torch
from torch_geometric.data import HeteroData
data = HeteroData()
data['paper'].num_nodes = 3
store = data['paper', 'cites', 'paper']
store.edge_index = torch.tensor([
[0, 1, 2],
[1, 2, 0],
])
store.pair_index = torch.tensor([
[10, 11, 12],
[20, 21, 22],
])
print('is_edge_attr:', store.is_edge_attr('pair_index'))
print('cat_dim:', data.__cat_dim__('pair_index', store.pair_index, store))
out = data.subgraph({'paper': torch.tensor([0, 1])})
print(out['paper', 'cites', 'paper'].pair_index)Observed on current master (cc678a392255a1467872f54582724b8dce434603):
is_edge_attr: True
cat_dim: -1
Traceback (most recent call last):
File "<string>", line 1, in <module>
File "torch_geometric/data/hetero_data.py", line 809, in subgraph
data[edge_type][key] = value[edge_mask]
~~~~~^^^^^^^^^^^
IndexError: The shape of the mask [3] at index 0 does not match the shape of the indexed tensor [2, 3] at index 0Expected: the only retained edge is 0 -> 1, so pair_index should be selected along its concatenation dimension and equal:
tensor([[10],
[20]])The equivalent edge-attribute paths in Data.subgraph() and HeteroData.edge_subgraph() already select along the attribute's configured concatenation dimension.
Versions
The reproducer uses CPU tensor operations only.
PyG version: 2.9.0 from source
PyG commit: cc678a392255a1467872f54582724b8dce434603
PyTorch version: 2.11.0+cu130
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A
OS: Microsoft Windows 11 Home Single Language (10.0.26200 64-bit)
GCC version: (Rev8, Built by MSYS2 project) 15.2.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A
Python version: 3.13.3 (tags/v3.13.3:6280bb5, Apr 8 2025, 14:47:33) [MSC v.1943 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-11-10.0.26200-SP0
Is CUDA available: True
CUDA runtime version: 12.9.86
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 5080 Laptop GPU
Nvidia driver version: Could not collect
cuDNN version: Could not collect
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
CPU: Intel(R) Core(TM) Ultra 9 275HX
Relevant packages:
numpy==2.2.6
torch==2.11.0+cu130
torchaudio==2.11.0
torchvision==0.26.0+cu130Disclosure
This report and its minimal reproducer were prepared with assistance from OpenAI Codex. I independently verified the behavior locally against the current master commit shown above.
Source: pyg-team/pytorch_geometric