HypergraphConv not Compatible with Torchscript Compilation

[chrisrosner]

🐛 Describe the bug

Trying to run torch.jit.script() on a HypergraphConv produces exceptions.

With use_attention=True

import torch
from torch_geometric.nn import HypergraphConv

in_channels, out_channels = (16, 32)
conv = HypergraphConv(in_channels, out_channels, use_attention=True)
script = torch.jit.script(conv)
assert script


RuntimeError: 
Variable 'alpha' previously had type NoneType but is now being assigned to a value of type Tensor
:
  File "/Users/chris/pytorch_geometric/torch_geometric/experimental.py", line 174
            x_i = x[hyperedge_index[0]]
            x_j = hyperedge_attr[hyperedge_index[1]]
            alpha = (torch.cat([x_i, x_j], dim=-1) * self.att).sum(dim=-1)
            ~~~~~ <--- HERE
            alpha = F.leaky_relu(alpha, self.negative_slope)
            if self.attention_mode == 'node':

With use_attention=False

import torch
from torch_geometric.nn import HypergraphConv

in_channels, out_channels = (16, 32)
conv = HypergraphConv(in_channels, out_channels, use_attention=True)
script = torch.jit.script(conv)
assert script

RuntimeError: 
Module 'HypergraphConv' has no attribute 'att' :
  File "/Users/chris/pytorch_geometric/torch_geometric/experimental.py", line 174
            x_i = x[hyperedge_index[0]]
            x_j = hyperedge_attr[hyperedge_index[1]]
            alpha = (torch.cat([x_i, x_j], dim=-1) * self.att).sum(dim=-1)
                                                     ~~~~~~~~ <--- HERE
            alpha = F.leaky_relu(alpha, self.negative_slope)
            if self.attention_mode == 'node':

Versions

Tried with:
pytorch_geometric 2.6.1
pytorch 2.2.2, and 2.8.0

Collecting environment information...
PyTorch version: 2.8.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.4.1 (arm64)
GCC version: Could not collect
Clang version: 17.0.0 (clang-1700.0.13.3)
CMake version: Could not collect
Libc version: N/A

Python version: 3.12.4 (main, Jan 13 2025, 14:59:08) [Clang 16.0.0 (clang-1600.0.26.6)] (64-bit runtime)
Python platform: macOS-15.4.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M4

Versions of relevant libraries:
[pip3] numpy==2.3.2
[pip3] torch==2.8.0
[pip3] torch-geometric==2.7.0
[conda] Could not collect

[chrisrosner]

Some tests that can be added to test_hypergraph_conv.py which demonstrates the issue(s):

def test_hypergraph_jit():
    in_channels, out_channels = (2, 3)
    conv = HypergraphConv(in_channels, out_channels)
    script = torch.jit.script(conv)
    output = script(torch.randn(4, in_channels),
                    torch.tensor([[0, 1, 2], [0, 0, 1]]),
                    hyperedge_attr=torch.randn(2, in_channels))
    assert output.size() == (4, out_channels)

def test_hypergraph_jit_with_attention():
    in_channels, out_channels = (2, 3)
    conv = HypergraphConv(in_channels, out_channels, use_attention=True)
    script = torch.jit.script(conv)
    output = script(torch.randn(4, in_channels),
                    torch.tensor([[0, 1, 2], [0, 0, 1]]),
                    hyperedge_attr=torch.randn(2, in_channels))
    assert output.size() == (4, out_channels * conv.heads)