Repository: pyg-team/pytorch_geometric
Pull request #7482 (2023-06-01 03:00:21)

Add `LocalFeatureStore` and `LocalGraphStore` helper initializations
This code belongs to the part of the whole distributed training for PyG.

This class will initialize the graph & node/edge feature data based on LocalGraphStore/LocalFeatureStore data structure. You can based on these initialize API to define your graph and node/edge features for your homo/hetero datasets.

We also include the unit test under /test/distributed folder to show how one OGB_MAG dataset will be initialize into the graph and feature by LocalDataset APIs.

Any comments please let us know. thanks

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7482?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7482](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7482?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (8662ea6) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/db460561919ca384ef6b10f0ed3b7a8a6dab95e2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (db46056) will **decrease** coverage by `0.27%`.
> The diff coverage is `51.78%`.

```diff
@@            Coverage Diff             @@
##           master    #7482      +/-   ##
==========================================
- Coverage   91.47%   91.20%   -0.27%     
==========================================
  Files         446      448       +2     
  Lines       24778    24945     +167     
==========================================
+ Hits        22666    22752      +86     
- Misses       2112     2193      +81     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7482?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/distributed/local\_dataset.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7482?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL2xvY2FsX2RhdGFzZXQucHk=) | `49.54% <49.54%> (ø)` | |
| [torch\_geometric/utils/tensor.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7482?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3RlbnNvci5weQ==) | `50.00% <50.00%> (ø)` | |
| [torch\_geometric/distributed/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7482?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/typing.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7482?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3R5cGluZy5weQ==) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/utils/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7482?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/utils/cross\_entropy.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7482?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2Nyb3NzX2VudHJvcHkucHk=) | `100.00% <100.00%> (ø)` | |

: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 ae0f0b87c..7e73babfe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,7 +17,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added a CPU-based and GPU-based `map_index` implementation ([#7493](https://github.com/pyg-team/pytorch_geometric/pull/7493))
 - Added the `AmazonBook` heterogeneous dataset ([#7483](https://github.com/pyg-team/pytorch_geometric/pull/7483))
 - Added hierarichial heterogeneous GraphSAGE example on OGB-MAG ([#7425](https://github.com/pyg-team/pytorch_geometric/pull/7425))
-- Added the `torch_geometric.distributed` package ([#7451](https://github.com/pyg-team/pytorch_geometric/pull/7451), [#7452](https://github.com/pyg-team/pytorch_geometric/pull/7452))
+- Added the `torch_geometric.distributed` package ([#7451](https://github.com/pyg-team/pytorch_geometric/pull/7451), [#7452](https://github.com/pyg-team/pytorch_geometric/pull/7452)), [#7482](https://github.com/pyg-team/pytorch_geometric/pull/7482))
 - Added the `GDELTLite` dataset ([#7442](https://github.com/pyg-team/pytorch_geometric/pull/7442))
 - Added the `approx_knn` function for approximated nearest neighbor search ([#7421](https://github.com/pyg-team/pytorch_geometric/pull/7421))
 - Added the `IGMCDataset` ([#7441](https://github.com/pyg-team/pytorch_geometric/pull/7441))
diff --git a/torch_geometric/distributed/local_feature_store.py b/torch_geometric/distributed/local_feature_store.py
index 5abbfa2f1..76442c005 100644
--- a/torch_geometric/distributed/local_feature_store.py
+++ b/torch_geometric/distributed/local_feature_store.py
@@ -1,12 +1,13 @@
 import copy
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
 
 from torch_geometric.data import FeatureStore, TensorAttr
 from torch_geometric.data.feature_store import _field_status
+from torch_geometric.typing import EdgeType, NodeType
 
 
 @dataclass
@@ -14,7 +15,7 @@ class LocalTensorAttr(TensorAttr):
     r"""Tensor attribute for storing features without :obj:`index`."""
     def __init__(
         self,
-        group_name: Optional[str] = _field_status.UNSET,
+        group_name: Optional[Union[NodeType, EdgeType]] = _field_status.UNSET,
         attr_name: Optional[str] = _field_status.UNSET,
         index=None,
     ):
@@ -27,35 +28,38 @@ class LocalFeatureStore(FeatureStore):
     def __init__(self):
         super().__init__(tensor_attr_cls=LocalTensorAttr)
 
-        self._feat: Dict[Tuple[str, str], Tensor] = {}
+        self._feat: Dict[Tuple[Union[NodeType, EdgeType], str], Tensor] = {}
 
         # Save the global node/edge IDs:
-        self._global_id: Dict[Tuple[str, str], Tensor] = {}
+        self._global_id: Dict[Union[NodeType, EdgeType], Tensor] = {}
 
         # Save the mapping from global node/edge IDs to indices in `_feat`:
-        self._global_id_to_index: Dict[Tuple[str, str], Tensor] = {}
+        self._global_id_to_index: Dict[Union[NodeType, EdgeType], Tensor] = {}
 
     @staticmethod
     def key(attr: TensorAttr) -> Tuple[str, str]:
         return (attr.group_name, attr.attr_name)
 
-    def put_global_id(self, global_id: Tensor, *args, **kwargs) -> bool:
-        attr = self._tensor_attr_cls.cast(*args, **kwargs)
-        self._global_id[self.key(attr)] = global_id
-        self._set_global_id_to_index(attr)
+    def put_global_id(
+        self,
+        global_id: Tensor,
+        group_name: Union[NodeType, EdgeType],
+    ) -> bool:
+        self._global_id[group_name] = global_id
+        self._set_global_id_to_index(group_name)
         return True
 
-    def get_global_id(self, *args, **kwargs) -> Optional[Tensor]:
-        attr = self._tensor_attr_cls.cast(*args, **kwargs)
-        return self._global_id.get(self.key(attr))
+    def get_global_id(
+        self,
+        group_name: Union[NodeType, EdgeType],
+    ) -> Optional[Tensor]:
+        return self._global_id.get(group_name)
 
-    def remove_global_id(self, *args, **kwargs) -> bool:
-        attr = self._tensor_attr_cls.cast(*args, **kwargs)
-        return self._global_id.pop(self.key(attr), None) is not None
+    def remove_global_id(self, group_name: Union[NodeType, EdgeType]) -> bool:
+        return self._global_id.pop(group_name) is not None
 
-    def _set_global_id_to_index(self, *args, **kwargs):
-        attr = self._tensor_attr_cls.cast(*args, **kwargs)
-        global_id = self.get_global_id(attr)
+    def _set_global_id_to_index(self, group_name: Union[NodeType, EdgeType]):
+        global_id = self.get_global_id(group_name)
 
         if global_id is None:
             return
@@ -64,7 +68,7 @@ class LocalFeatureStore(FeatureStore):
         global_id_to_index = global_id.new_full((int(global_id.max()) + 1, ),
                                                 fill_value=-1)
         global_id_to_index[global_id] = torch.arange(global_id.numel())
-        self._global_id_to_index[self.key(attr)] = global_id_to_index
+        self._global_id_to_index[group_name] = global_id_to_index
 
     def _put_tensor(self, tensor: Tensor, attr: TensorAttr) -> bool:
         assert attr.index is None
@@ -91,7 +95,7 @@ class LocalFeatureStore(FeatureStore):
         assert attr.index is not None
 
         attr = copy.copy(attr)
-        attr.index = self._global_id_to_index[self.key(attr)][attr.index]
+        attr.index = self._global_id_to_index[attr.group_name][attr.index]
 
         return self.get_tensor(attr)
 
@@ -100,3 +104,89 @@ class LocalFeatureStore(FeatureStore):
 
     def get_all_tensor_attrs(self) -> List[LocalTensorAttr]:
         return [self._tensor_attr_cls.cast(*key) for key in self._feat.keys()]
+
+    # Initialization ##########################################################
+
+    @classmethod
+    def from_data(
+        cls,
+        node_id: Tensor,
+        x: Optional[Tensor] = None,
+        y: Optional[Tensor] = None,
+        edge_id: Optional[Tensor] = None,
+        edge_attr: Optional[Tensor] = None,
+    ) -> 'LocalFeatureStore':
+        r"""Creates a local feature store from homogeneous :pyg:`PyG` tensors.
+
+        Args:
+            node_id (torch.Tensor): The global identifier for every local node.
+            x (torch.Tensor, optional): The node features.
+                (default: :obj:`None`)
+            y (torch.Tensor, optional): The node labels. (default: :obj:`None`)
+            edge_id (torch.Tensor, optional): The global identifier for every
+                local edge. (default: :obj:`None`)
+            edge_attr (torch.Tensor, optional): The edge features.
+                (default: :obj:`None`)
+        """
+        feat_store = cls()
+        feat_store.put_global_id(node_id, group_name=None)
+        if x is not None:
+            feat_store.put_tensor(x, group_name=None, attr_name='x')
+        if y is not None:
+            feat_store.put_tensor(y, group_name=None, attr_name='y')
+        if edge_id is not None:
+            feat_store.put_global_id(edge_id, group_name=(None, None))
+        if edge_attr is not None:
+            if edge_id is None:
+                raise ValueError("'edge_id' needs to be present in case "
+                                 "'edge_attr' is passed")
+            feat_store.put_tensor(edge_attr, group_name=(None, None),
+                                  attr_name='edge_attr')
+        return feat_store
+
+    @classmethod
+    def from_hetero_data(
+        cls,
+        node_id_dict: Dict[NodeType, Tensor],
+        x_dict: Optional[Dict[NodeType, Tensor]] = None,
+        y_dict: Optional[Dict[NodeType, Tensor]] = None,
+        edge_id_dict: Optional[Dict[EdgeType, Tensor]] = None,
+        edge_attr_dict: Optional[Dict[EdgeType, Tensor]] = None,
+    ) -> 'LocalFeatureStore':
+        r"""Creates a local graph store from heterogeneous :pyg:`PyG` tensors.
+
+        Args:
+            node_id_dict (Dict[NodeType, torch.Tensor]): The global identifier
+                for every local node of every node type.
+            x_dict (Dict[NodeType, torch.Tensor], optional): The node features
+                of node types. (default: :obj:`None`)
+            y_dict (Dict[NodeType, torch.Tensor], optional): The node labels of
+                node types. (default: :obj:`None`)
+            edge_id_dict (Dict[EdgeType, torch.Tensor], optional): The global
+                identifier for every local edge of edge types.
+                (default: :obj:`None`)
+            edge_attr_dict (Dict[EdgeType, torch.Tensor], optional): The edge
+                features of edge types. (default: :obj:`None`)
+        """
+        feat_store = cls()
+
+        for node_type, node_id in node_id_dict.items():
+            feat_store.put_global_id(node_id, group_name=node_type)
+        if x_dict is not None:
+            for node_type, x in x_dict.items():
+                feat_store.put_tensor(x, group_name=node_type, attr_name='x')
+        if y_dict is not None:
+            for node_type, y in y_dict.items():
+                feat_store.put_tensor(y, group_name=node_type, attr_name='y')
+        if edge_id_dict is not None:
+            for edge_type, edge_id in edge_id_dict.items():
+                feat_store.put_global_id(edge_id, group_name=edge_type)
+        if edge_attr_dict is not None:
+            for edge_type, edge_attr in edge_attr_dict.items():
+                if edge_id_dict is None or edge_type not in edge_id_dict:
+                    raise ValueError("'edge_id' needs to be present in case "
+                                     "'edge_attr' is passed")
+                feat_store.put_tensor(edge_attr, group_name=edge_type,
+                                      attr_name='edge_attr')
+
+        return feat_store
diff --git a/torch_geometric/distributed/local_graph_store.py b/torch_geometric/distributed/local_graph_store.py
index d83e301a1..ecdaf7dfd 100644
--- a/torch_geometric/distributed/local_graph_store.py
+++ b/torch_geometric/distributed/local_graph_store.py
@@ -3,7 +3,7 @@ from typing import Dict, List, Optional, Tuple
 from torch import Tensor
 
 from torch_geometric.data import EdgeAttr, GraphStore
-from torch_geometric.typing import EdgeTensorType
+from torch_geometric.typing import EdgeTensorType, EdgeType, NodeType
 
 
 class LocalGraphStore(GraphStore):
@@ -12,11 +12,12 @@ class LocalGraphStore(GraphStore):
     def __init__(self):
         super().__init__()
         self._edge_index: Dict[Tuple, EdgeTensorType] = {}
+        self._edge_attr: Dict[Tuple, EdgeAttr] = {}
         self._edge_id: Dict[Tuple, Tensor] = {}
 
     @staticmethod
     def key(attr: EdgeAttr) -> Tuple:
-        return (attr.edge_type, attr.layout.value, attr.is_sorted, attr.size)
+        return (attr.edge_type, attr.layout.value)
 
     def put_edge_id(self, edge_id: Tensor, *args, **kwargs) -> bool:
         edge_attr = self._edge_attr_cls.cast(*args, **kwargs)
@@ -34,13 +35,76 @@ class LocalGraphStore(GraphStore):
     def _put_edge_index(self, edge_index: EdgeTensorType,
                         edge_attr: EdgeAttr) -> bool:
         self._edge_index[self.key(edge_attr)] = edge_index
+        self._edge_attr[self.key(edge_attr)] = edge_attr
         return True
 
     def _get_edge_index(self, edge_attr: EdgeAttr) -> Optional[EdgeTensorType]:
         return self._edge_index.get(self.key(edge_attr), None)
 
     def _remove_edge_index(self, edge_attr: EdgeAttr) -> bool:
+        self._edge_attr.pop(self.key(edge_attr), None)
         return self._edge_index.pop(self.key(edge_attr), None) is not None
 
     def get_all_edge_attrs(self) -> List[EdgeAttr]:
-        return [EdgeAttr(*key) for key in self._edge_index.keys()]
+        return [self._edge_attr[key] for key in self._edge_index.keys()]
+
+    # Initialization ##########################################################
+
+    @classmethod
+    def from_data(
+        cls,
+        edge_id: Tensor,
+        edge_index: Tensor,
+        num_nodes: int,
+    ) -> 'LocalGraphStore':
+        r"""Creates a local graph store from a homogeneous :pyg:`PyG` graph.
+
+        Args:
+            edge_id (torch.Tensor): The global identifier for every local edge.
+            edge_index (torch.Tensor): The local edge indices.
+            num_nodes (int): The number of nodes in the local graph.
+        """
+        attr = dict(
+            edge_type=None,
+            layout='coo',
+            size=(num_nodes, num_nodes),
+        )
+
+        graph_store = cls()
+        graph_store.put_edge_index(edge_index, **attr)
+        graph_store.put_edge_id(edge_id, **attr)
+        return graph_store
+
+    @classmethod
+    def from_hetero_data(
+        cls,
+        edge_id_dict: Dict[EdgeType, Tensor],
+        edge_index_dict: Dict[EdgeType, Tensor],
+        num_nodes_dict: Dict[NodeType, int],
+    ) -> 'LocalGraphStore':
+        r"""Creates a local graph store from a heterogeneous :pyg:`PyG` graph.
+
+        Args:
+            edge_id_dict (Dict[EdgeType, torch.Tensor]): The global identifier
+                for every local edge of every edge type.
+            edge_index_dict (Dict[EdgeType, torch.Tensor]): The local edge
+                indices of every edge type.
+            num_nodes_dict (Dict[NodeType, int]): The number of nodes in the
+                local graph of every node type.
+        """
+        attr_dict = {}
+        for edge_type in edge_index_dict.keys():
+            src, _, dst = edge_type
+            attr_dict[edge_type] = dict(
+                edge_type=edge_type,
+                layout='coo',
+                size=(num_nodes_dict[src], num_nodes_dict[dst]),
+            )
+
+        graph_store = cls()
+        for edge_type, edge_index in edge_index_dict.items():
+            attr = attr_dict[edge_type]
+            edge_id = edge_id_dict[edge_type]
+            graph_store.put_edge_index(edge_index, **attr)
+            graph_store.put_edge_id(edge_id, **attr)
+        return graph_store


Test patch:
diff --git a/test/distributed/test_local_feature_store.py b/test/distributed/test_local_feature_store.py
index 1ccb15b3e..557ca91ca 100644
--- a/test/distributed/test_local_feature_store.py
+++ b/test/distributed/test_local_feature_store.py
@@ -18,15 +18,14 @@ def test_local_feature_store_global_id():
         [8, 8, 8],
     ])
 
-    kwargs = dict(group_name='part1', attr_name='feat')
-    part1_global_id = torch.tensor([1, 2, 3, 5, 8, 4])
-    part1_feat = feat[part1_global_id]
+    paper_global_id = torch.tensor([1, 2, 3, 5, 8, 4])
+    paper_feat = feat[paper_global_id]
 
-    store.put_global_id(part1_global_id, **kwargs)
-    store.put_tensor(part1_feat, **kwargs)
+    store.put_global_id(paper_global_id, group_name='paper')
+    store.put_tensor(paper_feat, group_name='paper', attr_name='feat')
 
-    out = store.get_tensor_from_global_id(index=torch.tensor([3, 8, 4]),
-                                          **kwargs)
+    out = store.get_tensor_from_global_id(group_name='paper', attr_name='feat',
+                                          index=torch.tensor([3, 8, 4]))
     assert torch.equal(out, feat[torch.tensor([3, 8, 4])])
 
 
@@ -45,15 +44,92 @@ def test_local_feature_store_utils():
         [8, 8, 8],
     ])
 
-    kwargs = dict(group_name='part1', attr_name='feat')
-    part1_global_id = torch.tensor([1, 2, 3, 5, 8, 4])
-    part1_feat = feat[part1_global_id]
+    paper_global_id = torch.tensor([1, 2, 3, 5, 8, 4])
+    paper_feat = feat[paper_global_id]
 
-    store.put_tensor(part1_feat, **kwargs)
+    store.put_tensor(paper_feat, group_name='paper', attr_name='feat')
 
     assert len(store.get_all_tensor_attrs()) == 1
     attr = store.get_all_tensor_attrs()[0]
-    assert attr.group_name == 'part1'
+    assert attr.group_name == 'paper'
     assert attr.attr_name == 'feat'
     assert attr.index is None
     assert store.get_tensor_size(attr) == (6, 3)
+
+
+def test_homogeneous_feature_store():
+    node_id = torch.randperm(6)
+    x = torch.randn(6, 32)
+    y = torch.randint(0, 2, (6, ))
+    edge_id = torch.randperm(12)
+    edge_attr = torch.randn(12, 16)
+
+    store = LocalFeatureStore.from_data(node_id, x, y, edge_id, edge_attr)
+
+    assert len(store.get_all_tensor_attrs()) == 3
+    attrs = store.get_all_tensor_attrs()
+
+    assert attrs[0].group_name is None
+    assert attrs[0].attr_name == 'x'
+    assert attrs[1].group_name is None
+    assert attrs[1].attr_name == 'y'
+    assert attrs[2].group_name == (None, None)
+    assert attrs[2].attr_name == 'edge_attr'
+
+    assert torch.equal(store.get_global_id(group_name=None), node_id)
+    assert torch.equal(store.get_tensor(group_name=None, attr_name='x'), x)
+    assert torch.equal(store.get_tensor(group_name=None, attr_name='y'), y)
+    assert torch.equal(store.get_global_id(group_name=(None, None)), edge_id)
+    assert torch.equal(
+        store.get_tensor(group_name=(None, None), attr_name='edge_attr'),
+        edge_attr,
+    )
+
+
+def test_heterogeneous_feature_store():
+    node_type = 'paper'
+    edge_type = ('paper', 'to', 'paper')
+    node_id_dict = {node_type: torch.randperm(6)}
+    x_dict = {node_type: torch.randn(6, 32)}
+    y_dict = {node_type: torch.randint(0, 2, (6, ))}
+    edge_id_dict = {edge_type: torch.randperm(12)}
+    edge_attr_dict = {edge_type: torch.randn(12, 16)}
+
+    store = LocalFeatureStore.from_hetero_data(
+        node_id_dict,
+        x_dict,
+        y_dict,
+        edge_id_dict,
+        edge_attr_dict,
+    )
+
+    assert len(store.get_all_tensor_attrs()) == 3
+    attrs = store.get_all_tensor_attrs()
+
+    assert attrs[0].group_name == node_type
+    assert attrs[0].attr_name == 'x'
+    assert attrs[1].group_name == node_type
+    assert attrs[1].attr_name == 'y'
+    assert attrs[2].group_name == edge_type
+    assert attrs[2].attr_name == 'edge_attr'
+
+    assert torch.equal(
+        store.get_global_id(group_name=node_type),
+        node_id_dict[node_type],
+    )
+    assert torch.equal(
+        store.get_tensor(group_name=node_type, attr_name='x'),
+        x_dict[node_type],
+    )
+    assert torch.equal(
+        store.get_tensor(group_name=node_type, attr_name='y'),
+        y_dict[node_type],
+    )
+    assert torch.equal(
+        store.get_global_id(group_name=edge_type),
+        edge_id_dict[edge_type],
+    )
+    assert torch.equal(
+        store.get_tensor(group_name=edge_type, attr_name='edge_attr'),
+        edge_attr_dict[edge_type],
+    )
diff --git a/test/distributed/test_local_graph_store.py b/test/distributed/test_local_graph_store.py
index 2047db8ef..ce83082af 100644
--- a/test/distributed/test_local_graph_store.py
+++ b/test/distributed/test_local_graph_store.py
@@ -8,7 +8,7 @@ def test_local_graph_store():
     graph_store = LocalGraphStore()
 
     edge_index = get_random_edge_index(100, 100, 300)
-    edge_id = torch.tensor([1, 2, 3, 5, 8, 4], dtype=torch.int64)
+    edge_id = torch.tensor([1, 2, 3, 5, 8, 4])
 
     graph_store.put_edge_index(edge_index, edge_type=None, layout='coo',
                                size=(100, 100))
@@ -24,3 +24,51 @@ def test_local_graph_store():
     graph_store.remove_edge_index(edge_attr)
     graph_store.remove_edge_id(edge_attr)
     assert len(graph_store.get_all_edge_attrs()) == 0
+
+
+def test_homogeneous_graph_store():
+    edge_id = torch.randperm(300)
+    edge_index = get_random_edge_index(100, 100, 300)
+
+    graph_store = LocalGraphStore.from_data(edge_id, edge_index, num_nodes=100)
+
+    assert len(graph_store.get_all_edge_attrs()) == 1
+    edge_attr = graph_store.get_all_edge_attrs()[0]
+    assert edge_attr.edge_type is None
+    assert edge_attr.layout.value == 'coo'
+    assert not edge_attr.is_sorted
+    assert edge_attr.size == (100, 100)
+
+    assert torch.equal(
+        graph_store.get_edge_id(edge_type=None, layout='coo'),
+        edge_id,
+    )
+    assert torch.equal(
+        graph_store.get_edge_index(edge_type=None, layout='coo'),
+        edge_index,
+    )
+
+
+def test_heterogeneous_graph_store():
+    edge_type = ('paper', 'to', 'paper')
+    edge_id_dict = {edge_type: torch.randperm(300)}
+    edge_index_dict = {edge_type: get_random_edge_index(100, 100, 300)}
+
+    graph_store = LocalGraphStore.from_hetero_data(
+        edge_id_dict, edge_index_dict, num_nodes_dict={'paper': 100})
+
+    assert len(graph_store.get_all_edge_attrs()) == 1
+    edge_attr = graph_store.get_all_edge_attrs()[0]
+    assert edge_attr.edge_type == edge_type
+    assert edge_attr.layout.value == 'coo'
+    assert not edge_attr.is_sorted
+    assert edge_attr.size == (100, 100)
+
+    assert torch.equal(
+        graph_store.get_edge_id(edge_type, layout='coo'),
+        edge_id_dict[edge_type],
+    )
+    assert torch.equal(
+        graph_store.get_edge_index(edge_type, layout='coo'),
+        edge_index_dict[edge_type],
+    )