torch_geometric.EdgeIndex
- class EdgeIndex(data: Any, *args: Any, sparse_size: Optional[Tuple[Optional[int], Optional[int]]] = None, sort_order: Optional[Union[str, SortOrder]] = None, is_undirected: bool = False, **kwargs: Any)[source]
基础类：
Tensor一个带有附加（元）数据的COO
edge_index张量。EdgeIndex是一个torch.Tensor，它保存了形状为[2, num_edges]的edge_index表示。边以稀疏 COO 格式的成对源节点和目标节点索引给出。虽然
EdgeIndex是torch.Tensor的一个子类，但它可以包含额外的（元）数据，即：sparse_size: 底层稀疏矩阵的大小sort_order: 排序顺序（如果存在），可以是按行或按列。is_undirected: 边是否是双向的。
此外，
EdgeIndex缓存数据以便快速进行CSR或CSC转换，前提是其表示已排序，例如其rowptr或colptr，或者用于从CSR到CSC或反之的置换向量。 缓存是根据需求填充的（例如，当调用EdgeIndex.sort_by()时），或者通过EdgeIndex.fill_cache_()明确请求时，并且在其生命周期内维护和调整（例如，当调用EdgeIndex.flip()时）。这种表示确保了在GNN消息传递方案中的最佳计算，同时保留了基于常规COO的PyG工作流程的易用性。
from torch_geometric import EdgeIndex edge_index = EdgeIndex( [[0, 1, 1, 2], [1, 0, 2, 1]] sparse_size=(3, 3), sort_order='row', is_undirected=True, device='cpu', ) >>> EdgeIndex([[0, 1, 1, 2], ... [1, 0, 2, 1]]) assert edge_index.is_sorted_by_row assert edge_index.is_undirected # Flipping order: edge_index = edge_index.flip(0) >>> EdgeIndex([[1, 0, 2, 1], ... [0, 1, 1, 2]]) assert edge_index.is_sorted_by_col assert edge_index.is_undirected # Filtering: mask = torch.tensor([True, True, True, False]) edge_index = edge_index[:, mask] >>> EdgeIndex([[1, 0, 2], ... [0, 1, 1]]) assert edge_index.is_sorted_by_col assert not edge_index.is_undirected # Sparse-Dense Matrix Multiplication: out = edge_index.flip(0) @ torch.randn(3, 16) assert out.size() == (3, 16)
- validate() EdgeIndex[source]
验证
EdgeIndex表示。特别是，它确保了
它只包含有效的索引。
排序顺序已正确设置。
如果指定为无向，则索引是双向的。
- Return type:
- sparse_size() Tuple[Optional[int], Optional[int]][source]
- sparse_size(dim: int) Optional[int]
底层稀疏矩阵的大小。 如果指定了
dim，则返回一个整数，表示该稀疏维度的大小。
- get_sparse_size() Size[source]
- get_sparse_size(dim: int) int
底层稀疏矩阵的大小。 当未明确设置时，自动计算并缓存。 如果指定了
dim，则返回一个整数，表示该稀疏维度的大小。
- sparse_resize_(num_rows: Optional[int], num_cols: Optional[int]) EdgeIndex[source]
分配或重新分配底层稀疏矩阵的大小。
- get_csr() Tuple[Tuple[Tensor, Tensor], Optional[Tensor]][source]
返回压缩的CSR表示
(rowptr, col), perm如果EdgeIndex已排序。
- get_csc() Tuple[Tuple[Tensor, Tensor], Optional[Tensor]][source]
返回压缩的CSC表示
(colptr, row), perm如果EdgeIndex已排序。
- sort_by(sort_order: Union[str, SortOrder], stable: bool = False) SortReturnType[source]
按行或列索引对元素进行排序。
- to_dense(value: Optional[Tensor] = None, fill_value: float = 0.0, dtype: Optional[dtype] = None) Tensor[source]
将
EdgeIndex转换为密集的torch.Tensor。警告
在存在重复边的情况下，行为是不确定的（将任意选择
value中的一个值）。为了确保确定性行为，可以考虑事先调用coalesce()。- Parameters:
value (torch.Tensor, optional) – 非零元素的值。 如果未指定，非零元素将被赋值为
1.0。（默认值：None）fill_value (float, optional) – 密集矩阵中剩余元素的填充值。（默认值：
0.0）dtype (torch.dtype, optional) – 返回张量的数据类型。（默认值：
None）
- Return type:
- to_sparse_coo(value: Optional[Tensor] = None) Tensor[source]
将
EdgeIndex转换为torch.sparse_coo_tensor。- Parameters:
value (torch.Tensor, optional) – 非零元素的值。 如果未指定，非零元素将被赋值为
1.0。（默认值：None）- Return type:
- to_sparse_csr(value: Optional[Tensor] = None) Tensor[source]
将
EdgeIndex转换为torch.sparse_csr_tensor。- Parameters:
value (torch.Tensor, optional) – 非零元素的值。 如果未指定，非零元素将被赋值为
1.0。（默认值：None）- Return type:
- to_sparse_csc(value: Optional[Tensor] = None) Tensor[source]
将
EdgeIndex转换为torch.sparse_csc_tensor。- Parameters:
value (torch.Tensor, optional) – 非零元素的值。 如果未指定，非零元素将被赋值为
1.0。（默认值：None）- Return type:
- to_sparse(*, layout: layout = torch.sparse_coo, value: Optional[Tensor] = None) Tensor[source]
将
EdgeIndex转换为torch.sparse张量。- Parameters:
布局 (torch.layout, 可选) – 所需的稀疏布局。可以是
torch.sparse_coo,torch.sparse_csr, 或torch.sparse_csc。 (默认:torch.sparse_coo)value (torch.Tensor, optional) – 非零元素的值。 如果未指定，非零元素将被赋值为
1.0。（默认值：None）
- Return type:
- to_sparse_tensor(value: Optional[Tensor] = None) SparseTensor[source]
将
EdgeIndex转换为torch_sparse.SparseTensor。 需要安装torch-sparse。- Parameters:
value (torch.Tensor, optional) – 非零元素的值。 (default:
None)- Return type:
SparseTensor
- matmul(other: EdgeIndex, input_value: Optional[Tensor] = None, other_value: Optional[Tensor] = None, reduce: Literal['sum', 'mean', 'amin', 'amax', 'add', 'min', 'max'] = 'sum', transpose: bool = False) Tuple[EdgeIndex, Tensor][source]
- matmul(other: Tensor, input_value: Optional[Tensor] = None, other_value: None = None, reduce: Literal['sum', 'mean', 'amin', 'amax', 'add', 'min', 'max'] = 'sum', transpose: bool = False) Tensor
执行矩阵
input和other的矩阵乘法。 如果input是一个 \((n \times m)\) 矩阵，而other是一个 \((m \times p)\) 张量，那么输出将是一个 \((n \times p)\) 张量。 有关更多信息，请参见torch.matmul()。input是一个稀疏矩阵，由EdgeIndex中的索引表示，而input_value对应于input中非零元素的值。 如果未指定，非零元素将被赋值为1.0。other可以是一个密集的torch.Tensor或者一个稀疏的EdgeIndex。 如果other是一个稀疏的EdgeIndex，那么other_value对应于其非零元素的值。此函数还接受一个可选的
reduce参数， 该参数允许指定一个可选的归约操作。 有关更多信息，请参见torch.sparse.mm()。最后，
transpose选项允许执行矩阵乘法，其中input将首先被转置，即：\[\textrm{input}^{\top} \cdot \textrm{other}\]- Parameters:
其他 (torch.Tensor 或 EdgeIndex) – 要相乘的第二个矩阵，可以是稀疏的或密集的。
input_value (torch.Tensor, optional) –
input中非零元素的值。 如果未指定，非零元素将被赋值为1.0。（默认值：None）other_value (torch.Tensor, optional) – 如果
other是稀疏的，则为非零元素的值。 如果未指定，非零元素将被赋值为1.0。（默认值：None）reduce (str, optional) – 归约操作，可以是
"sum"/"add","mean","min"/amin或"max"/amax. (默认:"sum")transpose (bool, 可选) – 如果设置为
True，将基于转置的input执行矩阵乘法。 (默认:False)
- Return type:
- sparse_narrow(dim: int, start: Union[int, Tensor], length: int) EdgeIndex[source]
返回一个新的
EdgeIndex，它是自身的缩小版本。缩小是通过将EdgeIndex解释为形状为(num_rows, num_cols)的稀疏矩阵来执行的。与
torch.narrow()相比，返回的张量不再共享相同的底层存储。- Parameters:
dim (int) – 沿着哪个维度进行缩小。
开始 (int 或 torch.Tensor) – 从该元素的索引开始缩小维度。
length (int) – 缩小维度的长度。
- Return type: