diff --git a/CHANGELOG.md b/CHANGELOG.md index ff558764e23d..36f855cc0ee4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ## [2.0.5] - 2022-MM-DD ### Added +- Added `size=None` explanation to jittable `MessagePassing` modules in the documentation ([#4850](https://github.com/pyg-team/pytorch_geometric/pull/4850)) - Added documentation to the `DataLoaderIterator` class ([#4838](https://github.com/pyg-team/pytorch_geometric/pull/4838)) - Added `GraphStore` support to `Data` and `HeteroData` ([#4816](https://github.com/pyg-team/pytorch_geometric/pull/4816)) - Added `FeatureStore` support to `Data` and `HeteroData` ([#4807](https://github.com/pyg-team/pytorch_geometric/pull/4807)) diff --git a/docs/source/notes/jit.rst b/docs/source/notes/jit.rst index 76999da5dcc1..54891c910570 100644 --- a/docs/source/notes/jit.rst +++ b/docs/source/notes/jit.rst @@ -99,7 +99,8 @@ However, if you want your own GNN module to be jittable, you need to account for def forward(self, x: Tensor, edge_index: Tensor, edge_weight: Optional[Tensor]) -> Tensor: - return self.propagate(edge_index, x=x, edge_weight=edge_weight) + return self.propagate(edge_index, x=x, edge_weight=edge_weight, + size=None) 2. Declaring the type of propagation arguments as a comment anywhere inside your module: @@ -115,4 +116,10 @@ However, if you want your own GNN module to be jittable, you need to account for edge_weight: Optional[Tensor]) -> Tensor: # propagate_type: (x: Tensor, edge_weight: Optional[Tensor]) - return self.propagate(edge_index, x=x, edge_weight=edge_weight) + return self.propagate(edge_index, x=x, edge_weight=edge_weight, + size=None) + +.. warning:: + + Importantly, due to TorchScript limitations, one also has to pass in the :obj:`size` attribute to :meth:`~torch_geometric.nn.conv.message_passing.MessagePassing.propagate`. + In most cases, this can be simply set to :obj:`None` in which case it will be automatically inferred.