prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  On Disk Node & Edge feature & file loader
  
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `python/dgl/graphbolt/feature_store.py`
  Content:
  ```
  1 """Feature store for GraphBolt."""
  2 import torch
  3 
  4 
  5 class FeatureStore:
  6     r"""Base class for feature store."""
  7 
  8     def __init__(self):
  9         pass
  10 
  11     def read(self, key: str, ids: torch.Tensor = None):
  12         """Read a feature from the feature store.
  13 
  14         Parameters
  15         ----------
  16         key : str
  17             The key that uniquely identifies the feature in the feature store.
  18         ids : torch.Tensor, optional
  19             The index of the feature. If specified, only the specified indices
  20             of the feature are read. If None, the entire feature is returned.
  21 
  22         Returns
  23         -------
  24         torch.Tensor
  25             The read feature.
  26         """
  27         raise NotImplementedError
  28 
  29     def update(self, key: str, value: torch.Tensor, ids: torch.Tensor = None):
  30         """Update a feature in the feature store.
  31 
  32         This function is used to update a feature in the feature store. The
  33         feature is identified by a unique key, and its value is specified using
  34         a tensor.
  35 
  36         Parameters
  37         ----------
  38         key : str
  39             The key that uniquely identifies the feature in the feature store.
  40         value : torch.Tensor
  41             The updated value of the feature.
  42         ids : torch.Tensor, optional
  43             The indices of the feature to update. If specified, only the
  44             specified indices of the feature will be updated. For the feature,
  45             the `ids[i]` row is updated to `value[i]`. So the indices and value
  46             must have the same length. If None, the entire feature will be
  47             updated.
  48         """
  49         raise NotImplementedError
  50 
  51 
  52 class InMemoryFeatureStore(FeatureStore):
  53     r"""In-memory key-value feature store, where the key is a string and value
  54     is Pytorch tensor."""
  55 
  56     def __init__(self, feature_dict: dict):
  57         """Initialize an in-memory feature store.
  58 
  59         The feature store is initialized with a dictionary of tensors, where the
  60         key is the name of a feature and the value is the tensor. The value can
  61         be multi-dimensional, where the first dimension is the index of the
  62         feature.
  63 
  64         Parameters
  65         ----------
  66         feature_dict : dict, optional
  67             A dictionary of tensors.
  68 
  69         Examples
  70         --------
  71         >>> import torch
  72         >>> feature_dict = {
  73         ...     "user": torch.arange(0, 5),
  74         ...     "item": torch.arange(0, 6),
  75         ...     "rel": torch.arange(0, 6).view(2, 3),
  76         ... }
  77         >>> feature_store = InMemoryFeatureStore(feature_dict)
  78         >>> feature_store.read("user", torch.tensor([0, 1, 2]))
  79         tensor([0, 1, 2])
  80         >>> feature_store.read("item", torch.tensor([0, 1, 2]))
  81         tensor([0, 1, 2])
  82         >>> feature_store.read("rel", torch.tensor([0]))
  83         tensor([[0, 1, 2]])
  84         >>> feature_store.update("user",
  85         ... torch.ones(3, dtype=torch.long), torch.tensor([0, 1, 2]))
  86         >>> feature_store.read("user", torch.tensor([0, 1, 2]))
  87         tensor([1, 1, 1])
  88         """
  89         super(InMemoryFeatureStore, self).__init__()
  90         assert isinstance(feature_dict, dict), (
  91             f"feature_dict in InMemoryFeatureStore must be dict, "
  92             f"but got {type(feature_dict)}."
  93         )
  94         for k, v in feature_dict.items():
  95             assert isinstance(
  96                 k, str
  97             ), f"Key in InMemoryFeatureStore must be str, but got {k}."
  98             assert isinstance(v, torch.Tensor), (
  99                 f"Value in InMemoryFeatureStore must be torch.Tensor,"
  100                 f"but got {v}."
  101             )
  102 
  103         self._feature_dict = feature_dict
  104 
  105     def read(self, key: str, ids: torch.Tensor = None):
  106         """Read a feature from the feature store by index.
  107 
  108         Parameters
  109         ----------
  110         key : str
  111             The key of the feature.
  112         ids : torch.Tensor, optional
  113             The index of the feature. If specified, only the specified indices
  114             of the feature are read. If None, the entire feature is returned.
  115 
  116         Returns
  117         -------
  118         torch.Tensor
  119             The read feature.
  120         """
  121         assert (
  122             key in self._feature_dict
  123         ), f"key {key} not in {self._feature_dict.keys()}"
  124         if ids is None:
  125             return self._feature_dict[key]
  126         return self._feature_dict[key][ids]
  127 
  128     def update(self, key: str, value: torch.Tensor, ids: torch.Tensor = None):
  129         """Update a feature in the feature store.
  130 
  131         This function is used to update a feature in the feature store. The
  132         feature is identified by a unique key, and its value is specified using
  133         a tensor.
  134 
  135         Parameters
  136         ----------
  137         key : str
  138             The key that uniquely identifies the feature in the feature store.
  139         value : torch.Tensor
  140             The updated value of the feature.
  141         ids : torch.Tensor, optional
  142             The indices of the feature to update. If specified, only the
  143             specified indices of the feature will be updated. For the feature,
  144             the `ids[i]` row is updated to `value[i]`. So the indices and value
  145             must have the same length. If None, the entire feature will be
  146             updated.
  147         """
  148         assert (
  149             key in self._feature_dict
  150         ), f"key {key} not in {self._feature_dict.keys()}"
  151         if ids is None:
  152             self._feature_dict[key] = value
  153         else:
  154             assert ids.shape[0] == value.shape[0], (
  155                 f"ids and value must have the same length, "
  156                 f"but got {ids.shape[0]} and {value.shape[0]}."
  157             )
  158             self._feature_dict[key][ids] = value
  159
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/python/dgl/graphbolt/feature_store.py b/python/dgl/graphbolt/feature_store.py
  --- a/python/dgl/graphbolt/feature_store.py
  +++ b/python/dgl/graphbolt/feature_store.py
  @@ -49,18 +49,20 @@
           raise NotImplementedError
   
   
  -class InMemoryFeatureStore(FeatureStore):
  -    r"""In-memory key-value feature store, where the key is a string and value
  -    is Pytorch tensor."""
  +class TorchBasedFeatureStore(FeatureStore):
  +    r"""Torch based key-value feature store, where the key are strings and
  +    values are Pytorch tensors."""
   
       def __init__(self, feature_dict: dict):
  -        """Initialize an in-memory feature store.
  +        """Initialize a torch based feature store.
   
           The feature store is initialized with a dictionary of tensors, where the
           key is the name of a feature and the value is the tensor. The value can
           be multi-dimensional, where the first dimension is the index of the
           feature.
   
  +        Note that the values can be in memory or on disk.
  +
           Parameters
           ----------
           feature_dict : dict, optional
  @@ -74,7 +76,7 @@
           ...     "item": torch.arange(0, 6),
           ...     "rel": torch.arange(0, 6).view(2, 3),
           ... }
  -        >>> feature_store = InMemoryFeatureStore(feature_dict)
  +        >>> feature_store = TorchBasedFeatureStore(feature_dict)
           >>> feature_store.read("user", torch.tensor([0, 1, 2]))
           tensor([0, 1, 2])
           >>> feature_store.read("item", torch.tensor([0, 1, 2]))
  @@ -85,18 +87,35 @@
           ... torch.ones(3, dtype=torch.long), torch.tensor([0, 1, 2]))
           >>> feature_store.read("user", torch.tensor([0, 1, 2]))
           tensor([1, 1, 1])
  +
  +        >>> import numpy as np
  +        >>> user = np.arange(0, 5)
  +        >>> item = np.arange(0, 6)
  +        >>> np.save("/tmp/user.npy", user)
  +        >>. np.save("/tmp/item.npy", item)
  +        >>> feature_dict = {
  +        ...     "user": torch.as_tensor(np.load("/tmp/user.npy",
  +        ...             mmap_mode="r+")),
  +        ...     "item": torch.as_tensor(np.load("/tmp/item.npy",
  +        ...             mmap_mode="r+")),
  +        ... }
  +        >>> feature_store = TorchBasedFeatureStore(feature_dict)
  +        >>> feature_store.read("user", torch.tensor([0, 1, 2]))
  +        tensor([0, 1, 2])
  +        >>> feature_store.read("item", torch.tensor([3, 4, 2]))
  +        tensor([3, 4, 2])
           """
  -        super(InMemoryFeatureStore, self).__init__()
  +        super(TorchBasedFeatureStore, self).__init__()
           assert isinstance(feature_dict, dict), (
  -            f"feature_dict in InMemoryFeatureStore must be dict, "
  +            f"feature_dict in TorchBasedFeatureStore must be dict, "
               f"but got {type(feature_dict)}."
           )
           for k, v in feature_dict.items():
               assert isinstance(
                   k, str
  -            ), f"Key in InMemoryFeatureStore must be str, but got {k}."
  +            ), f"Key in TorchBasedFeatureStore must be str, but got {k}."
               assert isinstance(v, torch.Tensor), (
  -                f"Value in InMemoryFeatureStore must be torch.Tensor,"
  +                f"Value in TorchBasedFeatureStore must be torch.Tensor,"
                   f"but got {v}."
               )
   
  @@ -105,6 +124,9 @@
       def read(self, key: str, ids: torch.Tensor = None):
           """Read a feature from the feature store by index.
   
  +        The returned feature is always in memory, no matter whether the feature
  +        to read is in memory or on disk.
  +
           Parameters
           ----------
           key : str
  
verification_info: {"golden_diff": "diff --git a/python/dgl/graphbolt/feature_store.py b/python/dgl/graphbolt/feature_store.py\n--- a/python/dgl/graphbolt/feature_store.py\n+++ b/python/dgl/graphbolt/feature_store.py\n@@ -49,18 +49,20 @@\n         raise NotImplementedError\n \n \n-class InMemoryFeatureStore(FeatureStore):\n-    r\"\"\"In-memory key-value feature store, where the key is a string and value\n-    is Pytorch tensor.\"\"\"\n+class TorchBasedFeatureStore(FeatureStore):\n+    r\"\"\"Torch based key-value feature store, where the key are strings and\n+    values are Pytorch tensors.\"\"\"\n \n     def __init__(self, feature_dict: dict):\n-        \"\"\"Initialize an in-memory feature store.\n+        \"\"\"Initialize a torch based feature store.\n \n         The feature store is initialized with a dictionary of tensors, where the\n         key is the name of a feature and the value is the tensor. The value can\n         be multi-dimensional, where the first dimension is the index of the\n         feature.\n \n+        Note that the values can be in memory or on disk.\n+\n         Parameters\n         ----------\n         feature_dict : dict, optional\n@@ -74,7 +76,7 @@\n         ...     \"item\": torch.arange(0, 6),\n         ...     \"rel\": torch.arange(0, 6).view(2, 3),\n         ... }\n-        >>> feature_store = InMemoryFeatureStore(feature_dict)\n+        >>> feature_store = TorchBasedFeatureStore(feature_dict)\n         >>> feature_store.read(\"user\", torch.tensor([0, 1, 2]))\n         tensor([0, 1, 2])\n         >>> feature_store.read(\"item\", torch.tensor([0, 1, 2]))\n@@ -85,18 +87,35 @@\n         ... torch.ones(3, dtype=torch.long), torch.tensor([0, 1, 2]))\n         >>> feature_store.read(\"user\", torch.tensor([0, 1, 2]))\n         tensor([1, 1, 1])\n+\n+        >>> import numpy as np\n+        >>> user = np.arange(0, 5)\n+        >>> item = np.arange(0, 6)\n+        >>> np.save(\"/tmp/user.npy\", user)\n+        >>. np.save(\"/tmp/item.npy\", item)\n+        >>> feature_dict = {\n+        ...     \"user\": torch.as_tensor(np.load(\"/tmp/user.npy\",\n+        ...             mmap_mode=\"r+\")),\n+        ...     \"item\": torch.as_tensor(np.load(\"/tmp/item.npy\",\n+        ...             mmap_mode=\"r+\")),\n+        ... }\n+        >>> feature_store = TorchBasedFeatureStore(feature_dict)\n+        >>> feature_store.read(\"user\", torch.tensor([0, 1, 2]))\n+        tensor([0, 1, 2])\n+        >>> feature_store.read(\"item\", torch.tensor([3, 4, 2]))\n+        tensor([3, 4, 2])\n         \"\"\"\n-        super(InMemoryFeatureStore, self).__init__()\n+        super(TorchBasedFeatureStore, self).__init__()\n         assert isinstance(feature_dict, dict), (\n-            f\"feature_dict in InMemoryFeatureStore must be dict, \"\n+            f\"feature_dict in TorchBasedFeatureStore must be dict, \"\n             f\"but got {type(feature_dict)}.\"\n         )\n         for k, v in feature_dict.items():\n             assert isinstance(\n                 k, str\n-            ), f\"Key in InMemoryFeatureStore must be str, but got {k}.\"\n+            ), f\"Key in TorchBasedFeatureStore must be str, but got {k}.\"\n             assert isinstance(v, torch.Tensor), (\n-                f\"Value in InMemoryFeatureStore must be torch.Tensor,\"\n+                f\"Value in TorchBasedFeatureStore must be torch.Tensor,\"\n                 f\"but got {v}.\"\n             )\n \n@@ -105,6 +124,9 @@\n     def read(self, key: str, ids: torch.Tensor = None):\n         \"\"\"Read a feature from the feature store by index.\n \n+        The returned feature is always in memory, no matter whether the feature\n+        to read is in memory or on disk.\n+\n         Parameters\n         ----------\n         key : str\n", "issue": "On Disk Node & Edge feature & file loader\n\n", "before_files": [{"content": "\"\"\"Feature store for GraphBolt.\"\"\"\nimport torch\n\n\nclass FeatureStore:\n    r\"\"\"Base class for feature store.\"\"\"\n\n    def __init__(self):\n        pass\n\n    def read(self, key: str, ids: torch.Tensor = None):\n        \"\"\"Read a feature from the feature store.\n\n        Parameters\n        ----------\n        key : str\n            The key that uniquely identifies the feature in the feature store.\n        ids : torch.Tensor, optional\n            The index of the feature. If specified, only the specified indices\n            of the feature are read. If None, the entire feature is returned.\n\n        Returns\n        -------\n        torch.Tensor\n            The read feature.\n        \"\"\"\n        raise NotImplementedError\n\n    def update(self, key: str, value: torch.Tensor, ids: torch.Tensor = None):\n        \"\"\"Update a feature in the feature store.\n\n        This function is used to update a feature in the feature store. The\n        feature is identified by a unique key, and its value is specified using\n        a tensor.\n\n        Parameters\n        ----------\n        key : str\n            The key that uniquely identifies the feature in the feature store.\n        value : torch.Tensor\n            The updated value of the feature.\n        ids : torch.Tensor, optional\n            The indices of the feature to update. If specified, only the\n            specified indices of the feature will be updated. For the feature,\n            the `ids[i]` row is updated to `value[i]`. So the indices and value\n            must have the same length. If None, the entire feature will be\n            updated.\n        \"\"\"\n        raise NotImplementedError\n\n\nclass InMemoryFeatureStore(FeatureStore):\n    r\"\"\"In-memory key-value feature store, where the key is a string and value\n    is Pytorch tensor.\"\"\"\n\n    def __init__(self, feature_dict: dict):\n        \"\"\"Initialize an in-memory feature store.\n\n        The feature store is initialized with a dictionary of tensors, where the\n        key is the name of a feature and the value is the tensor. The value can\n        be multi-dimensional, where the first dimension is the index of the\n        feature.\n\n        Parameters\n        ----------\n        feature_dict : dict, optional\n            A dictionary of tensors.\n\n        Examples\n        --------\n        >>> import torch\n        >>> feature_dict = {\n        ...     \"user\": torch.arange(0, 5),\n        ...     \"item\": torch.arange(0, 6),\n        ...     \"rel\": torch.arange(0, 6).view(2, 3),\n        ... }\n        >>> feature_store = InMemoryFeatureStore(feature_dict)\n        >>> feature_store.read(\"user\", torch.tensor([0, 1, 2]))\n        tensor([0, 1, 2])\n        >>> feature_store.read(\"item\", torch.tensor([0, 1, 2]))\n        tensor([0, 1, 2])\n        >>> feature_store.read(\"rel\", torch.tensor([0]))\n        tensor([[0, 1, 2]])\n        >>> feature_store.update(\"user\",\n        ... torch.ones(3, dtype=torch.long), torch.tensor([0, 1, 2]))\n        >>> feature_store.read(\"user\", torch.tensor([0, 1, 2]))\n        tensor([1, 1, 1])\n        \"\"\"\n        super(InMemoryFeatureStore, self).__init__()\n        assert isinstance(feature_dict, dict), (\n            f\"feature_dict in InMemoryFeatureStore must be dict, \"\n            f\"but got {type(feature_dict)}.\"\n        )\n        for k, v in feature_dict.items():\n            assert isinstance(\n                k, str\n            ), f\"Key in InMemoryFeatureStore must be str, but got {k}.\"\n            assert isinstance(v, torch.Tensor), (\n                f\"Value in InMemoryFeatureStore must be torch.Tensor,\"\n                f\"but got {v}.\"\n            )\n\n        self._feature_dict = feature_dict\n\n    def read(self, key: str, ids: torch.Tensor = None):\n        \"\"\"Read a feature from the feature store by index.\n\n        Parameters\n        ----------\n        key : str\n            The key of the feature.\n        ids : torch.Tensor, optional\n            The index of the feature. If specified, only the specified indices\n            of the feature are read. If None, the entire feature is returned.\n\n        Returns\n        -------\n        torch.Tensor\n            The read feature.\n        \"\"\"\n        assert (\n            key in self._feature_dict\n        ), f\"key {key} not in {self._feature_dict.keys()}\"\n        if ids is None:\n            return self._feature_dict[key]\n        return self._feature_dict[key][ids]\n\n    def update(self, key: str, value: torch.Tensor, ids: torch.Tensor = None):\n        \"\"\"Update a feature in the feature store.\n\n        This function is used to update a feature in the feature store. The\n        feature is identified by a unique key, and its value is specified using\n        a tensor.\n\n        Parameters\n        ----------\n        key : str\n            The key that uniquely identifies the feature in the feature store.\n        value : torch.Tensor\n            The updated value of the feature.\n        ids : torch.Tensor, optional\n            The indices of the feature to update. If specified, only the\n            specified indices of the feature will be updated. For the feature,\n            the `ids[i]` row is updated to `value[i]`. So the indices and value\n            must have the same length. If None, the entire feature will be\n            updated.\n        \"\"\"\n        assert (\n            key in self._feature_dict\n        ), f\"key {key} not in {self._feature_dict.keys()}\"\n        if ids is None:\n            self._feature_dict[key] = value\n        else:\n            assert ids.shape[0] == value.shape[0], (\n                f\"ids and value must have the same length, \"\n                f\"but got {ids.shape[0]} and {value.shape[0]}.\"\n            )\n            self._feature_dict[key][ids] = value\n", "path": "python/dgl/graphbolt/feature_store.py"}], "after_files": [{"content": "\"\"\"Feature store for GraphBolt.\"\"\"\nimport torch\n\n\nclass FeatureStore:\n    r\"\"\"Base class for feature store.\"\"\"\n\n    def __init__(self):\n        pass\n\n    def read(self, key: str, ids: torch.Tensor = None):\n        \"\"\"Read a feature from the feature store.\n\n        Parameters\n        ----------\n        key : str\n            The key that uniquely identifies the feature in the feature store.\n        ids : torch.Tensor, optional\n            The index of the feature. If specified, only the specified indices\n            of the feature are read. If None, the entire feature is returned.\n\n        Returns\n        -------\n        torch.Tensor\n            The read feature.\n        \"\"\"\n        raise NotImplementedError\n\n    def update(self, key: str, value: torch.Tensor, ids: torch.Tensor = None):\n        \"\"\"Update a feature in the feature store.\n\n        This function is used to update a feature in the feature store. The\n        feature is identified by a unique key, and its value is specified using\n        a tensor.\n\n        Parameters\n        ----------\n        key : str\n            The key that uniquely identifies the feature in the feature store.\n        value : torch.Tensor\n            The updated value of the feature.\n        ids : torch.Tensor, optional\n            The indices of the feature to update. If specified, only the\n            specified indices of the feature will be updated. For the feature,\n            the `ids[i]` row is updated to `value[i]`. So the indices and value\n            must have the same length. If None, the entire feature will be\n            updated.\n        \"\"\"\n        raise NotImplementedError\n\n\nclass TorchBasedFeatureStore(FeatureStore):\n    r\"\"\"Torch based key-value feature store, where the key are strings and\n    values are Pytorch tensors.\"\"\"\n\n    def __init__(self, feature_dict: dict):\n        \"\"\"Initialize a torch based feature store.\n\n        The feature store is initialized with a dictionary of tensors, where the\n        key is the name of a feature and the value is the tensor. The value can\n        be multi-dimensional, where the first dimension is the index of the\n        feature.\n\n        Note that the values can be in memory or on disk.\n\n        Parameters\n        ----------\n        feature_dict : dict, optional\n            A dictionary of tensors.\n\n        Examples\n        --------\n        >>> import torch\n        >>> feature_dict = {\n        ...     \"user\": torch.arange(0, 5),\n        ...     \"item\": torch.arange(0, 6),\n        ...     \"rel\": torch.arange(0, 6).view(2, 3),\n        ... }\n        >>> feature_store = TorchBasedFeatureStore(feature_dict)\n        >>> feature_store.read(\"user\", torch.tensor([0, 1, 2]))\n        tensor([0, 1, 2])\n        >>> feature_store.read(\"item\", torch.tensor([0, 1, 2]))\n        tensor([0, 1, 2])\n        >>> feature_store.read(\"rel\", torch.tensor([0]))\n        tensor([[0, 1, 2]])\n        >>> feature_store.update(\"user\",\n        ... torch.ones(3, dtype=torch.long), torch.tensor([0, 1, 2]))\n        >>> feature_store.read(\"user\", torch.tensor([0, 1, 2]))\n        tensor([1, 1, 1])\n\n        >>> import numpy as np\n        >>> user = np.arange(0, 5)\n        >>> item = np.arange(0, 6)\n        >>> np.save(\"/tmp/user.npy\", user)\n        >>. np.save(\"/tmp/item.npy\", item)\n        >>> feature_dict = {\n        ...     \"user\": torch.as_tensor(np.load(\"/tmp/user.npy\",\n        ...             mmap_mode=\"r+\")),\n        ...     \"item\": torch.as_tensor(np.load(\"/tmp/item.npy\",\n        ...             mmap_mode=\"r+\")),\n        ... }\n        >>> feature_store = TorchBasedFeatureStore(feature_dict)\n        >>> feature_store.read(\"user\", torch.tensor([0, 1, 2]))\n        tensor([0, 1, 2])\n        >>> feature_store.read(\"item\", torch.tensor([3, 4, 2]))\n        tensor([3, 4, 2])\n        \"\"\"\n        super(TorchBasedFeatureStore, self).__init__()\n        assert isinstance(feature_dict, dict), (\n            f\"feature_dict in TorchBasedFeatureStore must be dict, \"\n            f\"but got {type(feature_dict)}.\"\n        )\n        for k, v in feature_dict.items():\n            assert isinstance(\n                k, str\n            ), f\"Key in TorchBasedFeatureStore must be str, but got {k}.\"\n            assert isinstance(v, torch.Tensor), (\n                f\"Value in TorchBasedFeatureStore must be torch.Tensor,\"\n                f\"but got {v}.\"\n            )\n\n        self._feature_dict = feature_dict\n\n    def read(self, key: str, ids: torch.Tensor = None):\n        \"\"\"Read a feature from the feature store by index.\n\n        The returned feature is always in memory, no matter whether the feature\n        to read is in memory or on disk.\n\n        Parameters\n        ----------\n        key : str\n            The key of the feature.\n        ids : torch.Tensor, optional\n            The index of the feature. If specified, only the specified indices\n            of the feature are read. If None, the entire feature is returned.\n\n        Returns\n        -------\n        torch.Tensor\n            The read feature.\n        \"\"\"\n        assert (\n            key in self._feature_dict\n        ), f\"key {key} not in {self._feature_dict.keys()}\"\n        if ids is None:\n            return self._feature_dict[key]\n        return self._feature_dict[key][ids]\n\n    def update(self, key: str, value: torch.Tensor, ids: torch.Tensor = None):\n        \"\"\"Update a feature in the feature store.\n\n        This function is used to update a feature in the feature store. The\n        feature is identified by a unique key, and its value is specified using\n        a tensor.\n\n        Parameters\n        ----------\n        key : str\n            The key that uniquely identifies the feature in the feature store.\n        value : torch.Tensor\n            The updated value of the feature.\n        ids : torch.Tensor, optional\n            The indices of the feature to update. If specified, only the\n            specified indices of the feature will be updated. For the feature,\n            the `ids[i]` row is updated to `value[i]`. So the indices and value\n            must have the same length. If None, the entire feature will be\n            updated.\n        \"\"\"\n        assert (\n            key in self._feature_dict\n        ), f\"key {key} not in {self._feature_dict.keys()}\"\n        if ids is None:\n            self._feature_dict[key] = value\n        else:\n            assert ids.shape[0] == value.shape[0], (\n                f\"ids and value must have the same length, \"\n                f\"but got {ids.shape[0]} and {value.shape[0]}.\"\n            )\n            self._feature_dict[key][ids] = value\n", "path": "python/dgl/graphbolt/feature_store.py"}]}