torch_geometric.transforms.ToSparseTensor — pytorch_geometric documentation (original) (raw)

class ToSparseTensor(attr: Optional[str] = 'edge_weight', remove_edge_index: bool = True, fill_cache: bool = True, layout: Optional[int] = None)[source]

Bases: BaseTransform

Converts the edge_index attributes of a homogeneous or heterogeneous data object into a transposed torch_sparse.SparseTensor or PyTorch torch.sparse.Tensor object with key adj_t(functional name: to_sparse_tensor).

Note

In case of composing multiple transforms, it is best to convert thedata object via ToSparseTensor as late as possible, since there exist some transforms that are only able to operate ondata.edge_index for now.

Parameters:

• attr (str, optional) – The name of the attribute to add as a value to the SparseTensor ortorch.sparse.Tensor object (if present). (default: edge_weight)
• remove_edge_index (bool, optional) – If set to False, theedge_index tensor will not be removed. (default: True)
• fill_cache (bool, optional) – If set to True, will fill the underlying torch_sparse.SparseTensor cache (if used). (default: True)
• layout (torch.layout, optional) – Specifies the layout of the returned sparse tensor (None, torch.sparse_coo ortorch.sparse_csr). If set to None and the torch_sparse dependency is installed, will convert edge_index into atorch_sparse.SparseTensor object. If set to None and the torch_sparse dependency is not installed, will convert edge_index into atorch.sparse.Tensor object with layouttorch.sparse_csr. (default: None)