torch_geometric.nn.conv.CuGraphGATConv

class CuGraphGATConv(in_channels: int, out_channels: int, heads: int = 1, concat: bool = True, negative_slope: float = 0.2, bias: bool = True)[source]

Bases: CuGraphModule

The graph attentional operator from the “Graph Attention Networks” paper.

CuGraphGATConv is an optimized version of GATConv based on the cugraph-ops package that fuses message passing computation for accelerated execution and lower memory footprint.

forward(x: Tensor, csc: Tuple[Tensor, Tensor, int], max_num_neighbors: Optional[int] = None) → Tensor[source]

Runs the forward pass of the module.

Parameters

  • x (torch.Tensor) – The node features.

  • csc ((torch.Tensor, torch.Tensor, int)) – A tuple containing the CSC representation of a graph, given as a tuple of (row, colptr, num_src_nodes). Use the CuGraphModule.to_csc() method to convert an edge_index representation to the desired format.

  • max_num_neighbors (int, optional) – The maximum number of neighbors of a target node. It is only effective when operating in a bipartite graph. When not given, the value will be computed on-the-fly, leading to slightly worse performance. (default: None)

reset_parameters()[source]

Resets all learnable parameters of the module.