torch_geometric.transforms.ToSparseTensor

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 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 the data object via ToSparseTensor as late as possible, since there exist some transforms that are only able to operate on data.edge_index for now.

Parameters

  • attr (str, optional) – The name of the attribute to add as a value to the SparseTensor or torch.sparse.Tensor object (if present). (default: edge_weight)

  • remove_edge_index (bool, optional) – If set to False, the edge_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 or torch.sparse_csr). If set to None and the torch_sparse dependency is installed, will convert edge_index into a torch_sparse.SparseTensor object. If set to None and the torch_sparse dependency is not installed, will convert edge_index into a torch.sparse.Tensor object with layout torch.sparse_csr. (default: None)