tvm's interaction with pytorch_geometric

I’m trying to compile a graph neural network model written with PyTorch and an extension called torch_geometric, but it seems that tvm has limited support for external libraries it uses such as torch-scatter, torch-sparse, torch-cluster and torch-spline-conv. I’m very new to tvm so I’m not 100% sure if I’m using it correctly, but here’s the code to trigger the exception:

import tvm
from tvm import relay
import numpy as np
from import download_testdata
# PyTorch imports
import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T
from torch_geometric.nn import GATConv

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = GATConv(1433, 16, cached=False,
        self.conv2 = GATConv(16, 7, cached=False,

    def forward(self, x, edge_index):
        c1 = self.conv1(x, edge_index)
        rc1 = F.relu(c1)
        d1 = F.dropout(rc1,
        c2 = self.conv2(d1, edge_index)
        r = F.log_softmax(c2, dim=1)
        return r

dataset = 'Cora'
path = osp.join('..', 'data', dataset)
dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)

inp = (data.x.cuda(), data.edge_index.cuda())
scripted_model = torch.jit.trace(model, inp).eval()

input_name = 'input0'
shape_list = [(input_name, data.x.size())]
mod, params = relay.frontend.from_pytorch(scripted_model,

It seems that a range of ops are not currently supported, including:

 ['aten::_set_item', 'prim::ImplicitTensorToNum', 'aten::__range_length', 'aten::numel', 'aten::__is__', 'prim::unchecked_cast', 'aten::index', 'aten::dim', 'prim::dtype', 'torch_scatter::scatter_max', 'aten::scatter_add_', 'aten::__isnot__', 'aten::index_select']

#5133 already addresses prim::ImplicitTensorToNum I believe, and torch_scatter::scatter_max is specific to torch-scatter’s implementation.

I’m wondering if concerns like this align with the current direction for development. Thanks!

Converting graph NN sounds interesting, and in principle it should be possible.

But for your particular case, aten::_set_item and aten::index_put_ seem like inplace update op (e.g. a[i] = ... in python). Our PyTorch frontend doesn’t support such side-effecting ops. Maybe it is possible using Relay’s reference, but no one have tried so far.

1 Like

I tried to specifically avoid side-effecting in my implementation, so I believe they are caused by torch_geometric’s implementation of message passing layers. I also checked that scatter_add_ and index_select are indeed Pytorch native functions, so maybe they are not yet implemented?

In addition, when using GNN architectures like Graph Attention Networks, torch_geometric will use external ops like torch_scatter::scatter_max. So generally, is it advisable to add support for non-native ops?

Since the number of torch extension is unbounded and application specific, we don’t support any of them out of the box. But it is possible to define a custom converter yourself and use it from the frontend.

See this example of how to convert torchvision’s custom op.

1 Like