diff --git a/examples/contrib/pgm_explainer_graph_classification.py b/examples/contrib/pgm_explainer_graph_classification.py index 8960f1090f7f..cdd3acadee0c 100644 --- a/examples/contrib/pgm_explainer_graph_classification.py +++ b/examples/contrib/pgm_explainer_graph_classification.py @@ -1,5 +1,4 @@ -""" -This is an example of using the PGM explainer algorithm +"""This is an example of using the PGM explainer algorithm on a graph classification task """ import os.path as osp diff --git a/examples/contrib/pgm_explainer_node_classification.py b/examples/contrib/pgm_explainer_node_classification.py index 8b620a7d4223..3bb9c4a56624 100644 --- a/examples/contrib/pgm_explainer_node_classification.py +++ b/examples/contrib/pgm_explainer_node_classification.py @@ -1,5 +1,4 @@ -""" -This is an example of using the PGM explainer algorithm +"""This is an example of using the PGM explainer algorithm on a node classification task """ import os.path as osp diff --git a/examples/equilibrium_median.py b/examples/equilibrium_median.py index 2cd142efcac6..82cc32a5f05c 100644 --- a/examples/equilibrium_median.py +++ b/examples/equilibrium_median.py @@ -1,5 +1,4 @@ -r""" -Replicates the experiment from `"Deep Graph Infomax" +r"""Replicates the experiment from `"Deep Graph Infomax" `_ to try and teach `EquilibriumAggregation` to learn to take the median of a set of numbers diff --git a/examples/multi_gpu/distributed_sampling_xpu.py b/examples/multi_gpu/distributed_sampling_xpu.py index ebd3078abaa5..29f9a513ecc7 100644 --- a/examples/multi_gpu/distributed_sampling_xpu.py +++ b/examples/multi_gpu/distributed_sampling_xpu.py @@ -1,5 +1,4 @@ -""" -Distributed GAT training, targeting XPU devices. +"""Distributed GAT training, targeting XPU devices. PVC has 2 tiles, each reports itself as a separate device. DDP approach allows us to employ both tiles. diff --git a/examples/multi_gpu/multinode_multigpu_papers100m_gcn.py b/examples/multi_gpu/multinode_multigpu_papers100m_gcn.py index b8e3da188bf9..b08e9c7f244d 100644 --- a/examples/multi_gpu/multinode_multigpu_papers100m_gcn.py +++ b/examples/multi_gpu/multinode_multigpu_papers100m_gcn.py @@ -1,5 +1,4 @@ -""" -To run: +"""To run: srun -l -N --ntasks-per-node= \ --container-name=cont --container-image= \ --container-mounts=/ogb-papers100m/:/workspace/dataset diff --git a/examples/randlanet_classification.py b/examples/randlanet_classification.py index c4151af69a8d..c21b18757d0b 100644 --- a/examples/randlanet_classification.py +++ b/examples/randlanet_classification.py @@ -1,5 +1,4 @@ -""" -An adaptation of RandLA-Net to the classification task, which was not +"""An adaptation of RandLA-Net to the classification task, which was not addressed in the paper: RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds Reference: https://arxiv.org/abs/1911.11236 diff --git a/examples/randlanet_segmentation.py b/examples/randlanet_segmentation.py index 39602da37045..8f88d7e9988a 100644 --- a/examples/randlanet_segmentation.py +++ b/examples/randlanet_segmentation.py @@ -1,5 +1,4 @@ -""" -An implementation of RandLA-Net based on the paper: +"""An implementation of RandLA-Net based on the paper: RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds Reference: https://arxiv.org/abs/1911.11236 """ diff --git a/graphgym/custom_graphgym/encoder/example.py b/graphgym/custom_graphgym/encoder/example.py index d1edf30f7b88..c5cbd1ea82c5 100644 --- a/graphgym/custom_graphgym/encoder/example.py +++ b/graphgym/custom_graphgym/encoder/example.py @@ -9,8 +9,7 @@ @register_node_encoder('example') class ExampleNodeEncoder(torch.nn.Module): - """ - Provides an encoder for integer node features. + """Provides an encoder for integer node features. Args: num_classes (int): The number of classes for the embedding mapping to diff --git a/pyproject.toml b/pyproject.toml index 7443da07aa10..f8bbe34f3a57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,7 +131,6 @@ ignore = [ "D107", # TODO: Don't ignore "Missing docstring in __init__" "D200", # Ignore "One-line docstring should fit on one line" "D205", # Ignore "1 blank line required between summary line and description" - "D212", # Ignore "Multi-line docstring summary should start at the first line" "D415", # Ignore "First line should end with a period, question mark, or exclamation point" ] diff --git a/torch_geometric/data/feature_store.py b/torch_geometric/data/feature_store.py index 8f571472a7d2..093426e132e6 100644 --- a/torch_geometric/data/feature_store.py +++ b/torch_geometric/data/feature_store.py @@ -1,7 +1,7 @@ -r""" -This class defines the abstraction for a backend-agnostic feature store. The -goal of the feature store is to abstract away all node and edge feature memory -management so that varying implementations can allow for independent scale-out. +r"""This class defines the abstraction for a backend-agnostic feature store. +The goal of the feature store is to abstract away all node and edge feature +memory management so that varying implementations can allow for independent +scale-out. This particular feature store abstraction makes a few key assumptions: * The features we care about storing are node and edge features of a graph. diff --git a/torch_geometric/data/graph_store.py b/torch_geometric/data/graph_store.py index 4acfb77e84ad..cd591895d4ea 100644 --- a/torch_geometric/data/graph_store.py +++ b/torch_geometric/data/graph_store.py @@ -1,5 +1,4 @@ -r""" -This class defines the abstraction for a backend-agnostic graph store. The +r"""This class defines the abstraction for a backend-agnostic graph store. The goal of the graph store is to abstract away all graph edge index memory management so that varying implementations can allow for independent scale-out. diff --git a/torch_geometric/graphgym/config.py b/torch_geometric/graphgym/config.py index 94106da0e1f2..1cfc34193e9f 100644 --- a/torch_geometric/graphgym/config.py +++ b/torch_geometric/graphgym/config.py @@ -21,8 +21,7 @@ def set_cfg(cfg): - r""" - This function sets the default config value. + r"""This function sets the default config value. 1) Note that for an experiment, only part of the arguments will be used The remaining unused arguments won't affect anything. So feel free to register any argument in graphgym.contrib.config @@ -476,8 +475,7 @@ def assert_cfg(cfg): def dump_cfg(cfg): - r""" - Dumps the config to the output directory specified in + r"""Dumps the config to the output directory specified in :obj:`cfg.out_dir` Args: @@ -490,8 +488,7 @@ def dump_cfg(cfg): def load_cfg(cfg, args): - r""" - Load configurations from file system and command line + r"""Load configurations from file system and command line Args: cfg (CfgNode): Configuration node @@ -509,8 +506,7 @@ def makedirs_rm_exist(dir): def get_fname(fname): - r""" - Extract filename from file name path + r"""Extract filename from file name path Args: fname (str): Filename for the yaml format configuration file @@ -524,8 +520,7 @@ def get_fname(fname): def set_out_dir(out_dir, fname): - r""" - Create the directory for full experiment run + r"""Create the directory for full experiment run Args: out_dir (str): Directory for output, specified in :obj:`cfg.out_dir` @@ -541,8 +536,7 @@ def set_out_dir(out_dir, fname): def set_run_dir(out_dir): - r""" - Create the directory for each random seed experiment run + r"""Create the directory for each random seed experiment run Args: out_dir (str): Directory for output, specified in :obj:`cfg.out_dir` diff --git a/torch_geometric/graphgym/init.py b/torch_geometric/graphgym/init.py index f098c77467a0..79fe5311e15a 100644 --- a/torch_geometric/graphgym/init.py +++ b/torch_geometric/graphgym/init.py @@ -2,8 +2,7 @@ def init_weights(m): - r""" - Performs weight initialization + r"""Performs weight initialization Args: m (nn.Module): PyTorch module diff --git a/torch_geometric/graphgym/loader.py b/torch_geometric/graphgym/loader.py index 4fd5d49b27d2..094658df6f6d 100644 --- a/torch_geometric/graphgym/loader.py +++ b/torch_geometric/graphgym/loader.py @@ -49,8 +49,7 @@ def planetoid_dataset(name: str) -> Callable: def load_pyg(name, dataset_dir): - """ - Load PyG dataset objects. (More PyG datasets will be supported) + """Load PyG dataset objects. (More PyG datasets will be supported) Args: name (str): dataset name @@ -101,9 +100,7 @@ def set_dataset_attr(dataset, name, value, size): def load_ogb(name, dataset_dir): - r""" - - Load OGB dataset objects. + r"""Load OGB dataset objects. Args: @@ -173,9 +170,7 @@ def load_ogb(name, dataset_dir): def load_dataset(): - r""" - - Load dataset objects. + r"""Load dataset objects. Returns: PyG dataset object @@ -200,8 +195,7 @@ def load_dataset(): def set_dataset_info(dataset): - r""" - Set global dataset information + r"""Set global dataset information Args: dataset: PyG dataset object @@ -233,8 +227,7 @@ def set_dataset_info(dataset): def create_dataset(): - r""" - Create dataset object + r"""Create dataset object Returns: PyG dataset object @@ -311,8 +304,7 @@ def get_loader(dataset, sampler, batch_size, shuffle=True): def create_loader(): - """ - Create data loader object + """Create data loader object Returns: List of PyTorch data loaders diff --git a/torch_geometric/graphgym/logger.py b/torch_geometric/graphgym/logger.py index 38c04ecefb90..c58058e1637b 100644 --- a/torch_geometric/graphgym/logger.py +++ b/torch_geometric/graphgym/logger.py @@ -15,8 +15,7 @@ def set_printing(): - """ - Set up printing options + """Set up printing options """ logging.root.handlers = [] diff --git a/torch_geometric/graphgym/loss.py b/torch_geometric/graphgym/loss.py index 332196823ff6..380ebd4a66d2 100644 --- a/torch_geometric/graphgym/loss.py +++ b/torch_geometric/graphgym/loss.py @@ -6,8 +6,7 @@ def compute_loss(pred, true): - """ - Compute loss and prediction score + """Compute loss and prediction score Args: pred (torch.tensor): Unnormalized prediction diff --git a/torch_geometric/graphgym/models/transform.py b/torch_geometric/graphgym/models/transform.py index 802621292f15..7300659a9e32 100644 --- a/torch_geometric/graphgym/models/transform.py +++ b/torch_geometric/graphgym/models/transform.py @@ -4,8 +4,7 @@ def create_link_label(pos_edge_index, neg_edge_index): - """ - Create labels for link prediction, based on positive and negative edges + """Create labels for link prediction, based on positive and negative edges Args: pos_edge_index (torch.tensor): Positive edge index [2, num_edges] @@ -22,8 +21,7 @@ def create_link_label(pos_edge_index, neg_edge_index): def neg_sampling_transform(data): - """ - Do negative sampling for link prediction tasks + """Do negative sampling for link prediction tasks Args: data (torch_geometric.data): Input data object diff --git a/torch_geometric/graphgym/utils/agg_runs.py b/torch_geometric/graphgym/utils/agg_runs.py index b1b66535469e..06aa0827dd97 100644 --- a/torch_geometric/graphgym/utils/agg_runs.py +++ b/torch_geometric/graphgym/utils/agg_runs.py @@ -43,8 +43,7 @@ def join_list(l1, l2): def agg_dict_list(dict_list): - """ - Aggregate a list of dictionaries: mean + std + """Aggregate a list of dictionaries: mean + std Args: dict_list: list of dictionaries @@ -80,8 +79,7 @@ def rm_keys(dict, keys): def agg_runs(dir, metric_best='auto'): - r""" - Aggregate over different random seeds of a single experiment + r"""Aggregate over different random seeds of a single experiment Args: dir (str): Directory of the results, containing 1 experiment @@ -161,8 +159,7 @@ def agg_runs(dir, metric_best='auto'): def agg_batch(dir, metric_best='auto'): - r""" - Aggregate across results from multiple experiments via grid search + r"""Aggregate across results from multiple experiments via grid search Args: dir (str): Directory of the results, containing multiple experiments diff --git a/torch_geometric/nn/aggr/base.py b/torch_geometric/nn/aggr/base.py index 91d20d751fc1..2028a234748a 100644 --- a/torch_geometric/nn/aggr/base.py +++ b/torch_geometric/nn/aggr/base.py @@ -68,7 +68,8 @@ def forward( dim: int = -2, max_num_elements: Optional[int] = None, ) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. index (torch.Tensor, optional): The indices of elements for diff --git a/torch_geometric/nn/aggr/equilibrium.py b/torch_geometric/nn/aggr/equilibrium.py index 1357696a9f2f..46190c6ceb2e 100644 --- a/torch_geometric/nn/aggr/equilibrium.py +++ b/torch_geometric/nn/aggr/equilibrium.py @@ -48,8 +48,7 @@ def forward(self, x: Tensor, y: Tensor, index: Optional[Tensor], class MomentumOptimizer(torch.nn.Module): - r""" - Provides an inner loop optimizer for the implicitly defined output + r"""Provides an inner loop optimizer for the implicitly defined output layer. It is based on an unrolled Nesterov momentum algorithm. Args: diff --git a/torch_geometric/nn/attention/performer.py b/torch_geometric/nn/attention/performer.py index 77060808d797..e4a2fe7e672f 100644 --- a/torch_geometric/nn/attention/performer.py +++ b/torch_geometric/nn/attention/performer.py @@ -145,7 +145,7 @@ def __init__( self.dropout = torch.nn.Dropout(dropout) def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor: - r""" + r"""Forward pass. Args: x (torch.Tensor): Node feature tensor diff --git a/torch_geometric/nn/conv/han_conv.py b/torch_geometric/nn/conv/han_conv.py index c88bd2e127c7..6b061921a5a9 100644 --- a/torch_geometric/nn/conv/han_conv.py +++ b/torch_geometric/nn/conv/han_conv.py @@ -31,8 +31,7 @@ def group( class HANConv(MessagePassing): - r""" - The Heterogenous Graph Attention Operator from the + r"""The Heterogenous Graph Attention Operator from the `"Heterogenous Graph Attention Network" `_ paper. diff --git a/torch_geometric/nn/dense/dense_gat_conv.py b/torch_geometric/nn/dense/dense_gat_conv.py index a7202f27d322..5e11492a327d 100644 --- a/torch_geometric/nn/dense/dense_gat_conv.py +++ b/torch_geometric/nn/dense/dense_gat_conv.py @@ -55,7 +55,8 @@ def reset_parameters(self): def forward(self, x: Tensor, adj: Tensor, mask: Optional[Tensor] = None, add_loop: bool = True): - r""" + r"""Forward pass. + Args: x (torch.Tensor): Node feature tensor :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with diff --git a/torch_geometric/nn/dense/dense_gcn_conv.py b/torch_geometric/nn/dense/dense_gcn_conv.py index 0e75ad1cd0eb..471889428d36 100644 --- a/torch_geometric/nn/dense/dense_gcn_conv.py +++ b/torch_geometric/nn/dense/dense_gcn_conv.py @@ -39,7 +39,8 @@ def reset_parameters(self): def forward(self, x: Tensor, adj: Tensor, mask: OptTensor = None, add_loop: bool = True) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): Node feature tensor :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with diff --git a/torch_geometric/nn/dense/dense_gin_conv.py b/torch_geometric/nn/dense/dense_gin_conv.py index d33033e24794..2a19094d3782 100644 --- a/torch_geometric/nn/dense/dense_gin_conv.py +++ b/torch_geometric/nn/dense/dense_gin_conv.py @@ -32,7 +32,8 @@ def reset_parameters(self): def forward(self, x: Tensor, adj: Tensor, mask: Optional[Tensor] = None, add_loop: bool = True) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): Node feature tensor :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with diff --git a/torch_geometric/nn/dense/dense_graph_conv.py b/torch_geometric/nn/dense/dense_graph_conv.py index 720c2f739d3f..d9ed09cf790c 100644 --- a/torch_geometric/nn/dense/dense_graph_conv.py +++ b/torch_geometric/nn/dense/dense_graph_conv.py @@ -33,7 +33,8 @@ def reset_parameters(self): def forward(self, x: Tensor, adj: Tensor, mask: Optional[Tensor] = None) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): Node feature tensor :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with diff --git a/torch_geometric/nn/dense/dense_sage_conv.py b/torch_geometric/nn/dense/dense_sage_conv.py index 3a957f5da3ed..6e4f9b60f063 100644 --- a/torch_geometric/nn/dense/dense_sage_conv.py +++ b/torch_geometric/nn/dense/dense_sage_conv.py @@ -41,7 +41,8 @@ def reset_parameters(self): def forward(self, x: Tensor, adj: Tensor, mask: OptTensor = None) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): Node feature tensor :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with diff --git a/torch_geometric/nn/dense/dmon_pool.py b/torch_geometric/nn/dense/dmon_pool.py index 74a1265f7692..42a2dd47f7a0 100644 --- a/torch_geometric/nn/dense/dmon_pool.py +++ b/torch_geometric/nn/dense/dmon_pool.py @@ -81,7 +81,8 @@ def forward( adj: Tensor, mask: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: - r""" + r"""Forward pass. + Args: x (torch.Tensor): Node feature tensor :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py index 6e1fad32c1f8..55e2fe176ff0 100644 --- a/torch_geometric/nn/dense/linear.py +++ b/torch_geometric/nn/dense/linear.py @@ -137,7 +137,8 @@ def reset_parameters(self): reset_bias_(self.bias, self.in_channels, self.bias_initializer) def forward(self, x: Tensor) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The input features. """ @@ -257,7 +258,8 @@ def reset_parameters(self): self.kwargs.get('bias_initializer', None)) def forward(self, x: Tensor, type_vec: Tensor) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The input features. type_vec (torch.Tensor): A vector that maps each entry to a type. @@ -404,7 +406,8 @@ def forward( self, x_dict: Dict[str, Tensor], ) -> Dict[str, Tensor]: - r""" + r"""Forward pass. + Args: x_dict (Dict[Any, torch.Tensor]): A dictionary holding input features for each individual type. diff --git a/torch_geometric/nn/models/basic_gnn.py b/torch_geometric/nn/models/basic_gnn.py index e23bc197c91c..350545b11448 100644 --- a/torch_geometric/nn/models/basic_gnn.py +++ b/torch_geometric/nn/models/basic_gnn.py @@ -214,7 +214,8 @@ def forward( # noqa num_sampled_nodes_per_hop: Optional[List[int]] = None, num_sampled_edges_per_hop: Optional[List[int]] = None, ): - r""" + r"""Forward pass. + Args: x (torch.Tensor): The input node features. edge_index (torch.Tensor or SparseTensor): The edge indices. diff --git a/torch_geometric/nn/models/correct_and_smooth.py b/torch_geometric/nn/models/correct_and_smooth.py index 885cbc07a951..f2847755785f 100644 --- a/torch_geometric/nn/models/correct_and_smooth.py +++ b/torch_geometric/nn/models/correct_and_smooth.py @@ -80,7 +80,8 @@ def forward(self, y_soft: Tensor, *args) -> Tensor: # pragma: no cover def correct(self, y_soft: Tensor, y_true: Tensor, mask: Tensor, edge_index: Adj, edge_weight: OptTensor = None) -> Tensor: - r""" + r"""Forward pass. + Args: y_soft (torch.Tensor): The soft predictions :math:`\mathbf{Z}` obtained from a simple base predictor. @@ -124,7 +125,8 @@ def fix_input(x): def smooth(self, y_soft: Tensor, y_true: Tensor, mask: Tensor, edge_index: Adj, edge_weight: OptTensor = None) -> Tensor: - r""" + r"""Forward pass. + Args: y_soft (torch.Tensor): The corrected predictions :math:`\mathbf{Z}` obtained from :meth:`correct`. diff --git a/torch_geometric/nn/models/dimenet.py b/torch_geometric/nn/models/dimenet.py index f7f666534260..cca5b36b898c 100644 --- a/torch_geometric/nn/models/dimenet.py +++ b/torch_geometric/nn/models/dimenet.py @@ -676,7 +676,8 @@ def forward( pos: Tensor, batch: OptTensor = None, ) -> Tensor: - r""" + r"""Forward pass. + Args: z (torch.Tensor): Atomic number of each atom with shape :obj:`[num_atoms]`. diff --git a/torch_geometric/nn/models/graph_mixer.py b/torch_geometric/nn/models/graph_mixer.py index b1247754c97b..75680efd7267 100644 --- a/torch_geometric/nn/models/graph_mixer.py +++ b/torch_geometric/nn/models/graph_mixer.py @@ -37,7 +37,8 @@ def forward( edge_time: Tensor, seed_time: Tensor, ) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The input node features. edge_index (torch.Tensor): The edge indices. @@ -98,7 +99,8 @@ def reset_parameters(self): self.head_lin.reset_parameters() def forward(self, x: Tensor) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): Tensor of size :obj:`[*, num_tokens, in_channels]`. @@ -230,7 +232,8 @@ def forward( edge_time: Tensor, seed_time: Tensor, ) -> Tensor: - r""" + r"""Forward pass. + Args: edge_index (torch.Tensor): The edge indices. edge_attr (torch.Tensor): The edge features of shape diff --git a/torch_geometric/nn/models/jumping_knowledge.py b/torch_geometric/nn/models/jumping_knowledge.py index d62f37f8af7c..3dd1ba4d1b16 100644 --- a/torch_geometric/nn/models/jumping_knowledge.py +++ b/torch_geometric/nn/models/jumping_knowledge.py @@ -69,7 +69,8 @@ def reset_parameters(self): self.att.reset_parameters() def forward(self, xs: List[Tensor]) -> Tensor: - r""" + r"""Forward pass. + Args: xs (List[torch.Tensor]): List containing the layer-wise representations. diff --git a/torch_geometric/nn/models/label_prop.py b/torch_geometric/nn/models/label_prop.py index 810c5fdb5514..1a1310a9d00b 100644 --- a/torch_geometric/nn/models/label_prop.py +++ b/torch_geometric/nn/models/label_prop.py @@ -48,7 +48,8 @@ def forward( edge_weight: OptTensor = None, post_step: Optional[Callable[[Tensor], Tensor]] = None, ) -> Tensor: - r""" + r"""Forward pass. + Args: y (torch.Tensor): The ground-truth label information :math:`\mathbf{Y}`. diff --git a/torch_geometric/nn/models/meta.py b/torch_geometric/nn/models/meta.py index b9e8f3672749..64f676dab3f2 100644 --- a/torch_geometric/nn/models/meta.py +++ b/torch_geometric/nn/models/meta.py @@ -123,7 +123,8 @@ def forward( u: Optional[Tensor] = None, batch: Optional[Tensor] = None, ) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor]]: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The node features. edge_index (torch.Tensor): The edge indices. diff --git a/torch_geometric/nn/models/mlp.py b/torch_geometric/nn/models/mlp.py index 554d8afaf384..b8dcc6c9981e 100644 --- a/torch_geometric/nn/models/mlp.py +++ b/torch_geometric/nn/models/mlp.py @@ -200,7 +200,8 @@ def forward( batch_size: Optional[int] = None, return_emb: NoneType = None, ) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. batch (torch.Tensor, optional): The batch vector diff --git a/torch_geometric/nn/models/schnet.py b/torch_geometric/nn/models/schnet.py index 6230bd3e5037..c01c43ffa226 100644 --- a/torch_geometric/nn/models/schnet.py +++ b/torch_geometric/nn/models/schnet.py @@ -269,7 +269,8 @@ def from_qm9_pretrained( def forward(self, z: Tensor, pos: Tensor, batch: OptTensor = None) -> Tensor: - r""" + r"""Forward pass. + Args: z (torch.Tensor): Atomic number of each atom with shape :obj:`[num_atoms]`. @@ -342,7 +343,8 @@ def __init__(self, cutoff: float = 10.0, max_num_neighbors: int = 32): self.max_num_neighbors = max_num_neighbors def forward(self, pos: Tensor, batch: Tensor) -> Tuple[Tensor, Tensor]: - r""" + r"""Forward pass. + Args: pos (Tensor): Coordinates of each atom. batch (LongTensor, optional): Batch indices assigning each atom to diff --git a/torch_geometric/nn/norm/batch_norm.py b/torch_geometric/nn/norm/batch_norm.py index 7aa6e9fd5a13..4e4a7899a9a8 100644 --- a/torch_geometric/nn/norm/batch_norm.py +++ b/torch_geometric/nn/norm/batch_norm.py @@ -69,7 +69,8 @@ def reset_parameters(self): self.module.reset_parameters() def forward(self, x: Tensor) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. """ @@ -169,7 +170,8 @@ def reset_parameters(self): torch.nn.init.zeros_(self.bias) def forward(self, x: Tensor, type_vec: Tensor) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The input features. type_vec (torch.Tensor): A vector that maps each entry to a type. diff --git a/torch_geometric/nn/norm/diff_group_norm.py b/torch_geometric/nn/norm/diff_group_norm.py index 0172e35a604b..b69eeeb196f9 100644 --- a/torch_geometric/nn/norm/diff_group_norm.py +++ b/torch_geometric/nn/norm/diff_group_norm.py @@ -68,7 +68,8 @@ def reset_parameters(self): self.norm.reset_parameters() def forward(self, x: Tensor) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. """ diff --git a/torch_geometric/nn/norm/graph_norm.py b/torch_geometric/nn/norm/graph_norm.py index 3373b32f7d62..ccb5c24ad067 100644 --- a/torch_geometric/nn/norm/graph_norm.py +++ b/torch_geometric/nn/norm/graph_norm.py @@ -47,7 +47,8 @@ def reset_parameters(self): def forward(self, x: Tensor, batch: OptTensor = None, batch_size: Optional[int] = None) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. batch (torch.Tensor, optional): The batch vector diff --git a/torch_geometric/nn/norm/graph_size_norm.py b/torch_geometric/nn/norm/graph_size_norm.py index 567e85a661e1..92e1b1ccb234 100644 --- a/torch_geometric/nn/norm/graph_size_norm.py +++ b/torch_geometric/nn/norm/graph_size_norm.py @@ -21,7 +21,8 @@ def __init__(self): def forward(self, x: Tensor, batch: OptTensor = None, batch_size: Optional[int] = None) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. batch (torch.Tensor, optional): The batch vector diff --git a/torch_geometric/nn/norm/instance_norm.py b/torch_geometric/nn/norm/instance_norm.py index 00a892ec820a..ecc972dc03ab 100644 --- a/torch_geometric/nn/norm/instance_norm.py +++ b/torch_geometric/nn/norm/instance_norm.py @@ -54,7 +54,8 @@ def reset_parameters(self): def forward(self, x: Tensor, batch: OptTensor = None, batch_size: Optional[int] = None) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. batch (torch.Tensor, optional): The batch vector diff --git a/torch_geometric/nn/norm/layer_norm.py b/torch_geometric/nn/norm/layer_norm.py index 7fb56d71905f..edf620deb2d4 100644 --- a/torch_geometric/nn/norm/layer_norm.py +++ b/torch_geometric/nn/norm/layer_norm.py @@ -66,7 +66,8 @@ def reset_parameters(self): def forward(self, x: Tensor, batch: OptTensor = None, batch_size: Optional[int] = None) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. batch (torch.Tensor, optional): The batch vector @@ -171,7 +172,8 @@ def forward( type_vec: OptTensor = None, type_ptr: Optional[Union[Tensor, List[int]]] = None, ) -> Tensor: - r""" + r"""Forward pass. + .. note:: Either :obj:`type_vec` or :obj:`type_ptr` needs to be specified. In general, relying on :obj:`type_ptr` is more efficient in case diff --git a/torch_geometric/nn/norm/mean_subtraction_norm.py b/torch_geometric/nn/norm/mean_subtraction_norm.py index 220c79f82d17..ab0f5a468dbb 100644 --- a/torch_geometric/nn/norm/mean_subtraction_norm.py +++ b/torch_geometric/nn/norm/mean_subtraction_norm.py @@ -17,7 +17,8 @@ class MeanSubtractionNorm(torch.nn.Module): """ def forward(self, x: Tensor, batch: Optional[Tensor] = None, dim_size: Optional[int] = None) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. batch (torch.Tensor, optional): The batch vector diff --git a/torch_geometric/nn/norm/msg_norm.py b/torch_geometric/nn/norm/msg_norm.py index af4678e9cab3..a7fc26309fda 100644 --- a/torch_geometric/nn/norm/msg_norm.py +++ b/torch_geometric/nn/norm/msg_norm.py @@ -30,7 +30,8 @@ def reset_parameters(self): self.scale.data.fill_(1.0) def forward(self, x: Tensor, msg: Tensor, p: float = 2.0) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. msg (torch.Tensor): The message tensor :math:`\mathbf{M}`. diff --git a/torch_geometric/nn/norm/pair_norm.py b/torch_geometric/nn/norm/pair_norm.py index c548ddc6c242..809f78ba9d10 100644 --- a/torch_geometric/nn/norm/pair_norm.py +++ b/torch_geometric/nn/norm/pair_norm.py @@ -40,7 +40,8 @@ def __init__(self, scale: float = 1., scale_individually: bool = False, def forward(self, x: Tensor, batch: OptTensor = None, batch_size: Optional[int] = None) -> Tensor: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The source tensor. batch (torch.Tensor, optional): The batch vector diff --git a/torch_geometric/nn/pool/asap.py b/torch_geometric/nn/pool/asap.py index c87aad1bfbd5..d4abe8fe1267 100644 --- a/torch_geometric/nn/pool/asap.py +++ b/torch_geometric/nn/pool/asap.py @@ -89,7 +89,8 @@ def forward( edge_weight: Optional[Tensor] = None, batch: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor, Optional[Tensor], Tensor, Tensor]: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The node feature matrix. edge_index (torch.Tensor): The edge indices. diff --git a/torch_geometric/nn/pool/connect/base.py b/torch_geometric/nn/pool/connect/base.py index e3ddf3b06f1d..5fa3d283853b 100644 --- a/torch_geometric/nn/pool/connect/base.py +++ b/torch_geometric/nn/pool/connect/base.py @@ -75,7 +75,8 @@ def forward( edge_attr: Optional[Tensor] = None, batch: Optional[Tensor] = None, ) -> ConnectOutput: - r""" + r"""Forward pass. + Args: select_output (SelectOutput): The output of :class:`Select`. edge_index (torch.Tensor): The edge indices. diff --git a/torch_geometric/nn/pool/edge_pool.py b/torch_geometric/nn/pool/edge_pool.py index ffe411687229..7d7c9db36e89 100644 --- a/torch_geometric/nn/pool/edge_pool.py +++ b/torch_geometric/nn/pool/edge_pool.py @@ -110,7 +110,8 @@ def forward( edge_index: Tensor, batch: Tensor, ) -> Tuple[Tensor, Tensor, Tensor, UnpoolInfo]: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The node features. edge_index (torch.Tensor): The edge indices. diff --git a/torch_geometric/nn/pool/mem_pool.py b/torch_geometric/nn/pool/mem_pool.py index c0b161797ca4..c42923c2d19a 100644 --- a/torch_geometric/nn/pool/mem_pool.py +++ b/torch_geometric/nn/pool/mem_pool.py @@ -86,7 +86,8 @@ def forward( max_num_nodes: Optional[int] = None, batch_size: Optional[int] = None, ) -> Tuple[Tensor, Tensor]: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The node feature tensor of shape :math:`\mathbf{X} \in \mathbb{R}^{N \times F}` or diff --git a/torch_geometric/nn/pool/pan_pool.py b/torch_geometric/nn/pool/pan_pool.py index a8b1ccc6791e..65df76a79b3c 100644 --- a/torch_geometric/nn/pool/pan_pool.py +++ b/torch_geometric/nn/pool/pan_pool.py @@ -72,7 +72,8 @@ def forward( M: SparseTensor, batch: OptTensor = None, ) -> Tuple[Tensor, Tensor, Tensor, OptTensor, Tensor, Tensor]: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The node feature matrix. M (SparseTensor): The MET matrix :math:`\mathbf{M}`. diff --git a/torch_geometric/nn/pool/sag_pool.py b/torch_geometric/nn/pool/sag_pool.py index 8ae8c5068604..19a45936eff1 100644 --- a/torch_geometric/nn/pool/sag_pool.py +++ b/torch_geometric/nn/pool/sag_pool.py @@ -105,7 +105,8 @@ def forward( batch: OptTensor = None, attn: OptTensor = None, ) -> Tuple[Tensor, Tensor, OptTensor, OptTensor, Tensor, Tensor]: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The node feature matrix. edge_index (torch.Tensor): The edge indices. diff --git a/torch_geometric/nn/pool/topk_pool.py b/torch_geometric/nn/pool/topk_pool.py index 911a1275d9c2..1a4ea4ab7fdb 100644 --- a/torch_geometric/nn/pool/topk_pool.py +++ b/torch_geometric/nn/pool/topk_pool.py @@ -94,7 +94,8 @@ def forward( batch: Optional[Tensor] = None, attn: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor, OptTensor, OptTensor, Tensor, Tensor]: - r""" + r"""Forward pass. + Args: x (torch.Tensor): The node feature matrix. edge_index (torch.Tensor): The edge indices.