greatx.functional

to_sparse_tensor

Convert edge index to a torch_sparse.SparseTensor

to_dense_adj

Convert edge index to dense adjacency matrix torch.FloatTensor

to_sparse_adj

Convert edge index to sparse adjacency matrix torch.sparse.Tensor

spmm

Sparse-dense matrix multiplication.

drop_edge

DropEdge: Sampling edge using a uniform distribution from the "DropEdge: Towards Deep Graph Convolutional Networks on Node Classification" paper (ICLR'20)

drop_node

DropNode: Sampling node using a uniform distribution from the "Graph Contrastive Learning with Augmentations" paper (NeurIPS'20)

drop_path

DropPath: a structured form of greatx.functional.drop_edge from the "MaskGAE: Masked Graph Modeling Meets Graph Autoencoders" paper (arXiv'22)
to_sparse_tensor(edge_index: Tensor, edge_weight: Optional[Tensor] = None, num_nodes: Optional[int] = None, is_sorted: bool = False) → SparseTensor[source]

Convert edge index to a torch_sparse.SparseTensor

Parameters:

  • edge_index (torch.Tensor) – edge index with shape [2, M]

  • edge_weight (Optional[Tensor], optional) – edge weight with shape [M], by default None

  • num_nodes (Optional[int], optional) – the number of nodes in the graph, by default None

  • is_sorted (bool, optional) – whether the edge_index is sorted, by default False

Returns:

the output sparse adjacency matrix denoted as torch_sparse.SparseTensor, with shape [num_nodes, num_nodes]

Return type:

torch_sparse.SparseTensor

to_dense_adj(edge_index: Tensor, edge_weight: Optional[Tensor] = None, num_nodes: Optional[int] = None, fill_value: float = 1.0) → Tensor[source]

Convert edge index to dense adjacency matrix torch.FloatTensor

Parameters:

  • edge_index (torch.Tensor) – edge index with shape [2, M]

  • edge_weight (Optional[Tensor], optional) – edge weight with shape [M], by default None

  • num_nodes (Optional[int], optional) – the number of nodes in the graph, by default None

  • fill_value (float) – filling value for elements in the adjacency matrix where edges existed, by default 1.0

Returns:

output dense adjacency matrix with shape [num_nodes, num_nodes]

Return type:

torch.Tensor

to_sparse_adj(edge_index: Tensor, edge_attr: Optional[Tensor] = None, size: Optional[int] = None) → Tensor[source]

Convert edge index to sparse adjacency matrix torch.sparse.Tensor

Parameters:

  • edge_index (torch.Tensor) – edge index with shape [2, M]

  • edge_attr (Optional[Tensor], optional) – edge attributes, by default None

  • size (Optional[int], optional) –

    The size of the sparse matrix.

    If given as an integer, will create a quadratic sparse matrix. If set to None, will infer a quadratic sparse matrix based on edge_index.max() + 1. By default None

Returns:

  • torch.sparse.Tensor – output sparse adjacency matrix

  • Example

  • ——-

  • >>> edge_index = torch.tensor([[0, 1, 1, 2, 2, 3],

  • … [1, 0, 2, 1, 3, 2]])

  • >>> to_torch_coo_tensor(edge_index)

  • tensor(indices=tensor([[0, 1, 1, 2, 2, 3], – [1, 0, 2, 1, 3, 2]]), values=tensor([1., 1., 1., 1., 1., 1.]), size=(4, 4), nnz=6, layout=torch.sparse_coo)

spmm(x: Tensor, edge_index: Union[Tensor, SparseTensor], edge_weight: Optional[Tensor] = None, reduce: str = 'sum') → Tensor[source]

Sparse-dense matrix multiplication.

Parameters:

  • x (torch.Tensor) – the input dense 2D-matrix

  • edge_index (torch.Tensor) – the location of the non-zeros elements in the sparse matrix, denoted as edge_index with shape [2, M]

  • edge_weight (Optional[Tensor], optional) – the edge weight of the sparse matrix, by default None

  • reduce (str, optional) – reduction of the sparse matrix multiplication, including: ('mean', 'sum', 'add', 'max', 'min', 'median', 'sample_median') by default 'sum'

Returns:

the output result of the matrix multiplication.

Return type:

Tensor

Example

import torch
from greatx.functional import spmm

x = torch.randn(5, 2)
edge_index = torch.LongTensor([[1,2], [3,4]])
out1 = spmm(x, edge_index, reduce='sum')

# which is equivalent to:
A = torch.zeros(5, 5)
A[edge_index[0], edge_index[1]] = 1.0
out2 = torch.mm(A.t(), x)

assert torch.allclose(out1, out2)

# Also, it also supports :obj:`torch.sparse.Tensor`
# and :obj:`torch_sparse.SparseTensor`
A = A.to_sparse()
out3 = spmm(x, A.t())
assert torch.allclose(out1, out3)

A = SparseTensor.from_torch_sparse_coo_tensor(A)
out4 = spmm(x, A.t())
assert torch.allclose(out1, out4)

See also

spmm

drop_edge(edge_index: Tensor, edge_weight: Optional[Tensor] = None, p: float = 0.5, training: bool = True) → Tuple[Tensor, Optional[Tensor]][source]

DropEdge: Sampling edge using a uniform distribution from the “DropEdge: Towards Deep Graph Convolutional Networks on Node Classification” paper (ICLR’20)

Parameters:

  • edge_index (torch.Tensor) – the input edge index

  • edge_weight (Optional[Tensor], optional) – the input edge weight, by default None

  • p (float, optional) – the probability of dropping out on each edge, by default 0.5

  • training (bool, optional) – whether the model is during training, do nothing if training=True, by default True

Returns:

the output edge index and edge weight

Return type:

Tuple[Tensor, Optional[Tensor]]

Raises:

ValueError – p is out of range [0,1]

Example

from greatx.functional import drop_edge
edge_index = torch.LongTensor([[1, 2], [3,4]])
drop_edge(edge_index, p=0.5)

See also

greatx.nn.layers.DropEdge

drop_node(edge_index: Tensor, edge_weight: Optional[Tensor] = None, p: float = 0.5, training: bool = True, num_nodes: Optional[int] = None) → Tuple[Tensor, Optional[Tensor]][source]

DropNode: Sampling node using a uniform distribution from the “Graph Contrastive Learning with Augmentations” paper (NeurIPS’20)

Parameters:

  • edge_index (torch.Tensor) – the input edge index

  • edge_weight (Optional[Tensor], optional) – the input edge weight, by default None

  • p (float, optional) – the probability of dropping out on each node, by default 0.5

  • training (bool, optional) – whether the model is during training, do nothing if training=True, by default True

Returns:

the output edge index and edge weight

Return type:

Tuple[Tensor, Optional[Tensor]]

Raises:

ValueError – p is out of range [0,1]

Example

from greatx.functional import drop_node
edge_index = torch.LongTensor([[1, 2], [3,4]])
drop_node(edge_index, p=0.5)

See also

greatx.nn.layers.DropNode

drop_path(edge_index: Tensor, edge_weight: Optional[Tensor] = None, p: float = 0.5, walks_per_node: int = 1, walk_length: int = 3, num_nodes: Optional[int] = None, start: str = 'node', is_sorted: bool = False, training: bool = True) → Tuple[Tensor, Optional[Tensor]][source]

DropPath: a structured form of greatx.functional.drop_edge from the “MaskGAE: Masked Graph Modeling Meets Graph Autoencoders” paper (arXiv’22)

Parameters:

  • edge_index (torch.Tensor) – the input edge index

  • edge_weight (Optional[Tensor], optional) – the input edge weight, by default None

  • p (Optional[Union[float, Tensor]], optional) – If p is a float value - the percentage of nodes in the graph that chosen as root nodes to perform random walks. If p is torch.Tensor - a set of custom root nodes. By default, p=0.5.

  • walks_per_node (int, optional) – number of walks per node, by default 1

  • walk_length (int, optional) – number of walk length per node, by default 3

  • num_nodes (int, optional) – number of total nodes in the graph, by default None

  • start (string, optional) – the type of starting node chosen from node of edge, by default ‘node’

  • is_sorted (bool, optional) – whether the input edge_index is sorted

  • training (bool, optional) – whether the model is during training, do nothing if training=True, by default True

Returns:

the output edge index and edge weight

Return type:

Tuple[Tensor, Optional[Tensor]]

Raises:

Example

from greatx.functional import drop_path
edge_index = torch.LongTensor([[1, 2], [3,4]])
drop_path(edge_index, p=0.5)

drop_path(edge_index, p=torch.tensor([1,2])) # specify root nodes

See also

greatx.nn.layers.DropPath