Repository: pyg-team/pytorch_geometric
Pull request #6930 (2023-03-16 13:17:05)

Added PyTorch Sparse Tensor support for `ToSparseTensor`

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6930?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#6930](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6930?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4fd6c73) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/6587ea048c8454663c1c6cc8fafb40932202b1b9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (6587ea0) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #6930   +/-   ##
=======================================
  Coverage   91.50%   91.51%           
=======================================
  Files         429      429           
  Lines       23317    23323    +6     
=======================================
+ Hits        21337    21343    +6     
  Misses       1980     1980           
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6930?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/transforms/to\_sparse\_tensor.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6930?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvdG9fc3BhcnNlX3RlbnNvci5weQ==) | `97.56% <100.00%> (+0.41%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index aaabe8526..16fc8acd2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,7 +8,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Added the `DistMult` KGE model ([#6958](https://github.com/pyg-team/pytorch_geometric/pull/6958))
-- Added `HeteroData.set_value_dict` functionality ([#6961](https://github.com/pyg-team/pytorch_geometric/pull/6961))
+- Added `HeteroData.set_value_dict` functionality ([#6961](https://github.com/pyg-team/pytorch_geometric/pull/6961), [#6974](https://github.com/pyg-team/pytorch_geometric/pull/6974))
 - Added PyTorch >= 2.0 support ([#6934](https://github.com/pyg-team/pytorch_geometric/pull/6934))
 - Added PyTorch Lightning >= 2.0 support ([#6929](https://github.com/pyg-team/pytorch_geometric/pull/6929))
 - Added the `ComplEx` KGE model ([#6898](https://github.com/pyg-team/pytorch_geometric/pull/6898))
@@ -97,7 +97,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `Data.edge_subgraph` and `HeteroData.edge_subgraph` functionalities ([#6193](https://github.com/pyg-team/pytorch_geometric/pull/6193))
 - Added `input_time` option to `LightningNodeData` and `transform_sampler_output` to `NodeLoader` and `LinkLoader` ([#6187](https://github.com/pyg-team/pytorch_geometric/pull/6187))
 - Added `summary` for PyG/PyTorch models ([#5859](https://github.com/pyg-team/pytorch_geometric/pull/5859), [#6161](https://github.com/pyg-team/pytorch_geometric/pull/6161))
-- Started adding `torch.sparse` support to PyG ([#5906](https://github.com/pyg-team/pytorch_geometric/pull/5906), [#5944](https://github.com/pyg-team/pytorch_geometric/pull/5944), [#6003](https://github.com/pyg-team/pytorch_geometric/pull/6003), [#6033](https://github.com/pyg-team/pytorch_geometric/pull/6033), [#6514](https://github.com/pyg-team/pytorch_geometric/pull/6514), [#6532](https://github.com/pyg-team/pytorch_geometric/pull/6532), [#6748](https://github.com/pyg-team/pytorch_geometric/pull/6748), [#6847](https://github.com/pyg-team/pytorch_geometric/pull/6847), [#6868](https://github.com/pyg-team/pytorch_geometric/pull/6868), [#6874](https://github.com/pyg-team/pytorch_geometric/pull/6874), [#6897](https://github.com/pyg-team/pytorch_geometric/pull/6897), [#6932](https://github.com/pyg-team/pytorch_geometric/pull/6932), [#6936](https://github.com/pyg-team/pytorch_geometric/pull/6936), [#6937](https://github.com/pyg-team/pytorch_geometric/pull/6937), [#6939](https://github.com/pyg-team/pytorch_geometric/pull/6939), [#6947](https://github.com/pyg-team/pytorch_geometric/pull/6947), [#6950](https://github.com/pyg-team/pytorch_geometric/pull/6950), [#6951](https://github.com/pyg-team/pytorch_geometric/pull/6951), [#6957](https://github.com/pyg-team/pytorch_geometric/pull/6957))
+- Started adding `torch.sparse` support to PyG ([#5906](https://github.com/pyg-team/pytorch_geometric/pull/5906), [#5944](https://github.com/pyg-team/pytorch_geometric/pull/5944), [#6003](https://github.com/pyg-team/pytorch_geometric/pull/6003), [#6033](https://github.com/pyg-team/pytorch_geometric/pull/6033), [#6514](https://github.com/pyg-team/pytorch_geometric/pull/6514), [#6532](https://github.com/pyg-team/pytorch_geometric/pull/6532), [#6748](https://github.com/pyg-team/pytorch_geometric/pull/6748), [#6847](https://github.com/pyg-team/pytorch_geometric/pull/6847), [#6868](https://github.com/pyg-team/pytorch_geometric/pull/6868), [#6874](https://github.com/pyg-team/pytorch_geometric/pull/6874), [#6897](https://github.com/pyg-team/pytorch_geometric/pull/6897), [#6930](https://github.com/pyg-team/pytorch_geometric/pull/6930), [#6932](https://github.com/pyg-team/pytorch_geometric/pull/6932), [#6936](https://github.com/pyg-team/pytorch_geometric/pull/6936), [#6937](https://github.com/pyg-team/pytorch_geometric/pull/6937), [#6939](https://github.com/pyg-team/pytorch_geometric/pull/6939), [#6947](https://github.com/pyg-team/pytorch_geometric/pull/6947), [#6950](https://github.com/pyg-team/pytorch_geometric/pull/6950), [#6951](https://github.com/pyg-team/pytorch_geometric/pull/6951), [#6957](https://github.com/pyg-team/pytorch_geometric/pull/6957))
 - Add `inputs_channels` back in training benchmark ([#6154](https://github.com/pyg-team/pytorch_geometric/pull/6154))
 - Added support for dropping nodes in `utils.to_dense_batch` in case `max_num_nodes` is smaller than the number of nodes ([#6124](https://github.com/pyg-team/pytorch_geometric/pull/6124))
 - Added the RandLA-Net architecture as an example ([#5117](https://github.com/pyg-team/pytorch_geometric/pull/5117))
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index 4e7ccdb38..92f24f373 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -298,7 +298,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
 
            print(data['paper'].x)
         """
-        for k, v in value_dict.items():
+        for k, v in (value_dict or {}).items():
             self[k][key] = v
         return self
 
diff --git a/torch_geometric/explain/algorithm/captum_explainer.py b/torch_geometric/explain/algorithm/captum_explainer.py
index 1181a3067..64b4e6962 100644
--- a/torch_geometric/explain/algorithm/captum_explainer.py
+++ b/torch_geometric/explain/algorithm/captum_explainer.py
@@ -176,12 +176,8 @@ class CaptumExplainer(ExplainerAlgorithm):
             return Explanation(node_mask=node_mask, edge_mask=edge_mask)
 
         explanation = HeteroExplanation()
-        if node_mask is not None:
-            for node_type, mask in node_mask.items():
-                explanation.node_mask_dict[node_type] = mask
-        if edge_mask is not None:
-            for edge_type, mask in edge_mask.items():
-                explanation.edge_mask_dict[edge_type] = mask
+        explanation.set_value_dict('node_mask', node_mask)
+        explanation.set_value_dict('edge_mask', edge_mask)
         return explanation
 
     def supports(self) -> bool:
diff --git a/torch_geometric/explain/explainer.py b/torch_geometric/explain/explainer.py
index 2e9f2c5b0..1c8020b85 100644
--- a/torch_geometric/explain/explainer.py
+++ b/torch_geometric/explain/explainer.py
@@ -222,14 +222,13 @@ class Explainer:
                 explanation[key] = arg
 
         elif isinstance(explanation, HeteroExplanation):
-            assert isinstance(x, dict)
             # TODO Add `explanation._model_args`
-            for node_type, value in x.items():
-                explanation[node_type].x = value
+
+            assert isinstance(x, dict)
+            explanation.set_value_dict('x', x)
 
             assert isinstance(edge_index, dict)
-            for edge_type, value in edge_index.items():
-                explanation[edge_type].edge_index = value
+            explanation.set_value_dict('edge_index', edge_index)
 
             for key, arg in kwargs.items():  # Add remaining `kwargs`:
                 if isinstance(arg, dict):
@@ -237,8 +236,7 @@ class Explainer:
                     # while we only want to assign the `{attr_name}` to the
                     # `HeteroExplanation` object:
                     key = key[:-5] if key.endswith('_dict') else key
-                    for type_name, value in arg.items():
-                        explanation[type_name][key] = value
+                    explanation.set_value_dict(key, arg)
                 else:
                     explanation[key] = arg
 
diff --git a/torch_geometric/loader/link_loader.py b/torch_geometric/loader/link_loader.py
index b8fff41a9..88f13f368 100644
--- a/torch_geometric/loader/link_loader.py
+++ b/torch_geometric/loader/link_loader.py
@@ -237,14 +237,11 @@ class LinkLoader(torch.utils.data.DataLoader, AffinityMixin):
                 if 'n_id' not in data[key]:
                     data[key].n_id = node
 
-            if out.edge is not None:
-                for key, edge in out.edge.items():
-                    if 'e_id' not in data[key]:
-                        data[key].e_id = edge
-
-            if out.batch is not None:
-                for key, batch in out.batch.items():
-                    data[key].batch = batch
+            for key, edge in (out.edge or {}).items():
+                if 'e_id' not in data[key]:
+                    data[key].e_id = edge
+
+            data.set_value_dict('batch', out.batch)
 
             input_type = self.input_data.input_type
             data[input_type].input_id = out.metadata[0]
diff --git a/torch_geometric/loader/node_loader.py b/torch_geometric/loader/node_loader.py
index f3d9beee2..c2140293c 100644
--- a/torch_geometric/loader/node_loader.py
+++ b/torch_geometric/loader/node_loader.py
@@ -166,14 +166,9 @@ class NodeLoader(torch.utils.data.DataLoader, AffinityMixin):
                 if 'e_id' not in data[key]:
                     data[key].e_id = edge
 
-            for key, batch in (out.batch or {}).items():
-                data[key].batch = batch
-
-            for key, value in (out.num_sampled_nodes or {}).items():
-                data[key].num_sampled_nodes = value
-
-            for key, value in (out.num_sampled_edges or {}).items():
-                data[key].num_sampled_edges = value
+            data.set_value_dict('batch', out.batch)
+            data.set_value_dict('num_sampled_nodes', out.num_sampled_nodes)
+            data.set_value_dict('num_sampled_edges', out.num_sampled_edges)
 
             input_type = self.input_data.input_type
             data[input_type].input_id = out.metadata[0]
diff --git a/torch_geometric/transforms/to_sparse_tensor.py b/torch_geometric/transforms/to_sparse_tensor.py
index 030daa918..cfa7ed742 100644
--- a/torch_geometric/transforms/to_sparse_tensor.py
+++ b/torch_geometric/transforms/to_sparse_tensor.py
@@ -1,42 +1,72 @@
 from typing import Optional, Union
 
+import torch
+from torch import Tensor
+
+import torch_geometric.typing
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
 from torch_geometric.typing import SparseTensor
-from torch_geometric.utils import sort_edge_index
+from torch_geometric.utils import (
+    sort_edge_index,
+    to_torch_coo_tensor,
+    to_torch_csr_tensor,
+)
 
 
 @functional_transform('to_sparse_tensor')
 class ToSparseTensor(BaseTransform):
     r"""Converts the :obj:`edge_index` attributes of a homogeneous or
-    heterogeneous data object into a (transposed)
-    :class:`torch_sparse.SparseTensor` object with key :obj:`adj_t`
+    heterogeneous data object into a **transposed**
+    :class:`torch_sparse.SparseTensor` or :pytorch:`PyTorch`
+    :class:`torch.sparse.Tensor` object with key :obj:`adj_t`
     (functional name: :obj:`to_sparse_tensor`).
 
     .. note::
 
         In case of composing multiple transforms, it is best to convert the
-        :obj:`data` object to a :class:`~torch_sparse.SparseTensor` as late as
-        possible, since there exist some transforms that are only able to
-        operate on :obj:`data.edge_index` for now.
+        :obj:`data` object via :class:`ToSparseTensor` as late as possible,
+        since there exist some transforms that are only able to operate on
+        :obj:`data.edge_index` for now.
 
     Args:
         attr (str, optional): The name of the attribute to add as a value to
-            the :class:`~torch_sparse.SparseTensor` object (if present).
+            the :class:`~torch_sparse.SparseTensor` or
+            :class:`torch.sparse.Tensor` object (if present).
             (default: :obj:`edge_weight`)
         remove_edge_index (bool, optional): If set to :obj:`False`, the
             :obj:`edge_index` tensor will not be removed.
             (default: :obj:`True`)
-        fill_cache (bool, optional): If set to :obj:`False`, will not fill the
-            underlying :class:`~torch_sparse.SparseTensor` cache.
+        fill_cache (bool, optional): If set to :obj:`True`, will fill the
+            underlying :class:`torch_sparse.SparseTensor` cache (if used).
             (default: :obj:`True`)
+        layout (torch.layout, optional): Specifies the layout of the returned
+            sparse tensor (:obj:`None`, :obj:`torch.sparse_coo` or
+            :obj:`torch.sparse_csr`).
+            If set to :obj:`None` and the :obj:`torch_sparse` dependency is
+            installed, will convert :obj:`edge_index` into a
+            :class:`torch_sparse.SparseTensor` object.
+            If set to :obj:`None` and the :obj:`torch_sparse` dependency is
+            not installed, will convert :obj:`edge_index` into a
+            :class:`torch.sparse.Tensor` object with layout
+            :obj:`torch.sparse_csr`. (default: :obj:`None`)
     """
-    def __init__(self, attr: Optional[str] = 'edge_weight',
-                 remove_edge_index: bool = True, fill_cache: bool = True):
+    def __init__(
+        self,
+        attr: Optional[str] = 'edge_weight',
+        remove_edge_index: bool = True,
+        fill_cache: bool = True,
+        layout: Optional[int] = None,
+    ):
+        if layout not in {None, torch.sparse_coo, torch.sparse_csr}:
+            raise ValueError(f"Unexpected sparse tensor layout "
+                             f"(got '{layout}')")
+
         self.attr = attr
         self.remove_edge_index = remove_edge_index
         self.fill_cache = fill_cache
+        self.layout = layout
 
     def __call__(
         self,
@@ -65,19 +95,50 @@ class ToSparseTensor(BaseTransform):
             for key, value in zip(keys, values):
                 store[key] = value
 
-            store.adj_t = SparseTensor(
-                row=store.edge_index[1], col=store.edge_index[0],
-                value=None if self.attr is None or self.attr not in store else
-                store[self.attr], sparse_sizes=store.size()[::-1],
-                is_sorted=True, trust_data=True)
+            layout = self.layout
+            size = store.size()[::-1]
+            value: Optional[Tensor] = None
+            if self.attr is not None and self.attr in store:
+                value = store[self.attr]
+
+            if layout is None and torch_geometric.typing.WITH_TORCH_SPARSE:
+                store.adj_t = SparseTensor(
+                    row=store.edge_index[1],
+                    col=store.edge_index[0],
+                    value=value,
+                    sparse_sizes=size,
+                    is_sorted=True,
+                    trust_data=True,
+                )
+
+            # TODO Multi-dimensional edge attributes only supported for COO.
+            elif ((value is not None and value.dim() > 1)
+                  or layout == torch.sparse_coo):
+                store.adj_t = to_torch_coo_tensor(
+                    store.edge_index.flip([0]),
+                    edge_attr=value,
+                    size=size,
+                )
+
+            elif layout is None or layout == torch.sparse_csr:
+                store.adj_t = to_torch_csr_tensor(
+                    store.edge_index.flip([0]),
+                    edge_attr=value,
+                    size=size,
+                )
 
             if self.remove_edge_index:
                 del store['edge_index']
                 if self.attr is not None and self.attr in store:
                     del store[self.attr]
 
-            if self.fill_cache:  # Pre-process some important attributes.
+            if self.fill_cache and isinstance(store.adj_t, SparseTensor):
+                # Pre-process some important attributes.
                 store.adj_t.storage.rowptr()
                 store.adj_t.storage.csr2csc()
 
         return data
+
+    def __repr__(self):
+        return (f'{self.__class__.__name__}(attr={self.attr}, '
+                f'layout={self.layout})')


Test patch:
diff --git a/test/transforms/test_to_sparse_tensor.py b/test/transforms/test_to_sparse_tensor.py
index 71d4cc1fd..e0ba20b13 100644
--- a/test/transforms/test_to_sparse_tensor.py
+++ b/test/transforms/test_to_sparse_tensor.py
@@ -1,12 +1,15 @@
+import pytest
 import torch
 
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.transforms import ToSparseTensor
 
 
-def test_to_sparse_tensor():
-    transform = ToSparseTensor()
-    assert str(transform) == 'ToSparseTensor()'
+@pytest.mark.parametrize('layout', [None, torch.sparse_coo, torch.sparse_csr])
+def test_to_sparse_tensor_basic(layout):
+    transform = ToSparseTensor(layout=layout)
+    assert str(transform) == (f'ToSparseTensor(attr=edge_weight, '
+                              f'layout={layout})')
 
     edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
     edge_weight = torch.randn(edge_index.size(1))
@@ -17,12 +20,24 @@ def test_to_sparse_tensor():
     data = Data(edge_index=edge_index, edge_weight=edge_weight,
                 edge_attr=edge_attr, num_nodes=3)
     data = transform(data)
+
     assert len(data) == 3
-    assert data.adj_t.storage.row().tolist() == [0, 1, 1, 2]
-    assert data.adj_t.storage.col().tolist() == [1, 0, 2, 1]
-    assert data.adj_t.storage.value().tolist() == edge_weight[perm].tolist()
-    assert data.edge_attr.tolist() == edge_attr[perm].tolist()
     assert data.num_nodes == 3
+    assert torch.equal(data.edge_attr, edge_attr[perm])
+    assert 'adj_t' in data
+
+    if layout is None:  # `torch_sparse.SparseTensor`.
+        row, col, value = data.adj_t.coo()
+        assert row.tolist() == [0, 1, 1, 2]
+        assert col.tolist() == [1, 0, 2, 1]
+        assert torch.equal(value, edge_weight[perm])
+    else:
+        adj_t = data.adj_t
+        assert adj_t.layout == layout
+        if layout != torch.sparse_coo:
+            adj_t = adj_t.to_sparse_coo()
+        assert adj_t.indices().tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+        assert torch.equal(adj_t.values(), edge_weight[perm])
 
 
 def test_to_sparse_tensor_and_keep_edge_index():
@@ -35,17 +50,15 @@ def test_to_sparse_tensor_and_keep_edge_index():
     data = Data(edge_index=edge_index, edge_weight=edge_weight,
                 edge_attr=edge_attr, num_nodes=3)
     data = ToSparseTensor(remove_edge_index=False)(data)
+
     assert len(data) == 5
-    assert data.adj_t.storage.row().tolist() == [0, 1, 1, 2]
-    assert data.adj_t.storage.col().tolist() == [1, 0, 2, 1]
-    assert data.adj_t.storage.value().tolist() == edge_weight[perm].tolist()
-    assert data.edge_index.tolist() == edge_index[:, perm].tolist()
-    assert data.edge_weight.tolist() == edge_weight[perm].tolist()
-    assert data.edge_attr.tolist() == edge_attr[perm].tolist()
-    assert data.num_nodes == 3
+    assert torch.equal(data.edge_index, edge_index[:, perm])
+    assert torch.equal(data.edge_weight, edge_weight[perm])
+    assert torch.equal(data.edge_attr, edge_attr[perm])
 
 
-def test_hetero_to_sparse_tensor():
+@pytest.mark.parametrize('layout', [None, torch.sparse_coo, torch.sparse_csr])
+def test_hetero_to_sparse_tensor(layout):
     edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
 
     data = HeteroData()
@@ -53,13 +66,33 @@ def test_hetero_to_sparse_tensor():
     data['w'].num_nodes = 3
     data['v', 'v'].edge_index = edge_index
     data['v', 'w'].edge_index = edge_index
-    data = ToSparseTensor()(data)
-    assert data['v', 'v'].adj_t.storage.row().tolist() == [0, 1, 1, 2]
-    assert data['v', 'v'].adj_t.storage.col().tolist() == [1, 0, 2, 1]
-    assert data['v', 'v'].adj_t.storage.value() is None
-    assert data['v', 'w'].adj_t.storage.row().tolist() == [0, 1, 1, 2]
-    assert data['v', 'w'].adj_t.storage.col().tolist() == [1, 0, 2, 1]
-    assert data['v', 'w'].adj_t.storage.value() is None
+
+    data = ToSparseTensor(layout=layout)(data)
+
+    if layout is None:  # `torch_sparse.SparseTensor`.
+        row, col, value = data['v', 'v'].adj_t.coo()
+        assert row.tolist() == [0, 1, 1, 2]
+        assert col.tolist() == [1, 0, 2, 1]
+        assert value is None
+
+        row, col, value = data['v', 'w'].adj_t.coo()
+        assert row.tolist() == [0, 1, 1, 2]
+        assert col.tolist() == [1, 0, 2, 1]
+        assert value is None
+    else:
+        adj_t = data['v', 'v'].adj_t
+        assert adj_t.layout == layout
+        if layout != torch.sparse_coo:
+            adj_t = adj_t.to_sparse_coo()
+        assert adj_t.indices().tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+        assert adj_t.values().tolist() == [1., 1., 1., 1.]
+
+        adj_t = data['v', 'w'].adj_t
+        assert adj_t.layout == layout
+        if layout != torch.sparse_coo:
+            adj_t = adj_t.to_sparse_coo()
+        assert adj_t.indices().tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+        assert adj_t.values().tolist() == [1., 1., 1., 1.]
 
 
 def test_to_sparse_tensor_num_nodes_equals_num_edges():
@@ -74,10 +107,8 @@ def test_to_sparse_tensor_num_nodes_equals_num_edges():
     data = Data(x=x, edge_index=edge_index, edge_weight=edge_weight,
                 edge_attr=edge_attr, y=y)
     data = ToSparseTensor()(data)
+
     assert len(data) == 4
-    assert data.x.tolist() == [0, 1, 2, 3]
-    assert data.adj_t.storage.row().tolist() == [0, 1, 1, 2]
-    assert data.adj_t.storage.col().tolist() == [1, 0, 2, 1]
-    assert data.adj_t.storage.value().tolist() == edge_weight[perm].tolist()
-    assert data.edge_attr.tolist() == edge_attr[perm].tolist()
-    assert data.y.tolist() == [0, 1, 2, 3]
+    assert torch.equal(data.x, x)
+    assert torch.equal(data.y, y)
+    assert torch.equal(data.edge_attr, edge_attr[perm])