Add additional kwargs to GATConv

python/cugraph-dgl/cugraph_dgl/nn/conv/gatconv.py

@@ -11,7 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 from cugraph_dgl.nn.conv.base import BaseConv, SparseGraph
 from cugraph.utilities.utils import import_optional
@@ -186,6 +186,7 @@ def forward(
  nfeat: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
  efeat: Optional[torch.Tensor] = None,
  max_in_degree: Optional[int] = None,
+ **kwargs: Any,
  ) -> torch.Tensor:
  r"""Forward computation.
 
@@ -204,6 +205,8 @@ def forward(
  from a neighbor sampler, the value should be set to the corresponding
  :attr:`fanout`. This option is used to invoke the MFG-variant of
  cugraph-ops kernel.
+ **kwargs : Any
+ Additional arguments of `pylibcugraphops.pytorch.operators.mha_gat_n2n`.
 
  Returns
  -------
@@ -232,6 +235,8 @@ def forward(
  _graph = self.get_cugraph_ops_CSC(
  g, is_bipartite=bipartite, max_in_degree=max_in_degree
  )
+ if kwargs.get("deterministic_dgrad", False):
+ _graph.add_reverse_graph()
 
  if bipartite:
  nfeat = (self.feat_drop(nfeat[0]), self.feat_drop(nfeat[1]))
@@ -273,6 +278,7 @@ def forward(
  negative_slope=self.negative_slope,
  concat_heads=self.concat,
  edge_feat=efeat,
+ **kwargs,
  )[: g.num_dst_nodes()]
 
  if self.concat:

python/cugraph-dgl/cugraph_dgl/nn/conv/gatv2conv.py

@@ -13,7 +13,7 @@
 
 from typing import Optional, Union
 
-from cugraph_dgl.nn.conv.base import BaseConv, SparseGraph
+from cugraph_dgl.nn.conv.base import Any, BaseConv, SparseGraph
 from cugraph.utilities.utils import import_optional
 
 dgl = import_optional("dgl")
@@ -150,6 +150,7 @@ def forward(
  nfeat: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
  efeat: Optional[torch.Tensor] = None,
  max_in_degree: Optional[int] = None,
+ **kwargs: Any,
  ) -> torch.Tensor:
  r"""Forward computation.
 
@@ -166,6 +167,8 @@ def forward(
  from a neighbor sampler, the value should be set to the corresponding
  :attr:`fanout`. This option is used to invoke the MFG-variant of
  cugraph-ops kernel.
+ **kwargs : Any
+ Additional arguments of `pylibcugraphops.pytorch.operators.mha_gat_v2_n2n`.
 
  Returns
  -------
@@ -196,6 +199,8 @@ def forward(
  _graph = self.get_cugraph_ops_CSC(
  g, is_bipartite=graph_bipartite, max_in_degree=max_in_degree
  )
+ if kwargs.get("deterministic_dgrad", False):
+ _graph.add_reverse_graph()
 
  if nfeat_bipartite:
  nfeat = (self.feat_drop(nfeat[0]), self.feat_drop(nfeat[1]))
@@ -228,6 +233,7 @@ def forward(
  negative_slope=self.negative_slope,
  concat_heads=self.concat,
  edge_feat=efeat,
+ **kwargs,
  )[: g.num_dst_nodes()]
 
  if self.concat:

python/cugraph-pyg/cugraph_pyg/nn/conv/gat_conv.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -11,7 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 from cugraph.utilities.utils import import_optional
 from pylibcugraphops.pytorch.operators import mha_gat_n2n
@@ -162,6 +162,7 @@ def forward(
  csc: Tuple[torch.Tensor, torch.Tensor, int],
  edge_attr: Optional[torch.Tensor] = None,
  max_num_neighbors: Optional[int] = None,
+ **kwargs: Any,
  ) -> torch.Tensor:
  r"""Runs the forward pass of the module.
 
@@ -178,11 +179,15 @@ def forward(
  of a destination node. When enabled, it allows models to use
  the message-flow-graph primitives in cugraph-ops.
  (default: :obj:`None`)
+ **kwargs : Additional arguments of
+ `pylibcugraphops.pytorch.operators.mha_gat_n2n`.
  """
  bipartite = not isinstance(x, torch.Tensor)
  graph = self.get_cugraph(
  csc, bipartite=bipartite, max_num_neighbors=max_num_neighbors
  )
+ if kwargs.get("deterministic_dgrad", False):
+ graph.add_reverse_graph()
 
  if edge_attr is not None:
  if self.lin_edge is None:
@@ -220,6 +225,7 @@ def forward(
  negative_slope=self.negative_slope,
  concat_heads=self.concat,
  edge_feat=edge_attr,
+ **kwargs,
  )
 
  if self.bias is not None:

python/cugraph-pyg/cugraph_pyg/nn/conv/gatv2_conv.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -11,7 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 from cugraph.utilities.utils import import_optional
 from pylibcugraphops.pytorch.operators import mha_gat_v2_n2n
@@ -174,6 +174,7 @@ def forward(
  x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
  csc: Tuple[torch.Tensor, torch.Tensor, int],
  edge_attr: Optional[torch.Tensor] = None,
+ **kwargs: Any,
  ) -> torch.Tensor:
  r"""Runs the forward pass of the module.
 
@@ -186,9 +187,13 @@ def forward(
  :meth:`to_csc` method to convert an :obj:`edge_index`
  representation to the desired format.
  edge_attr: (torch.Tensor, optional) The edge features.
+ **kwargs : Additional arguments of
+ `pylibcugraphops.pytorch.operators.mha_gat_v2_n2n`.
  """
  bipartite = not isinstance(x, torch.Tensor) or not self.share_weights
  graph = self.get_cugraph(csc, bipartite=bipartite)
+ if kwargs.get("deterministic_dgrad", False):
+ graph.add_reverse_graph()
 
  if edge_attr is not None:
  if self.lin_edge is None:
@@ -217,6 +222,7 @@ def forward(
  negative_slope=self.negative_slope,
  concat_heads=self.concat,
  edge_feat=edge_attr,
+ **kwargs,
  )
 
  if self.bias is not None: