We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[GraphBolt] New Storage API : create_heterograph_from_shared_memory


--- 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/graph_storage/csc_sampling_graph.py`
Content:
```
1 """CSC format sampling graph."""
2 # pylint: disable= invalid-name
3 import os
4 import tarfile
5 import tempfile
6 from typing import Dict, Optional, Tuple
7 
8 import torch
9 
10 
11 class GraphMetadata:
12     r"""Class for metadata of csc sampling graph."""
13 
14     def __init__(
15         self,
16         node_type_to_id: Dict[str, int],
17         edge_type_to_id: Dict[Tuple[str, str, str], int],
18     ):
19         """Initialize the GraphMetadata object.
20 
21         Parameters
22         ----------
23         node_type_to_id : Dict[str, int]
24             Dictionary from node types to node type IDs.
25         edge_type_to_id : Dict[Tuple[str, str, str], int]
26             Dictionary from edge types to edge type IDs.
27 
28         Raises
29         ------
30         AssertionError
31             If any of the assertions fail.
32         """
33 
34         node_types = list(node_type_to_id.keys())
35         edge_types = list(edge_type_to_id.keys())
36         node_type_ids = list(node_type_to_id.values())
37         edge_type_ids = list(edge_type_to_id.values())
38 
39         # Validate node_type_to_id.
40         assert all(
41             isinstance(x, str) for x in node_types
42         ), "Node type name should be string."
43         assert all(
44             isinstance(x, int) for x in node_type_ids
45         ), "Node type id should be int."
46         assert len(node_type_ids) == len(
47             set(node_type_ids)
48         ), "Multiple node types shoud not be mapped to a same id."
49         # Validate edge_type_to_id.
50         for edge_type in edge_types:
51             src, edge, dst = edge_type
52             assert isinstance(edge, str), "Edge type name should be string."
53             assert (
54                 src in node_types
55             ), f"Unrecognized node type {src} in edge type {edge_type}"
56             assert (
57                 dst in node_types
58             ), f"Unrecognized node type {dst} in edge type {edge_type}"
59         assert all(
60             isinstance(x, int) for x in edge_type_ids
61         ), "Edge type id should be int."
62         assert len(edge_type_ids) == len(
63             set(edge_type_ids)
64         ), "Multiple edge types shoud not be mapped to a same id."
65 
66         self.node_type_to_id = node_type_to_id
67         self.edge_type_to_id = edge_type_to_id
68 
69 
70 class CSCSamplingGraph:
71     r"""Class for CSC sampling graph."""
72 
73     def __repr__(self):
74         return _csc_sampling_graph_str(self)
75 
76     def __init__(
77         self, c_csc_graph: torch.ScriptObject, metadata: Optional[GraphMetadata]
78     ):
79         self._c_csc_graph = c_csc_graph
80         self._metadata = metadata
81 
82     @property
83     def num_nodes(self) -> int:
84         """Returns the number of nodes in the graph.
85 
86         Returns
87         -------
88         int
89             The number of rows in the dense format.
90         """
91         return self._c_csc_graph.num_nodes()
92 
93     @property
94     def num_edges(self) -> int:
95         """Returns the number of edges in the graph.
96 
97         Returns
98         -------
99         int
100             The number of edges in the graph.
101         """
102         return self._c_csc_graph.num_edges()
103 
104     @property
105     def csc_indptr(self) -> torch.tensor:
106         """Returns the indices pointer in the CSC graph.
107 
108         Returns
109         -------
110         torch.tensor
111             The indices pointer in the CSC graph. An integer tensor with
112             shape `(num_nodes+1,)`.
113         """
114         return self._c_csc_graph.csc_indptr()
115 
116     @property
117     def indices(self) -> torch.tensor:
118         """Returns the indices in the CSC graph.
119 
120         Returns
121         -------
122         torch.tensor
123             The indices in the CSC graph. An integer tensor with shape
124             `(num_edges,)`.
125 
126         Notes
127         -------
128         It is assumed that edges of each node are already sorted by edge type
129         ids.
130         """
131         return self._c_csc_graph.indices()
132 
133     @property
134     def node_type_offset(self) -> Optional[torch.Tensor]:
135         """Returns the node type offset tensor if present.
136 
137         Returns
138         -------
139         torch.Tensor or None
140             If present, returns a 1D integer tensor of shape
141             `(num_node_types + 1,)`. The tensor is in ascending order as nodes
142             of the same type have continuous IDs, and larger node IDs are
143             paired with larger node type IDs. The first value is 0 and last
144             value is the number of nodes. And nodes with IDs between
145             `node_type_offset_[i]~node_type_offset_[i+1]` are of type id 'i'.
146 
147         """
148         return self._c_csc_graph.node_type_offset()
149 
150     @property
151     def type_per_edge(self) -> Optional[torch.Tensor]:
152         """Returns the edge type tensor if present.
153 
154         Returns
155         -------
156         torch.Tensor or None
157             If present, returns a 1D integer tensor of shape (num_edges,)
158             containing the type of each edge in the graph.
159         """
160         return self._c_csc_graph.type_per_edge()
161 
162     @property
163     def metadata(self) -> Optional[GraphMetadata]:
164         """Returns the metadata of the graph.
165 
166         Returns
167         -------
168         GraphMetadata or None
169             If present, returns the metadata of the graph.
170         """
171         return self._metadata
172 
173     def in_subgraph(self, nodes: torch.Tensor) -> torch.ScriptObject:
174         """Return the subgraph induced on the inbound edges of the given nodes.
175 
176         An in subgraph is equivalent to creating a new graph using the incoming
177         edges of the given nodes.
178 
179         Parameters
180         ----------
181         nodes : torch.Tensor
182             The nodes to form the subgraph which are type agnostic.
183 
184         Returns
185         -------
186         SampledSubgraph
187             The in subgraph.
188         """
189         # Ensure nodes is 1-D tensor.
190         assert nodes.dim() == 1, "Nodes should be 1-D tensor."
191         # Ensure that there are no duplicate nodes.
192         assert len(torch.unique(nodes)) == len(
193             nodes
194         ), "Nodes cannot have duplicate values."
195         return self._c_csc_graph.in_subgraph(nodes)
196 
197 
198 def from_csc(
199     csc_indptr: torch.Tensor,
200     indices: torch.Tensor,
201     node_type_offset: Optional[torch.tensor] = None,
202     type_per_edge: Optional[torch.tensor] = None,
203     metadata: Optional[GraphMetadata] = None,
204 ) -> CSCSamplingGraph:
205     """Create a CSCSamplingGraph object from a CSC representation.
206 
207     Parameters
208     ----------
209     csc_indptr : torch.Tensor
210         Pointer to the start of each row in the `indices`. An integer tensor
211         with shape `(num_nodes+1,)`.
212     indices : torch.Tensor
213         Column indices of the non-zero elements in the CSC graph. An integer
214         tensor with shape `(num_edges,)`.
215     node_type_offset : Optional[torch.tensor], optional
216         Offset of node types in the graph, by default None.
217     type_per_edge : Optional[torch.tensor], optional
218         Type ids of each edge in the graph, by default None.
219     metadata: Optional[GraphMetadata], optional
220         Metadata of the graph, by default None.
221     Returns
222     -------
223     CSCSamplingGraph
224         The created CSCSamplingGraph object.
225 
226     Examples
227     --------
228     >>> ntypes = {'n1': 0, 'n2': 1, 'n3': 2}
229     >>> etypes = {('n1', 'e1', 'n2'): 0, ('n1', 'e2', 'n3'): 1}
230     >>> metadata = graphbolt.GraphMetadata(ntypes, etypes)
231     >>> csc_indptr = torch.tensor([0, 2, 5, 7])
232     >>> indices = torch.tensor([1, 3, 0, 1, 2, 0, 3])
233     >>> node_type_offset = torch.tensor([0, 1, 2, 3])
234     >>> type_per_edge = torch.tensor([0, 1, 0, 1, 1, 0, 0])
235     >>> graph = graphbolt.from_csc(csc_indptr, indices, node_type_offset, \
236     >>>                            type_per_edge, metadata)
237     >>> print(graph)
238     CSCSamplingGraph(csc_indptr=tensor([0, 2, 5, 7]),
239                      indices=tensor([1, 3, 0, 1, 2, 0, 3]),
240                      num_nodes=3, num_edges=7)
241     """
242     if metadata and metadata.node_type_to_id and node_type_offset is not None:
243         assert len(metadata.node_type_to_id) + 1 == node_type_offset.size(
244             0
245         ), "node_type_offset length should be |ntypes| + 1."
246     return CSCSamplingGraph(
247         torch.ops.graphbolt.from_csc(
248             csc_indptr, indices, node_type_offset, type_per_edge
249         ),
250         metadata,
251     )
252 
253 
254 def _csc_sampling_graph_str(graph: CSCSamplingGraph) -> str:
255     """Internal function for converting a csc sampling graph to string
256     representation.
257     """
258     csc_indptr_str = str(graph.csc_indptr)
259     indices_str = str(graph.indices)
260     meta_str = f"num_nodes={graph.num_nodes}, num_edges={graph.num_edges}"
261     prefix = f"{type(graph).__name__}("
262 
263     def _add_indent(_str, indent):
264         lines = _str.split("\n")
265         lines = [lines[0]] + [" " * indent + line for line in lines[1:]]
266         return "\n".join(lines)
267 
268     final_str = (
269         "csc_indptr="
270         + _add_indent(csc_indptr_str, len("csc_indptr="))
271         + ",\n"
272         + "indices="
273         + _add_indent(indices_str, len("indices="))
274         + ",\n"
275         + meta_str
276         + ")"
277     )
278 
279     final_str = prefix + _add_indent(final_str, len(prefix))
280     return final_str
281 
282 
283 def load_csc_sampling_graph(filename):
284     """Load CSCSamplingGraph from tar file."""
285     with tempfile.TemporaryDirectory() as temp_dir:
286         with tarfile.open(filename, "r") as archive:
287             archive.extractall(temp_dir)
288         graph_filename = os.path.join(temp_dir, "csc_sampling_graph.pt")
289         metadata_filename = os.path.join(temp_dir, "metadata.pt")
290         return CSCSamplingGraph(
291             torch.ops.graphbolt.load_csc_sampling_graph(graph_filename),
292             torch.load(metadata_filename),
293         )
294 
295 
296 def save_csc_sampling_graph(graph, filename):
297     """Save CSCSamplingGraph to tar file."""
298     with tempfile.TemporaryDirectory() as temp_dir:
299         graph_filename = os.path.join(temp_dir, "csc_sampling_graph.pt")
300         torch.ops.graphbolt.save_csc_sampling_graph(
301             graph._c_csc_graph, graph_filename
302         )
303         metadata_filename = os.path.join(temp_dir, "metadata.pt")
304         torch.save(graph.metadata, metadata_filename)
305         with tarfile.open(filename, "w") as archive:
306             archive.add(
307                 graph_filename, arcname=os.path.basename(graph_filename)
308             )
309             archive.add(
310                 metadata_filename, arcname=os.path.basename(metadata_filename)
311             )
312     print(f"CSCSamplingGraph has been saved to {filename}.")
313
```
--- 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()

```

diff --git a/python/dgl/graphbolt/graph_storage/csc_sampling_graph.py b/python/dgl/graphbolt/graph_storage/csc_sampling_graph.py
--- a/python/dgl/graphbolt/graph_storage/csc_sampling_graph.py
+++ b/python/dgl/graphbolt/graph_storage/csc_sampling_graph.py
@@ -194,6 +194,24 @@
         ), "Nodes cannot have duplicate values."
         return self._c_csc_graph.in_subgraph(nodes)
 
+    def copy_to_shared_memory(self, shared_memory_name: str):
+        """Copy the graph to shared memory.
+
+        Parameters
+        ----------
+        shared_memory_name : str
+            Name of the shared memory.
+
+        Returns
+        -------
+        CSCSamplingGraph
+            The copied CSCSamplingGraph object on shared memory.
+        """
+        return CSCSamplingGraph(
+            self._c_csc_graph.copy_to_shared_memory(shared_memory_name),
+            self._metadata,
+        )
+
 
 def from_csc(
     csc_indptr: torch.Tensor,
@@ -251,6 +269,28 @@
     )
 
 
+def load_from_shared_memory(
+    shared_memory_name: str,
+    metadata: Optional[GraphMetadata] = None,
+) -> CSCSamplingGraph:
+    """Load a CSCSamplingGraph object from shared memory.
+
+    Parameters
+    ----------
+    shared_memory_name : str
+        Name of the shared memory.
+
+    Returns
+    -------
+    CSCSamplingGraph
+        The loaded CSCSamplingGraph object on shared memory.
+    """
+    return CSCSamplingGraph(
+        torch.ops.graphbolt.load_from_shared_memory(shared_memory_name),
+        metadata,
+    )
+
+
 def _csc_sampling_graph_str(graph: CSCSamplingGraph) -> str:
     """Internal function for converting a csc sampling graph to string
     representation.

{"golden_diff": "diff --git a/python/dgl/graphbolt/graph_storage/csc_sampling_graph.py b/python/dgl/graphbolt/graph_storage/csc_sampling_graph.py\n--- a/python/dgl/graphbolt/graph_storage/csc_sampling_graph.py\n+++ b/python/dgl/graphbolt/graph_storage/csc_sampling_graph.py\n@@ -194,6 +194,24 @@\n         ), \"Nodes cannot have duplicate values.\"\n         return self._c_csc_graph.in_subgraph(nodes)\n \n+    def copy_to_shared_memory(self, shared_memory_name: str):\n+        \"\"\"Copy the graph to shared memory.\n+\n+        Parameters\n+        ----------\n+        shared_memory_name : str\n+            Name of the shared memory.\n+\n+        Returns\n+        -------\n+        CSCSamplingGraph\n+            The copied CSCSamplingGraph object on shared memory.\n+        \"\"\"\n+        return CSCSamplingGraph(\n+            self._c_csc_graph.copy_to_shared_memory(shared_memory_name),\n+            self._metadata,\n+        )\n+\n \n def from_csc(\n     csc_indptr: torch.Tensor,\n@@ -251,6 +269,28 @@\n     )\n \n \n+def load_from_shared_memory(\n+    shared_memory_name: str,\n+    metadata: Optional[GraphMetadata] = None,\n+) -> CSCSamplingGraph:\n+    \"\"\"Load a CSCSamplingGraph object from shared memory.\n+\n+    Parameters\n+    ----------\n+    shared_memory_name : str\n+        Name of the shared memory.\n+\n+    Returns\n+    -------\n+    CSCSamplingGraph\n+        The loaded CSCSamplingGraph object on shared memory.\n+    \"\"\"\n+    return CSCSamplingGraph(\n+        torch.ops.graphbolt.load_from_shared_memory(shared_memory_name),\n+        metadata,\n+    )\n+\n+\n def _csc_sampling_graph_str(graph: CSCSamplingGraph) -> str:\n     \"\"\"Internal function for converting a csc sampling graph to string\n     representation.\n", "issue": "[GraphBolt] New Storage API : create_heterograph_from_shared_memory\n\n", "before_files": [{"content": "\"\"\"CSC format sampling graph.\"\"\"\n# pylint: disable= invalid-name\nimport os\nimport tarfile\nimport tempfile\nfrom typing import Dict, Optional, Tuple\n\nimport torch\n\n\nclass GraphMetadata:\n    r\"\"\"Class for metadata of csc sampling graph.\"\"\"\n\n    def __init__(\n        self,\n        node_type_to_id: Dict[str, int],\n        edge_type_to_id: Dict[Tuple[str, str, str], int],\n    ):\n        \"\"\"Initialize the GraphMetadata object.\n\n        Parameters\n        ----------\n        node_type_to_id : Dict[str, int]\n            Dictionary from node types to node type IDs.\n        edge_type_to_id : Dict[Tuple[str, str, str], int]\n            Dictionary from edge types to edge type IDs.\n\n        Raises\n        ------\n        AssertionError\n            If any of the assertions fail.\n        \"\"\"\n\n        node_types = list(node_type_to_id.keys())\n        edge_types = list(edge_type_to_id.keys())\n        node_type_ids = list(node_type_to_id.values())\n        edge_type_ids = list(edge_type_to_id.values())\n\n        # Validate node_type_to_id.\n        assert all(\n            isinstance(x, str) for x in node_types\n        ), \"Node type name should be string.\"\n        assert all(\n            isinstance(x, int) for x in node_type_ids\n        ), \"Node type id should be int.\"\n        assert len(node_type_ids) == len(\n            set(node_type_ids)\n        ), \"Multiple node types shoud not be mapped to a same id.\"\n        # Validate edge_type_to_id.\n        for edge_type in edge_types:\n            src, edge, dst = edge_type\n            assert isinstance(edge, str), \"Edge type name should be string.\"\n            assert (\n                src in node_types\n            ), f\"Unrecognized node type {src} in edge type {edge_type}\"\n            assert (\n                dst in node_types\n            ), f\"Unrecognized node type {dst} in edge type {edge_type}\"\n        assert all(\n            isinstance(x, int) for x in edge_type_ids\n        ), \"Edge type id should be int.\"\n        assert len(edge_type_ids) == len(\n            set(edge_type_ids)\n        ), \"Multiple edge types shoud not be mapped to a same id.\"\n\n        self.node_type_to_id = node_type_to_id\n        self.edge_type_to_id = edge_type_to_id\n\n\nclass CSCSamplingGraph:\n    r\"\"\"Class for CSC sampling graph.\"\"\"\n\n    def __repr__(self):\n        return _csc_sampling_graph_str(self)\n\n    def __init__(\n        self, c_csc_graph: torch.ScriptObject, metadata: Optional[GraphMetadata]\n    ):\n        self._c_csc_graph = c_csc_graph\n        self._metadata = metadata\n\n    @property\n    def num_nodes(self) -> int:\n        \"\"\"Returns the number of nodes in the graph.\n\n        Returns\n        -------\n        int\n            The number of rows in the dense format.\n        \"\"\"\n        return self._c_csc_graph.num_nodes()\n\n    @property\n    def num_edges(self) -> int:\n        \"\"\"Returns the number of edges in the graph.\n\n        Returns\n        -------\n        int\n            The number of edges in the graph.\n        \"\"\"\n        return self._c_csc_graph.num_edges()\n\n    @property\n    def csc_indptr(self) -> torch.tensor:\n        \"\"\"Returns the indices pointer in the CSC graph.\n\n        Returns\n        -------\n        torch.tensor\n            The indices pointer in the CSC graph. An integer tensor with\n            shape `(num_nodes+1,)`.\n        \"\"\"\n        return self._c_csc_graph.csc_indptr()\n\n    @property\n    def indices(self) -> torch.tensor:\n        \"\"\"Returns the indices in the CSC graph.\n\n        Returns\n        -------\n        torch.tensor\n            The indices in the CSC graph. An integer tensor with shape\n            `(num_edges,)`.\n\n        Notes\n        -------\n        It is assumed that edges of each node are already sorted by edge type\n        ids.\n        \"\"\"\n        return self._c_csc_graph.indices()\n\n    @property\n    def node_type_offset(self) -> Optional[torch.Tensor]:\n        \"\"\"Returns the node type offset tensor if present.\n\n        Returns\n        -------\n        torch.Tensor or None\n            If present, returns a 1D integer tensor of shape\n            `(num_node_types + 1,)`. The tensor is in ascending order as nodes\n            of the same type have continuous IDs, and larger node IDs are\n            paired with larger node type IDs. The first value is 0 and last\n            value is the number of nodes. And nodes with IDs between\n            `node_type_offset_[i]~node_type_offset_[i+1]` are of type id 'i'.\n\n        \"\"\"\n        return self._c_csc_graph.node_type_offset()\n\n    @property\n    def type_per_edge(self) -> Optional[torch.Tensor]:\n        \"\"\"Returns the edge type tensor if present.\n\n        Returns\n        -------\n        torch.Tensor or None\n            If present, returns a 1D integer tensor of shape (num_edges,)\n            containing the type of each edge in the graph.\n        \"\"\"\n        return self._c_csc_graph.type_per_edge()\n\n    @property\n    def metadata(self) -> Optional[GraphMetadata]:\n        \"\"\"Returns the metadata of the graph.\n\n        Returns\n        -------\n        GraphMetadata or None\n            If present, returns the metadata of the graph.\n        \"\"\"\n        return self._metadata\n\n    def in_subgraph(self, nodes: torch.Tensor) -> torch.ScriptObject:\n        \"\"\"Return the subgraph induced on the inbound edges of the given nodes.\n\n        An in subgraph is equivalent to creating a new graph using the incoming\n        edges of the given nodes.\n\n        Parameters\n        ----------\n        nodes : torch.Tensor\n            The nodes to form the subgraph which are type agnostic.\n\n        Returns\n        -------\n        SampledSubgraph\n            The in subgraph.\n        \"\"\"\n        # Ensure nodes is 1-D tensor.\n        assert nodes.dim() == 1, \"Nodes should be 1-D tensor.\"\n        # Ensure that there are no duplicate nodes.\n        assert len(torch.unique(nodes)) == len(\n            nodes\n        ), \"Nodes cannot have duplicate values.\"\n        return self._c_csc_graph.in_subgraph(nodes)\n\n\ndef from_csc(\n    csc_indptr: torch.Tensor,\n    indices: torch.Tensor,\n    node_type_offset: Optional[torch.tensor] = None,\n    type_per_edge: Optional[torch.tensor] = None,\n    metadata: Optional[GraphMetadata] = None,\n) -> CSCSamplingGraph:\n    \"\"\"Create a CSCSamplingGraph object from a CSC representation.\n\n    Parameters\n    ----------\n    csc_indptr : torch.Tensor\n        Pointer to the start of each row in the `indices`. An integer tensor\n        with shape `(num_nodes+1,)`.\n    indices : torch.Tensor\n        Column indices of the non-zero elements in the CSC graph. An integer\n        tensor with shape `(num_edges,)`.\n    node_type_offset : Optional[torch.tensor], optional\n        Offset of node types in the graph, by default None.\n    type_per_edge : Optional[torch.tensor], optional\n        Type ids of each edge in the graph, by default None.\n    metadata: Optional[GraphMetadata], optional\n        Metadata of the graph, by default None.\n    Returns\n    -------\n    CSCSamplingGraph\n        The created CSCSamplingGraph object.\n\n    Examples\n    --------\n    >>> ntypes = {'n1': 0, 'n2': 1, 'n3': 2}\n    >>> etypes = {('n1', 'e1', 'n2'): 0, ('n1', 'e2', 'n3'): 1}\n    >>> metadata = graphbolt.GraphMetadata(ntypes, etypes)\n    >>> csc_indptr = torch.tensor([0, 2, 5, 7])\n    >>> indices = torch.tensor([1, 3, 0, 1, 2, 0, 3])\n    >>> node_type_offset = torch.tensor([0, 1, 2, 3])\n    >>> type_per_edge = torch.tensor([0, 1, 0, 1, 1, 0, 0])\n    >>> graph = graphbolt.from_csc(csc_indptr, indices, node_type_offset, \\\n    >>>                            type_per_edge, metadata)\n    >>> print(graph)\n    CSCSamplingGraph(csc_indptr=tensor([0, 2, 5, 7]),\n                     indices=tensor([1, 3, 0, 1, 2, 0, 3]),\n                     num_nodes=3, num_edges=7)\n    \"\"\"\n    if metadata and metadata.node_type_to_id and node_type_offset is not None:\n        assert len(metadata.node_type_to_id) + 1 == node_type_offset.size(\n            0\n        ), \"node_type_offset length should be |ntypes| + 1.\"\n    return CSCSamplingGraph(\n        torch.ops.graphbolt.from_csc(\n            csc_indptr, indices, node_type_offset, type_per_edge\n        ),\n        metadata,\n    )\n\n\ndef _csc_sampling_graph_str(graph: CSCSamplingGraph) -> str:\n    \"\"\"Internal function for converting a csc sampling graph to string\n    representation.\n    \"\"\"\n    csc_indptr_str = str(graph.csc_indptr)\n    indices_str = str(graph.indices)\n    meta_str = f\"num_nodes={graph.num_nodes}, num_edges={graph.num_edges}\"\n    prefix = f\"{type(graph).__name__}(\"\n\n    def _add_indent(_str, indent):\n        lines = _str.split(\"\\n\")\n        lines = [lines[0]] + [\" \" * indent + line for line in lines[1:]]\n        return \"\\n\".join(lines)\n\n    final_str = (\n        \"csc_indptr=\"\n        + _add_indent(csc_indptr_str, len(\"csc_indptr=\"))\n        + \",\\n\"\n        + \"indices=\"\n        + _add_indent(indices_str, len(\"indices=\"))\n        + \",\\n\"\n        + meta_str\n        + \")\"\n    )\n\n    final_str = prefix + _add_indent(final_str, len(prefix))\n    return final_str\n\n\ndef load_csc_sampling_graph(filename):\n    \"\"\"Load CSCSamplingGraph from tar file.\"\"\"\n    with tempfile.TemporaryDirectory() as temp_dir:\n        with tarfile.open(filename, \"r\") as archive:\n            archive.extractall(temp_dir)\n        graph_filename = os.path.join(temp_dir, \"csc_sampling_graph.pt\")\n        metadata_filename = os.path.join(temp_dir, \"metadata.pt\")\n        return CSCSamplingGraph(\n            torch.ops.graphbolt.load_csc_sampling_graph(graph_filename),\n            torch.load(metadata_filename),\n        )\n\n\ndef save_csc_sampling_graph(graph, filename):\n    \"\"\"Save CSCSamplingGraph to tar file.\"\"\"\n    with tempfile.TemporaryDirectory() as temp_dir:\n        graph_filename = os.path.join(temp_dir, \"csc_sampling_graph.pt\")\n        torch.ops.graphbolt.save_csc_sampling_graph(\n            graph._c_csc_graph, graph_filename\n        )\n        metadata_filename = os.path.join(temp_dir, \"metadata.pt\")\n        torch.save(graph.metadata, metadata_filename)\n        with tarfile.open(filename, \"w\") as archive:\n            archive.add(\n                graph_filename, arcname=os.path.basename(graph_filename)\n            )\n            archive.add(\n                metadata_filename, arcname=os.path.basename(metadata_filename)\n            )\n    print(f\"CSCSamplingGraph has been saved to {filename}.\")\n", "path": "python/dgl/graphbolt/graph_storage/csc_sampling_graph.py"}], "after_files": [{"content": "\"\"\"CSC format sampling graph.\"\"\"\n# pylint: disable= invalid-name\nimport os\nimport tarfile\nimport tempfile\nfrom typing import Dict, Optional, Tuple\n\nimport torch\n\n\nclass GraphMetadata:\n    r\"\"\"Class for metadata of csc sampling graph.\"\"\"\n\n    def __init__(\n        self,\n        node_type_to_id: Dict[str, int],\n        edge_type_to_id: Dict[Tuple[str, str, str], int],\n    ):\n        \"\"\"Initialize the GraphMetadata object.\n\n        Parameters\n        ----------\n        node_type_to_id : Dict[str, int]\n            Dictionary from node types to node type IDs.\n        edge_type_to_id : Dict[Tuple[str, str, str], int]\n            Dictionary from edge types to edge type IDs.\n\n        Raises\n        ------\n        AssertionError\n            If any of the assertions fail.\n        \"\"\"\n\n        node_types = list(node_type_to_id.keys())\n        edge_types = list(edge_type_to_id.keys())\n        node_type_ids = list(node_type_to_id.values())\n        edge_type_ids = list(edge_type_to_id.values())\n\n        # Validate node_type_to_id.\n        assert all(\n            isinstance(x, str) for x in node_types\n        ), \"Node type name should be string.\"\n        assert all(\n            isinstance(x, int) for x in node_type_ids\n        ), \"Node type id should be int.\"\n        assert len(node_type_ids) == len(\n            set(node_type_ids)\n        ), \"Multiple node types shoud not be mapped to a same id.\"\n        # Validate edge_type_to_id.\n        for edge_type in edge_types:\n            src, edge, dst = edge_type\n            assert isinstance(edge, str), \"Edge type name should be string.\"\n            assert (\n                src in node_types\n            ), f\"Unrecognized node type {src} in edge type {edge_type}\"\n            assert (\n                dst in node_types\n            ), f\"Unrecognized node type {dst} in edge type {edge_type}\"\n        assert all(\n            isinstance(x, int) for x in edge_type_ids\n        ), \"Edge type id should be int.\"\n        assert len(edge_type_ids) == len(\n            set(edge_type_ids)\n        ), \"Multiple edge types shoud not be mapped to a same id.\"\n\n        self.node_type_to_id = node_type_to_id\n        self.edge_type_to_id = edge_type_to_id\n\n\nclass CSCSamplingGraph:\n    r\"\"\"Class for CSC sampling graph.\"\"\"\n\n    def __repr__(self):\n        return _csc_sampling_graph_str(self)\n\n    def __init__(\n        self, c_csc_graph: torch.ScriptObject, metadata: Optional[GraphMetadata]\n    ):\n        self._c_csc_graph = c_csc_graph\n        self._metadata = metadata\n\n    @property\n    def num_nodes(self) -> int:\n        \"\"\"Returns the number of nodes in the graph.\n\n        Returns\n        -------\n        int\n            The number of rows in the dense format.\n        \"\"\"\n        return self._c_csc_graph.num_nodes()\n\n    @property\n    def num_edges(self) -> int:\n        \"\"\"Returns the number of edges in the graph.\n\n        Returns\n        -------\n        int\n            The number of edges in the graph.\n        \"\"\"\n        return self._c_csc_graph.num_edges()\n\n    @property\n    def csc_indptr(self) -> torch.tensor:\n        \"\"\"Returns the indices pointer in the CSC graph.\n\n        Returns\n        -------\n        torch.tensor\n            The indices pointer in the CSC graph. An integer tensor with\n            shape `(num_nodes+1,)`.\n        \"\"\"\n        return self._c_csc_graph.csc_indptr()\n\n    @property\n    def indices(self) -> torch.tensor:\n        \"\"\"Returns the indices in the CSC graph.\n\n        Returns\n        -------\n        torch.tensor\n            The indices in the CSC graph. An integer tensor with shape\n            `(num_edges,)`.\n\n        Notes\n        -------\n        It is assumed that edges of each node are already sorted by edge type\n        ids.\n        \"\"\"\n        return self._c_csc_graph.indices()\n\n    @property\n    def node_type_offset(self) -> Optional[torch.Tensor]:\n        \"\"\"Returns the node type offset tensor if present.\n\n        Returns\n        -------\n        torch.Tensor or None\n            If present, returns a 1D integer tensor of shape\n            `(num_node_types + 1,)`. The tensor is in ascending order as nodes\n            of the same type have continuous IDs, and larger node IDs are\n            paired with larger node type IDs. The first value is 0 and last\n            value is the number of nodes. And nodes with IDs between\n            `node_type_offset_[i]~node_type_offset_[i+1]` are of type id 'i'.\n\n        \"\"\"\n        return self._c_csc_graph.node_type_offset()\n\n    @property\n    def type_per_edge(self) -> Optional[torch.Tensor]:\n        \"\"\"Returns the edge type tensor if present.\n\n        Returns\n        -------\n        torch.Tensor or None\n            If present, returns a 1D integer tensor of shape (num_edges,)\n            containing the type of each edge in the graph.\n        \"\"\"\n        return self._c_csc_graph.type_per_edge()\n\n    @property\n    def metadata(self) -> Optional[GraphMetadata]:\n        \"\"\"Returns the metadata of the graph.\n\n        Returns\n        -------\n        GraphMetadata or None\n            If present, returns the metadata of the graph.\n        \"\"\"\n        return self._metadata\n\n    def in_subgraph(self, nodes: torch.Tensor) -> torch.ScriptObject:\n        \"\"\"Return the subgraph induced on the inbound edges of the given nodes.\n\n        An in subgraph is equivalent to creating a new graph using the incoming\n        edges of the given nodes.\n\n        Parameters\n        ----------\n        nodes : torch.Tensor\n            The nodes to form the subgraph which are type agnostic.\n\n        Returns\n        -------\n        SampledSubgraph\n            The in subgraph.\n        \"\"\"\n        # Ensure nodes is 1-D tensor.\n        assert nodes.dim() == 1, \"Nodes should be 1-D tensor.\"\n        # Ensure that there are no duplicate nodes.\n        assert len(torch.unique(nodes)) == len(\n            nodes\n        ), \"Nodes cannot have duplicate values.\"\n        return self._c_csc_graph.in_subgraph(nodes)\n\n    def copy_to_shared_memory(self, shared_memory_name: str):\n        \"\"\"Copy the graph to shared memory.\n\n        Parameters\n        ----------\n        shared_memory_name : str\n            Name of the shared memory.\n\n        Returns\n        -------\n        CSCSamplingGraph\n            The copied CSCSamplingGraph object on shared memory.\n        \"\"\"\n        return CSCSamplingGraph(\n            self._c_csc_graph.copy_to_shared_memory(shared_memory_name),\n            self._metadata,\n        )\n\n\ndef from_csc(\n    csc_indptr: torch.Tensor,\n    indices: torch.Tensor,\n    node_type_offset: Optional[torch.tensor] = None,\n    type_per_edge: Optional[torch.tensor] = None,\n    metadata: Optional[GraphMetadata] = None,\n) -> CSCSamplingGraph:\n    \"\"\"Create a CSCSamplingGraph object from a CSC representation.\n\n    Parameters\n    ----------\n    csc_indptr : torch.Tensor\n        Pointer to the start of each row in the `indices`. An integer tensor\n        with shape `(num_nodes+1,)`.\n    indices : torch.Tensor\n        Column indices of the non-zero elements in the CSC graph. An integer\n        tensor with shape `(num_edges,)`.\n    node_type_offset : Optional[torch.tensor], optional\n        Offset of node types in the graph, by default None.\n    type_per_edge : Optional[torch.tensor], optional\n        Type ids of each edge in the graph, by default None.\n    metadata: Optional[GraphMetadata], optional\n        Metadata of the graph, by default None.\n    Returns\n    -------\n    CSCSamplingGraph\n        The created CSCSamplingGraph object.\n\n    Examples\n    --------\n    >>> ntypes = {'n1': 0, 'n2': 1, 'n3': 2}\n    >>> etypes = {('n1', 'e1', 'n2'): 0, ('n1', 'e2', 'n3'): 1}\n    >>> metadata = graphbolt.GraphMetadata(ntypes, etypes)\n    >>> csc_indptr = torch.tensor([0, 2, 5, 7])\n    >>> indices = torch.tensor([1, 3, 0, 1, 2, 0, 3])\n    >>> node_type_offset = torch.tensor([0, 1, 2, 3])\n    >>> type_per_edge = torch.tensor([0, 1, 0, 1, 1, 0, 0])\n    >>> graph = graphbolt.from_csc(csc_indptr, indices, node_type_offset, \\\n    >>>                            type_per_edge, metadata)\n    >>> print(graph)\n    CSCSamplingGraph(csc_indptr=tensor([0, 2, 5, 7]),\n                     indices=tensor([1, 3, 0, 1, 2, 0, 3]),\n                     num_nodes=3, num_edges=7)\n    \"\"\"\n    if metadata and metadata.node_type_to_id and node_type_offset is not None:\n        assert len(metadata.node_type_to_id) + 1 == node_type_offset.size(\n            0\n        ), \"node_type_offset length should be |ntypes| + 1.\"\n    return CSCSamplingGraph(\n        torch.ops.graphbolt.from_csc(\n            csc_indptr, indices, node_type_offset, type_per_edge\n        ),\n        metadata,\n    )\n\n\ndef load_from_shared_memory(\n    shared_memory_name: str,\n    metadata: Optional[GraphMetadata] = None,\n) -> CSCSamplingGraph:\n    \"\"\"Load a CSCSamplingGraph object from shared memory.\n\n    Parameters\n    ----------\n    shared_memory_name : str\n        Name of the shared memory.\n\n    Returns\n    -------\n    CSCSamplingGraph\n        The loaded CSCSamplingGraph object on shared memory.\n    \"\"\"\n    return CSCSamplingGraph(\n        torch.ops.graphbolt.load_from_shared_memory(shared_memory_name),\n        metadata,\n    )\n\n\ndef _csc_sampling_graph_str(graph: CSCSamplingGraph) -> str:\n    \"\"\"Internal function for converting a csc sampling graph to string\n    representation.\n    \"\"\"\n    csc_indptr_str = str(graph.csc_indptr)\n    indices_str = str(graph.indices)\n    meta_str = f\"num_nodes={graph.num_nodes}, num_edges={graph.num_edges}\"\n    prefix = f\"{type(graph).__name__}(\"\n\n    def _add_indent(_str, indent):\n        lines = _str.split(\"\\n\")\n        lines = [lines[0]] + [\" \" * indent + line for line in lines[1:]]\n        return \"\\n\".join(lines)\n\n    final_str = (\n        \"csc_indptr=\"\n        + _add_indent(csc_indptr_str, len(\"csc_indptr=\"))\n        + \",\\n\"\n        + \"indices=\"\n        + _add_indent(indices_str, len(\"indices=\"))\n        + \",\\n\"\n        + meta_str\n        + \")\"\n    )\n\n    final_str = prefix + _add_indent(final_str, len(prefix))\n    return final_str\n\n\ndef load_csc_sampling_graph(filename):\n    \"\"\"Load CSCSamplingGraph from tar file.\"\"\"\n    with tempfile.TemporaryDirectory() as temp_dir:\n        with tarfile.open(filename, \"r\") as archive:\n            archive.extractall(temp_dir)\n        graph_filename = os.path.join(temp_dir, \"csc_sampling_graph.pt\")\n        metadata_filename = os.path.join(temp_dir, \"metadata.pt\")\n        return CSCSamplingGraph(\n            torch.ops.graphbolt.load_csc_sampling_graph(graph_filename),\n            torch.load(metadata_filename),\n        )\n\n\ndef save_csc_sampling_graph(graph, filename):\n    \"\"\"Save CSCSamplingGraph to tar file.\"\"\"\n    with tempfile.TemporaryDirectory() as temp_dir:\n        graph_filename = os.path.join(temp_dir, \"csc_sampling_graph.pt\")\n        torch.ops.graphbolt.save_csc_sampling_graph(\n            graph._c_csc_graph, graph_filename\n        )\n        metadata_filename = os.path.join(temp_dir, \"metadata.pt\")\n        torch.save(graph.metadata, metadata_filename)\n        with tarfile.open(filename, \"w\") as archive:\n            archive.add(\n                graph_filename, arcname=os.path.basename(graph_filename)\n            )\n            archive.add(\n                metadata_filename, arcname=os.path.basename(metadata_filename)\n            )\n    print(f\"CSCSamplingGraph has been saved to {filename}.\")\n", "path": "python/dgl/graphbolt/graph_storage/csc_sampling_graph.py"}]}