How to use TorchScript to save Heterogeneous Graphs?

[wumian0929]

I'm trying to use TorchScript to save a self-defined heterogeneous graph model. The conv operator and the forward() function is defined as follows,
conv = pyg_nn.HeteroConv({ ('node1', 'inner', 'node1'): pyg_nn.LEConv(-1, hidden_channels, add_self_loops=True).jittable(), ('node2', 'belong', 'node1'): pyg_nn.LEConv(-1, hidden_channels, add_self_loops=False).jittable(), ('node2', 'inner', 'node2'): pyg_nn.LEConv(-1, hidden_channels, add_self_loops=False).jittable(), }, aggr='sum')
def forward(self, x_dict, edge_index_dict, output_adj_matrix): for conv in self.convs: x_dict = conv(x_dict, edge_index_dict) x_dict = {key: x.relu() for key, x in x_dict.items()} y_pred = torch.matmul(x_dict["node1"].T, output_adj_matrix) y_pred = self.trans(y_pred.T) return y_pred
As the type of x_dict and edge_index are dict rather than tensor, the torch.jit raised a NotSupportedError:

File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/torch/jit/frontend.py", line 359, in build_param_list
raise NotSupportedError(ctx_range, _vararg_kwarg_err)
torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/torch_geometric/nn/conv/hetero_conv.py", line 76
edge_index_dict: Dict[EdgeType, Adj],
*args_dict,
**kwargs_dict,
~~~~~~~~~~~~ <--- HERE
) -> Dict[NodeType, Tensor]:
r"""

How to declare the dict type to convert the model a TorchScript program? Looking forward to your reply.

[rusty1s]

Currently, HeteroConv does not support TorchScript since it uses *args and **kwargs logic internally. You can probably bypass this limitation by defining a custom HeteroConv wrapper without this logic. Would that work for you?

[wumian0929]

I'll try to define the heterogeneous module to figure this out. Thanks a lot for your help!