Global operations on batches with bipartite data + num_nodes

[thistlillo]

I am working with bipartite graphs, with training based on mini-batches of examples.

This is a MWE that highlights the problem I am encountering in my 'real' code:

import random
import torch
import torch.nn as nn
from torch_geometric.data import Data, Batch
from torch_geometric.loader import DataLoader
import torch_geometric as pyg

additional_tensor_keys = ["fid", "l", "h"]
class BipartiteData(Data):
    def __inc__(self, key, value, *args, **kwargs):
        if key == 'edge_index':
            inc = torch.tensor([[self.x_left.size(0)], [self.x_right.size(0)]])
        else:
            inc = super().__inc__(key, value, *args, **kwargs)
        # print(f"__inc__: {key} -> {inc}")
        return inc
    #<

    def __cat_dim__(self, key, value, *args, **kwargs):
        if key in additional_tensor_keys:
            return None
        return super().__cat_dim__(key, value, *args, **kwargs)

    @property
    def num_nodes(self):
        return self.x_left.size(0) + self.x_right.size(0)
    #<
#<


def generate_bipartite_example(n_left, n_right, fid, l, h):
    edges = []
    edge_features = []
    for l in range(n_left):
        num_connections = random.randint(1, n_right)
        connected_right = random.sample(range(n_right), num_connections)
        for r in connected_right:
            edges.append([l, r])  # constraint → variable
            edge_features.append(1)
    edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()  # [2, num_edges]
    edge_features = torch.tensor(edge_features, dtype=torch.float).view(-1, 1)  # [num_edges, 1]

    # Feature tensors
    constraint_x = torch.randn(n_left, 32)
    variable_x = torch.randn(n_right, 32)
    bd = BipartiteData(
        x_left=constraint_x,
        x_right=variable_x,
        edge_index=edge_index,
        x_edge=edge_features,
    )
    bd.n_left = n_left
    bd.n_right = n_right
    bd.fid = torch.tensor(fid, dtype=torch.long)
    bd.l = torch.tensor(l, dtype=torch.float)
    bd.h = torch.tensor(h, dtype=torch.float)
    return bd
#<

def generate_graphs(n_graphs, max_left, max_right):
    graphs = [
        generate_bipartite_example(
            n_left=random.randint(3, max_left),
            n_right=random.randint(3, max_right),
            fid=i, l=1.0, h=1.0
        ) for i in range(n_graphs)
    ]
    return graphs
#<

n_graphs = 3
max_right = 5
max_left = 8
graphs = generate_graphs(n_graphs, max_left, max_right)

for g in graphs:
    print(g)

dataloader = DataLoader(graphs, batch_size=2)
iter_loader = iter(dataloader)
batch = next(iter_loader)
print(batch)

# build a model that received a bipartite graph, process left and right features, then perform a global graph-level op

class Example(nn.Module):
    def __init__(self, in_channels_left, in_channels_right, out_channels):
        super(Example, self).__init__()
        self.left_p = nn.Linear(in_channels_left, out_channels)
        self.right_p = nn.Linear(in_channels_right, out_channels)
        self.pool = pyg.nn.global_add_pool
    
    def forward(self, batch):
        x_left = self.left_p(batch.x_left)
        x_right = self.right_p(batch.x_right)
        
        out1 = self.pool(x_left, batch.batch)
        out2 = self.pool(x_right, batch.batch)
        return out
    
model = Example(32, 32, 16)
out = model(batch)
print("first output:", out.shape)


batch = next(iter_loader)
out = model(batch)
print("second output:", out.shape)

whose output is:

BipartiteData(edge_index=[2, 22], x_left=[7, 32], x_right=[5, 32], x_edge=[22, 1], n_left=7, n_right=5, fid=0, l=6.0, h=1.0)
BipartiteData(edge_index=[2, 13], x_left=[6, 32], x_right=[3, 32], x_edge=[13, 1], n_left=6, n_right=3, fid=1, l=5.0, h=1.0)
BipartiteData(edge_index=[2, 17], x_left=[8, 32], x_right=[4, 32], x_edge=[17, 1], n_left=8, n_right=4, fid=2, l=7.0, h=1.0)
BipartiteDataBatch(edge_index=[2, 35], x_left=[13, 32], x_right=[8, 32], x_edge=[35, 1], n_left=[2], n_right=[2], fid=[2], l=[2], h=[2], batch=[13], ptr=[3])
[/home/USER/workspace/.../.venv/lib/python3.10/site-packages/torch_geometric/data/storage.py:452): UserWarning: Unable to accurately infer 'num_nodes' from the attribute set '{'l', 'edge_index', 'fid', 'n_left', 'x_left', 'h', 'x_edge', 'n_right', 'x_right'}'. Please explicitly set 'num_nodes' as an attribute of 'data' to suppress this warning
  warnings.warn(
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[1], line 102
     99         return out
    101 model = Example(32, 32, 16)
--> 102 out = model(batch)
    103 print("first output:", out.shape)
    106 batch = next(iter_loader)

File ~/workspace/project/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module._wrapped_call_impl(self, *args, **kwargs)
   1734     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735 else:
-> 1736     return self._call_impl(*args, **kwargs)

File ~/workspace/project/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module._call_impl(self, *args, **kwargs)
   1742 # If we don't have any hooks, we want to skip the rest of the logic in
   1743 # this function, and just call forward.
   1744 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1745         or _global_backward_pre_hooks or _global_backward_hooks
   1746         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747     return forward_call(*args, **kwargs)
   1749 result = None
   1750 called_always_called_hooks = set()

Cell In[1], line 98, in Example.forward(self, batch)
     95 x_right = self.right_p(batch.x_right)
     97 out1 = self.pool(x_left, batch.batch)
---> 98 out2 = self.pool(x_right, batch.batch)
     99 return out

File ~/workspace/project/.venv/lib/python3.10/site-packages/torch_geometric/nn/pool/glob.py:34, in global_add_pool(x, batch, size)
     32 if batch is None:
     33     return x.sum(dim=dim, keepdim=x.dim() <= 2)
---> 34 return scatter(x, batch, dim=dim, dim_size=size, reduce='sum')

File ~/workspace/project/.venv/lib/python3.10/site-packages/torch_geometric/utils/_scatter.py:74, in scatter(src, index, dim, dim_size, reduce)
     72 # For "sum" and "mean" reduction, we make use of `scatter_add_`:
     73 if reduce == 'sum' or reduce == 'add':
---> 74     index = broadcast(index, src, dim)
     75     return src.new_zeros(size).scatter_add_(dim, index, src)
     77 if reduce == 'mean':

File ~/workspace/project/.venv/lib/python3.10/site-packages/torch_geometric/utils/_scatter.py:195, in broadcast(src, ref, dim)
    193 dim = ref.dim() + dim if dim < 0 else dim
    194 size = ((1, ) * dim) + (-1, ) + ((1, ) * (ref.dim() - dim - 1))
--> 195 return src.view(size).expand_as(ref)

RuntimeError: The expanded size of the tensor (8) must match the existing size (13) at non-singleton dimension 0.  Target sizes: [8, 16].  Tensor sizes: [13, 1]

I do not understand two things:

1. Why do I get "UserWarning: Unable to accurately infer 'num_nodes' from the attribute set '{'l', 'x_right', 'fid', 'h', 'x_edge', 'n_right', 'n_left', 'edge_index', 'x_left'}'. Please explicitly set 'num_nodes' as an attribute of 'data' to suppress this warning
   warnings.warn(" ? Shouldn't the property @property def num_nodes(self): ... avoid that warning?

2. This error message: The expanded size of the tensor (8) must match the existing size (13) at non-singleton dimension 0. Target sizes: [8, 16]. Tensor sizes: [13, 1]

How am I supposed to perform global, graph-level operations on the two sets of nodes in the bipartite graph?

Using:

torch==2.5.0
torch-cluster==1.6.3+pt25cu124
torch-geometric==2.6.1
torch-scatter==2.1.2+pt25cu124
torch-sparse==0.6.18+pt25cu124
torch-spline-conv==1.2.2+pt25cu124
torchaudio==2.5.0
torchvision==0.20.0

with Python 3.10.15
on Linux.