--- a
+++ b/torchdrug/data/graph.py
@@ -0,0 +1,1851 @@
+import math
+import warnings
+from functools import reduce
+from collections import defaultdict
+
+import networkx as nx
+
+from matplotlib import pyplot as plt
+import torch
+from torch_scatter import scatter_add, scatter_min
+
+from torchdrug import core, utils
+from torchdrug.data import Dictionary
+from torchdrug.utils import pretty
+
+plt.switch_backend("agg")
+
+
+class Graph(core._MetaContainer):
+    r"""
+    Basic container for sparse graphs.
+
+    To batch graphs with variadic sizes, use :meth:`data.Graph.pack <torchdrug.data.Graph.pack>`.
+    This will return a PackedGraph object with the following block diagonal adjacency matrix.
+
+    .. math::
+
+        \begin{bmatrix}
+            A_1    & \cdots & 0      \\
+            \vdots & \ddots & \vdots \\
+            0      & \cdots & A_n
+        \end{bmatrix}
+
+    where :math:`A_i` is the adjacency of :math:`i`-th graph.
+
+    You may register dynamic attributes for each graph.
+    The registered attributes will be automatically processed during packing.
+
+    .. warning::
+
+        This class doesn't enforce any order on the edges.
+
+    Example::
+
+        >>> graph = data.Graph(torch.randint(10, (30, 2)))
+        >>> with graph.node():
+        >>>     graph.my_node_attr = torch.rand(10, 5, 5)
+
+    Parameters:
+        edge_list (array_like, optional): list of edges of shape :math:`(|E|, 2)` or :math:`(|E|, 3)`.
+            Each tuple is (node_in, node_out) or (node_in, node_out, relation).
+        edge_weight (array_like, optional): edge weights of shape :math:`(|E|,)`
+        num_node (int, optional): number of nodes.
+            By default, it will be inferred from the largest id in `edge_list`
+        num_relation (int, optional): number of relations
+        node_feature (array_like, optional): node features of shape :math:`(|V|, ...)`
+        edge_feature (array_like, optional): edge features of shape :math:`(|E|, ...)`
+        graph_feature (array_like, optional): graph feature of any shape
+    """
+
+    _meta_types = {"node", "edge", "graph", "node reference", "edge reference", "graph reference"}
+
+    def __init__(self, edge_list=None, edge_weight=None, num_node=None, num_relation=None,
+                 node_feature=None, edge_feature=None, graph_feature=None, **kwargs):
+        super(Graph, self).__init__(**kwargs)
+        # edge_list: N * [h, t] or N * [h, t, r]
+        edge_list, num_edge = self._standarize_edge_list(edge_list, num_relation)
+        edge_weight = self._standarize_edge_weight(edge_weight, edge_list)
+
+        num_node = self._standarize_num_node(num_node, edge_list)
+        num_relation = self._standarize_num_relation(num_relation, edge_list)
+
+        self._edge_list = edge_list
+        self._edge_weight = edge_weight
+        self.num_node = num_node
+        self.num_edge = num_edge
+        self.num_relation = num_relation
+
+        if node_feature is not None:
+            with self.node():
+                self.node_feature = torch.as_tensor(node_feature, device=self.device)
+        if edge_feature is not None:
+            with self.edge():
+                self.edge_feature = torch.as_tensor(edge_feature, device=self.device)
+        if graph_feature is not None:
+            with self.graph():
+                self.graph_feature = torch.as_tensor(graph_feature, device=self.device)
+
+    def node(self):
+        """
+        Context manager for node attributes.
+        """
+        return self.context("node")
+
+    def edge(self):
+        """
+        Context manager for edge attributes.
+        """
+        return self.context("edge")
+
+    def graph(self):
+        """
+        Context manager for graph attributes.
+        """
+        return self.context("graph")
+
+    def node_reference(self):
+        """
+        Context manager for node references.
+        """
+        return self.context("node reference")
+
+    def edge_reference(self):
+        """
+        Context manager for edge references.
+        """
+        return self.context("edge reference")
+
+    def graph_reference(self):
+        """
+        Context manager for graph references.
+        """
+        return self.context("graph reference")
+
+    def _check_attribute(self, key, value):
+        for type in self._meta_contexts:
+            if "reference" in type:
+                if value.dtype != torch.long:
+                    raise TypeError("Tensors used as reference must be long tensors")
+            if type == "node":
+                if len(value) != self.num_node:
+                    raise ValueError("Expect node attribute `%s` to have shape (%d, *), but found %s" %
+                                     (key, self.num_node, value.shape))
+            elif type == "edge":
+                if len(value) != self.num_edge:
+                    raise ValueError("Expect edge attribute `%s` to have shape (%d, *), but found %s" %
+                                     (key, self.num_edge, value.shape))
+            elif type == "node reference":
+                is_valid = (value >= -1) & (value < self.num_node)
+                if not is_valid.all():
+                    error_value = value[~is_valid]
+                    raise ValueError("Expect node reference in [-1, %d), but found %d" %
+                                     (self.num_node, error_value[0]))
+            elif type == "edge reference":
+                is_valid = (value >= -1) & (value < self.num_edge)
+                if not is_valid.all():
+                    error_value = value[~is_valid]
+                    raise ValueError("Expect edge reference in [-1, %d), but found %d" %
+                                     (self.num_edge, error_value[0]))
+            elif type == "graph reference":
+                is_valid = (value >= -1) & (value < self.batch_size)
+                if not is_valid.all():
+                    error_value = value[~is_valid]
+                    raise ValueError("Expect graph reference in [-1, %d), but found %d" %
+                                     (self.batch_size, error_value[0]))
+
+    def __setattr__(self, key, value):
+        if hasattr(self, "meta_dict"):
+            self._check_attribute(key, value)
+        super(Graph, self).__setattr__(key, value)
+
+    def _standarize_edge_list(self, edge_list, num_relation):
+        if edge_list is not None and len(edge_list):
+            if isinstance(edge_list, torch.Tensor) and edge_list.dtype != torch.long:
+                try:
+                    edge_list = torch.LongTensor(edge_list)
+                except TypeError:
+                    raise TypeError("Can't convert `edge_list` to torch.long")
+            else:
+                edge_list = torch.as_tensor(edge_list, dtype=torch.long)
+        else:
+            num_element = 2 if num_relation is None else 3
+            if isinstance(edge_list, torch.Tensor):
+                device = edge_list.device
+            else:
+                device = "cpu"
+            edge_list = torch.zeros(0, num_element, dtype=torch.long, device=device)
+        if (edge_list < 0).any():
+            raise ValueError("`edge_list` should only contain non-negative indexes")
+        num_edge = torch.tensor(len(edge_list), device=edge_list.device)
+        return edge_list, num_edge
+
+    def _standarize_edge_weight(self, edge_weight, edge_list):
+        if edge_weight is not None:
+            edge_weight = torch.as_tensor(edge_weight, dtype=torch.float, device=edge_list.device)
+            if len(edge_list) != len(edge_weight):
+                raise ValueError("`edge_list` and `edge_weight` should be the same size, but found %d and %d"
+                                 % (len(edge_list), len(edge_weight)))
+        else:
+            edge_weight = torch.ones(len(edge_list), device=edge_list.device)
+        return edge_weight
+
+    def _standarize_num_node(self, num_node, edge_list):
+        if num_node is None:
+            num_node = self._maybe_num_node(edge_list)
+        num_node = torch.as_tensor(num_node, device=edge_list.device)
+        if (edge_list[:, :2] >= num_node).any():
+            raise ValueError("`num_node` is %d, but found node %d in `edge_list`" % (num_node, edge_list[:, :2].max()))
+        return num_node
+
+    def _standarize_num_relation(self, num_relation, edge_list):
+        if num_relation is None and edge_list.shape[1] > 2:
+            num_relation = self._maybe_num_relation(edge_list)
+        if num_relation is not None:
+            num_relation = torch.as_tensor(num_relation, device=edge_list.device)
+            if edge_list.shape[1] <= 2:
+                raise ValueError("`num_relation` is provided, but the number of dims of `edge_list` is less than 3.")
+            elif (edge_list[:, 2] >= num_relation).any():
+                raise ValueError("`num_relation` is %d, but found relation %d in `edge_list`" % (num_relation, edge_list[:, 2].max()))
+        return num_relation
+
+    def _maybe_num_node(self, edge_list):
+        warnings.warn("_maybe_num_node() is used to determine the number of nodes. "
+                      "This may underestimate the count if there are isolated nodes.")
+        if len(edge_list):
+            return edge_list[:, :2].max().item() + 1
+        else:
+            return 0
+
+    def _maybe_num_relation(self, edge_list):
+        warnings.warn("_maybe_num_relation() is used to determine the number of relations. "
+                      "This may underestimate the count if there are unseen relations.")
+        return edge_list[:, 2].max().item() + 1
+
+    def _standarize_index(self, index, count):
+        if isinstance(index, slice):
+            start = index.start or 0
+            if start < 0:
+                start += count
+            stop = index.stop or count
+            if stop < 0:
+                stop += count
+            step = index.step or 1
+            index = torch.arange(start, stop, step, device=self.device)
+        else:
+            index = torch.as_tensor(index, device=self.device)
+            if index.ndim == 0:
+                index = index.unsqueeze(0)
+            if index.dtype == torch.bool:
+                if index.shape != (count,):
+                    raise IndexError("Invalid mask. Expect mask to have shape %s, but found %s" %
+                                     ((int(count),), tuple(index.shape)))
+                index = index.nonzero().squeeze(-1)
+            else:
+                index = index.long()
+                max_index = -1 if len(index) == 0 else index.max().item()
+                if max_index >= count:
+                    raise IndexError("Invalid index. Expect index smaller than %d, but found %d" % (count, max_index))
+        return index
+
+    def _get_mapping(self, index, count):
+        index = self._standarize_index(index, count)
+        if (index.bincount() > 1).any():
+            raise ValueError("Can't create mapping for duplicate index")
+        mapping = -torch.ones(count + 1, dtype=torch.long, device=self.device)
+        mapping[index] = torch.arange(len(index), device=self.device)
+        return mapping
+
+    def _get_repeat_pack_offsets(self, num_xs, repeats):
+        new_num_xs = num_xs.repeat_interleave(repeats)
+        cum_repeats_shifted = repeats.cumsum(0) - repeats
+        new_num_xs[cum_repeats_shifted] -= num_xs
+        offsets = new_num_xs.cumsum(0)
+        return offsets
+
+    @classmethod
+    def from_dense(cls, adjacency, node_feature=None, edge_feature=None):
+        """
+        Create a sparse graph from a dense adjacency matrix.
+        For zero entries in the adjacency matrix, their edge features will be ignored.
+
+        Parameters:
+            adjacency (array_like): adjacency matrix of shape :math:`(|V|, |V|)` or :math:`(|V|, |V|, |R|)`
+            node_feature (array_like): node features of shape :math:`(|V|, ...)`
+            edge_feature (array_like): edge features of shape :math:`(|V|, |V|, ...)` or :math:`(|V|, |V|, |R|, ...)`
+        """
+        adjacency = torch.as_tensor(adjacency)
+        if adjacency.shape[0] != adjacency.shape[1]:
+            raise ValueError("`adjacency` should be a square matrix, but found %d and %d" % adjacency.shape[:2])
+
+        edge_list = adjacency.nonzero()
+        edge_weight = adjacency[tuple(edge_list.t())]
+        num_node = adjacency.shape[0]
+        num_relation = adjacency.shape[2] if adjacency.ndim > 2 else None
+        if edge_feature is not None:
+            edge_feature = torch.as_tensor(edge_feature)
+            edge_feature = edge_feature[tuple(edge_list.t())]
+
+        return cls(edge_list, edge_weight, num_node, num_relation, node_feature, edge_feature)
+
+    def connected_components(self):
+        """
+        Split this graph into connected components.
+
+        Returns:
+            (PackedGraph, LongTensor): connected components, number of connected components per graph
+        """
+        node_in, node_out = self.edge_list.t()[:2]
+        range = torch.arange(self.num_node, device=self.device)
+        node_in, node_out = torch.cat([node_in, node_out, range]), torch.cat([node_out, node_in, range])
+
+        # find connected component
+        # O(|E|d), d is the diameter of the graph
+        min_neighbor = torch.arange(self.num_node, device=self.device)
+        last = torch.zeros_like(min_neighbor)
+        while not torch.equal(min_neighbor, last):
+            last = min_neighbor
+            min_neighbor = scatter_min(min_neighbor[node_out], node_in, dim_size=self.num_node)[0]
+        anchor = torch.unique(min_neighbor)
+        num_cc = self.node2graph[anchor].bincount(minlength=self.batch_size)
+        return self.split(min_neighbor), num_cc
+
+    def split(self, node2graph):
+        """
+        Split a graph into multiple disconnected graphs.
+
+        Parameters:
+            node2graph (array_like): ID of the graph each node belongs to
+
+        Returns:
+            PackedGraph
+        """
+        node2graph = torch.as_tensor(node2graph, dtype=torch.long, device=self.device)
+        # coalesce arbitrary graph IDs to [0, n)
+        _, node2graph = torch.unique(node2graph, return_inverse=True)
+        num_graph = node2graph.max() + 1
+        index = node2graph.argsort()
+        mapping = torch.zeros_like(index)
+        mapping[index] = torch.arange(len(index), device=self.device)
+
+        node_in, node_out = self.edge_list.t()[:2]
+        edge_mask = node2graph[node_in] == node2graph[node_out]
+        edge2graph = node2graph[node_in]
+        edge_index = edge2graph.argsort()
+        edge_index = edge_index[edge_mask[edge_index]]
+
+        prepend = torch.tensor([-1], device=self.device)
+        is_first_node = torch.diff(node2graph[index], prepend=prepend) > 0
+        graph_index = self.node2graph[index[is_first_node]]
+
+        edge_list = self.edge_list.clone()
+        edge_list[:, :2] = mapping[edge_list[:, :2]]
+
+        num_nodes = node2graph.bincount(minlength=num_graph)
+        num_edges = edge2graph[edge_index].bincount(minlength=num_graph)
+
+        num_cum_nodes = num_nodes.cumsum(0)
+        offsets = (num_cum_nodes - num_nodes)[edge2graph[edge_index]]
+
+        data_dict, meta_dict = self.data_mask(index, edge_index, graph_index=graph_index, exclude="graph reference")
+
+        return self.packed_type(edge_list[edge_index], edge_weight=self.edge_weight[edge_index], num_nodes=num_nodes,
+                                num_edges=num_edges, num_relation=self.num_relation, offsets=offsets,
+                                meta_dict=meta_dict, **data_dict)
+
+    @classmethod
+    def pack(cls, graphs):
+        """
+        Pack a list of graphs into a PackedGraph object.
+
+        Parameters:
+            graphs (list of Graph): list of graphs
+
+        Returns:
+            PackedGraph
+        """
+        edge_list = []
+        edge_weight = []
+        num_nodes = []
+        num_edges = []
+        num_relation = -1
+        num_cum_node = 0
+        num_cum_edge = 0
+        num_graph = 0
+        data_dict = defaultdict(list)
+        meta_dict = graphs[0].meta_dict
+        for graph in graphs:
+            edge_list.append(graph.edge_list)
+            edge_weight.append(graph.edge_weight)
+            num_nodes.append(graph.num_node)
+            num_edges.append(graph.num_edge)
+            for k, v in graph.data_dict.items():
+                for type in meta_dict[k]:
+                    if type == "graph":
+                        v = v.unsqueeze(0)
+                    elif type == "node reference":
+                        v = v + num_cum_node
+                    elif type == "edge reference":
+                        v = v + num_cum_edge
+                    elif type == "graph reference":
+                        v = v + num_graph
+                data_dict[k].append(v)
+            if num_relation == -1:
+                num_relation = graph.num_relation
+            elif num_relation != graph.num_relation:
+                raise ValueError("Inconsistent `num_relation` in graphs. Expect %d but got %d."
+                                 % (num_relation, graph.num_relation))
+            num_cum_node += graph.num_node
+            num_cum_edge += graph.num_edge
+            num_graph += 1
+
+        edge_list = torch.cat(edge_list)
+        edge_weight = torch.cat(edge_weight)
+        data_dict = {k: torch.cat(v) for k, v in data_dict.items()}
+
+        return cls.packed_type(edge_list, edge_weight=edge_weight, num_nodes=num_nodes, num_edges=num_edges,
+                               num_relation=num_relation, meta_dict=meta_dict, **data_dict)
+
+    def repeat(self, count):
+        """
+        Repeat this graph.
+
+        Parameters:
+            count (int): number of repetitions
+
+        Returns:
+            PackedGraph
+        """
+        edge_list = self.edge_list.repeat(count, 1)
+        edge_weight = self.edge_weight.repeat(count)
+        num_nodes = [self.num_node] * count
+        num_edges = [self.num_edge] * count
+        num_relation = self.num_relation
+
+        data_dict = {}
+        for k, v in self.data_dict.items():
+            if "graph" in self.meta_dict[k]:
+                v = v.unsqueeze(0)
+            shape = [1] * v.ndim
+            shape[0] = count
+            length = len(v)
+            v = v.repeat(shape)
+            for type in self.meta_dict[k]:
+                if type == "node reference":
+                    offsets = torch.arange(count, device=self.device) * self.num_node
+                    v = v + offsets.repeat_interleave(length)
+                elif type == "edge reference":
+                    offsets = torch.arange(count, device=self.device) * self.num_edge
+                    v = v + offsets.repeat_interleave(length)
+                elif type == "graph reference":
+                    offsets = torch.arange(count, device=self.device)
+                    v = v + offsets.repeat_interleave(length)
+            data_dict[k] = v
+
+        return self.packed_type(edge_list, edge_weight=edge_weight, num_nodes=num_nodes, num_edges=num_edges,
+                                num_relation=num_relation, meta_dict=self.meta_dict, **data_dict)
+
+    def get_edge(self, edge):
+        """
+        Get the weight of of an edge.
+
+        Parameters:
+            edge (array_like): index of shape :math:`(2,)` or :math:`(3,)`
+
+        Returns:
+            Tensor: weight of the edge
+        """
+        if len(edge) != self.edge_list.shape[1]:
+            raise ValueError("Incorrect edge index. Expect %d axes but got %d axes"
+                             % (self.edge_list.shape[1], len(edge)))
+
+        edge_index, num_match = self.match(edge)
+        return self.edge_weight[edge_index].sum()
+
+    def match(self, pattern):
+        """
+        Return all matched indexes for each pattern. Support patterns with ``-1`` as the wildcard.
+
+        Parameters:
+            pattern (array_like): index of shape :math:`(N, 2)` or :math:`(N, 3)`
+
+        Returns:
+            (LongTensor, LongTensor): matched indexes, number of matches per edge
+
+        Examples::
+
+            >>> graph = data.Graph([[0, 1], [1, 0], [1, 2], [2, 1], [2, 0], [0, 2]])
+            >>> index, num_match = graph.match([[0, -1], [1, 2]])
+            >>> assert (index == torch.tensor([0, 5, 2])).all()
+            >>> assert (num_match == torch.tensor([2, 1])).all()
+
+        """
+        if len(pattern) == 0:
+            index = num_match = torch.zeros(0, dtype=torch.long, device=self.device)
+            return index, num_match
+
+        if not hasattr(self, "edge_inverted_index"):
+            self.edge_inverted_index = {}
+        pattern = torch.as_tensor(pattern, dtype=torch.long, device=self.device)
+        if pattern.ndim == 1:
+            pattern = pattern.unsqueeze(0)
+        mask = pattern != -1
+        scale = 2 ** torch.arange(pattern.shape[-1], device=self.device)
+        query_type = (mask * scale).sum(dim=-1)
+        query_index = query_type.argsort()
+        num_query = query_type.unique(return_counts=True)[1]
+        query_ends = num_query.cumsum(0)
+        query_starts = query_ends - num_query
+        mask_set = mask[query_index[query_starts]].tolist()
+
+        type_ranges = []
+        type_orders = []
+        # get matched range for each query type
+        for i, mask in enumerate(mask_set):
+            query_type = tuple(mask)
+            type_index = query_index[query_starts[i]: query_ends[i]]
+            type_edge = pattern[type_index][:, mask]
+            if query_type not in self.edge_inverted_index:
+                self.edge_inverted_index[query_type] = self._build_edge_inverted_index(mask)
+            inverted_range, order = self.edge_inverted_index[query_type]
+            ranges = inverted_range.get(type_edge, default=0)
+            type_ranges.append(ranges)
+            type_orders.append(order)
+        ranges = torch.cat(type_ranges)
+        orders = torch.stack(type_orders)
+        types = torch.arange(len(mask_set), device=self.device)
+        types = types.repeat_interleave(num_query)
+
+        # reorder matched ranges according to the query order
+        ranges = scatter_add(ranges, query_index, dim=0, dim_size=len(pattern))
+        types = scatter_add(types, query_index, dim_size=len(pattern))
+        # convert range to indexes
+        starts, ends = ranges.t()
+        num_match = ends - starts
+        offsets = num_match.cumsum(0) - num_match
+        types = types.repeat_interleave(num_match)
+        ranges = torch.arange(num_match.sum(), device=self.device)
+        ranges = ranges + (starts - offsets).repeat_interleave(num_match)
+        index = orders[types, ranges]
+
+        return index, num_match
+
+    def _build_edge_inverted_index(self, mask):
+        keys = self.edge_list[:, mask]
+        base = torch.tensor(self.shape, device=self.device)
+        base = base[mask]
+        max = reduce(int.__mul__, base.tolist())
+        if max > torch.iinfo(torch.int64).max:
+            raise ValueError("Fail to build an inverted index table based on sorting. "
+                             "The graph is too large.")
+        scale = base.cumprod(0)
+        scale = torch.div(scale[-1], scale, rounding_mode="floor")
+        key = (keys * scale).sum(dim=-1)
+        order = key.argsort()
+        num_keys = key.unique(return_counts=True)[1]
+        ends = num_keys.cumsum(0)
+        starts = ends - num_keys
+        ranges = torch.stack([starts, ends], dim=-1)
+        keys_set = keys[order[starts]]
+        inverted_range = Dictionary(keys_set, ranges)
+        return inverted_range, order
+
+    def __getitem__(self, index):
+        # why do we check tuple?
+        # case 1: x[0, 1] is parsed as (0, 1)
+        # case 2: x[[0, 1]] is parsed as [0, 1]
+        if not isinstance(index, tuple):
+            index = (index,)
+        index = list(index)
+
+        while len(index) < 2:
+            index.append(slice(None))
+        if len(index) > 2:
+            raise ValueError("Graph has only 2 axis, but %d axis is indexed" % len(index))
+
+        if all([isinstance(axis_index, int) for axis_index in index]):
+            return self.get_edge(index)
+
+        edge_list = self.edge_list.clone()
+        for i, axis_index in enumerate(index):
+            axis_index = self._standarize_index(axis_index, self.num_node)
+            mapping = -torch.ones(self.num_node, dtype=torch.long, device=self.device)
+            mapping[axis_index] = axis_index
+            edge_list[:, i] = mapping[edge_list[:, i]]
+        edge_index = (edge_list >= 0).all(dim=-1)
+
+        return self.edge_mask(edge_index)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def batch_size(self):
+        """Batch size."""
+        return 1
+
+    def subgraph(self, index):
+        """
+        Return a subgraph based on the specified nodes.
+        Equivalent to :meth:`node_mask(index, compact=True) <node_mask>`.
+
+        Parameters:
+            index (array_like): node index
+
+        Returns:
+            Graph
+
+        See also:
+            :meth:`Graph.node_mask`
+        """
+        return self.node_mask(index, compact=True)
+
+    def data_mask(self, node_index=None, edge_index=None, graph_index=None, include=None, exclude=None):
+        data_dict, meta_dict = self.data_by_meta(include, exclude)
+        node_mapping = None
+        edge_mapping = None
+        graph_mapping = None
+        for k, v in data_dict.items():
+            for type in meta_dict[k]:
+                if type == "node" and node_index is not None:
+                    v = v[node_index]
+                elif type == "edge" and edge_index is not None:
+                    v = v[edge_index]
+                elif type == "graph" and graph_index is not None:
+                    v = v.unsqueeze(0)[graph_index]
+                elif type == "node reference" and node_index is not None:
+                    if node_mapping is None:
+                        node_mapping = self._get_mapping(node_index, self.num_node)
+                    v = node_mapping[v]
+                elif type == "edge reference" and edge_index is not None:
+                    if edge_mapping is None:
+                        edge_mapping = self._get_mapping(edge_index, self.num_edge)
+                    v = edge_mapping[v]
+                elif type == "graph reference" and graph_index is not None:
+                    if graph_mapping is None:
+                        graph_mapping = self._get_mapping(graph_index, self.batch_size)
+                    v = graph_mapping[v]
+            data_dict[k] = v
+
+        return data_dict, meta_dict
+
+    def node_mask(self, index, compact=False):
+        """
+        Return a masked graph based on the specified nodes.
+
+        This function can also be used to re-order the nodes.
+
+        Parameters:
+            index (array_like): node index
+            compact (bool, optional): compact node ids or not
+
+        Returns:
+            Graph
+
+        Examples::
+
+            >>> graph = data.Graph.from_dense(torch.eye(3))
+            >>> assert graph.node_mask([1, 2]).adjacency.shape == (3, 3)
+            >>> assert graph.node_mask([1, 2], compact=True).adjacency.shape == (2, 2)
+
+        """
+        index = self._standarize_index(index, self.num_node)
+        mapping = -torch.ones(self.num_node, dtype=torch.long, device=self.device)
+        if compact:
+            mapping[index] = torch.arange(len(index), device=self.device)
+            num_node = len(index)
+        else:
+            mapping[index] = index
+            num_node = self.num_node
+
+        edge_list = self.edge_list.clone()
+        edge_list[:, :2] = mapping[edge_list[:, :2]]
+        edge_index = (edge_list[:, :2] >= 0).all(dim=-1)
+
+        if compact:
+            data_dict, meta_dict = self.data_mask(index, edge_index)
+        else:
+            data_dict, meta_dict = self.data_mask(edge_index=edge_index)
+
+        return type(self)(edge_list[edge_index], edge_weight=self.edge_weight[edge_index], num_node=num_node,
+                          num_relation=self.num_relation, meta_dict=meta_dict, **data_dict)
+
+    def compact(self):
+        """
+        Remove isolated nodes and compact node ids.
+
+        Returns:
+            Graph
+        """
+        index = self.degree_out + self.degree_in > 0
+        return self.subgraph(index)
+
+    def edge_mask(self, index):
+        """
+        Return a masked graph based on the specified edges.
+
+        This function can also be used to re-order the edges.
+
+        Parameters:
+            index (array_like): edge index
+
+        Returns:
+            Graph
+        """
+        index = self._standarize_index(index, self.num_edge)
+        data_dict, meta_dict = self.data_mask(edge_index=index)
+
+        return type(self)(self.edge_list[index], edge_weight=self.edge_weight[index], num_node=self.num_node,
+                          num_relation=self.num_relation, meta_dict=meta_dict, **data_dict)
+
+    def line_graph(self):
+        """
+        Construct a line graph of this graph.
+        The node feature of the line graph is inherited from the edge feature of the original graph.
+
+        In the line graph, each node corresponds to an edge in the original graph.
+        For a pair of edges (a, b) and (b, c) that share the same intermediate node in the original graph,
+        there is a directed edge (a, b) -> (b, c) in the line graph.
+
+        Returns:
+            Graph
+        """
+        node_in, node_out = self.edge_list.t()[:2]
+        edge_index = torch.arange(self.num_edge, device=self.device)
+        edge_in = edge_index[node_out.argsort()]
+        edge_out = edge_index[node_in.argsort()]
+
+        degree_in = node_in.bincount(minlength=self.num_node)
+        degree_out = node_out.bincount(minlength=self.num_node)
+        size = degree_out * degree_in
+        starts = (size.cumsum(0) - size).repeat_interleave(size)
+        range = torch.arange(size.sum(), device=self.device)
+        # each node u has degree_out[u] * degree_in[u] local edges
+        local_index = range - starts
+        local_inner_size = degree_in.repeat_interleave(size)
+        edge_in_offset = (degree_out.cumsum(0) - degree_out).repeat_interleave(size)
+        edge_out_offset = (degree_in.cumsum(0) - degree_in).repeat_interleave(size)
+        edge_in_index = torch.div(local_index, local_inner_size, rounding_mode="floor") + edge_in_offset
+        edge_out_index = local_index % local_inner_size + edge_out_offset
+
+        edge_in = edge_in[edge_in_index]
+        edge_out = edge_out[edge_out_index]
+        edge_list = torch.stack([edge_in, edge_out], dim=-1)
+        node_feature = getattr(self, "edge_feature", None)
+        num_node = self.num_edge
+        num_edge = size.sum()
+
+        return Graph(edge_list, num_node=num_node, num_edge=num_edge, node_feature=node_feature)
+
+    def full(self):
+        """
+        Return a fully connected graph over the nodes.
+
+        Returns:
+            Graph
+        """
+        index = torch.arange(self.num_node, device=self.device)
+        if self.num_relation:
+            edge_list = torch.meshgrid(index, index, torch.arange(self.num_relation, device=self.device))
+        else:
+            edge_list = torch.meshgrid(index, index)
+        edge_list = torch.stack(edge_list).flatten(1)
+        edge_weight = torch.ones(len(edge_list))
+
+        data_dict, meta_dict = self.data_by_meta(exclude="edge")
+
+        return type(self)(edge_list, edge_weight=edge_weight, num_node=self.num_node, num_relation=self.num_relation,
+                          meta_dict=meta_dict, **data_dict)
+
+    def directed(self, order=None):
+        """
+        Mask the edges to create a directed graph.
+        Edges that go from a node index to a larger or equal node index will be kept.
+
+        Parameters:
+            order (Tensor, optional): topological order of the nodes
+        """
+        node_in, node_out = self.edge_list.t()[:2]
+        if order is not None:
+            edge_index = order[node_in] <= order[node_out]
+        else:
+            edge_index = node_in <= node_out
+
+        return self.edge_mask(edge_index)
+
+    def undirected(self, add_inverse=False):
+        """
+        Flip all the edges to create an undirected graph.
+
+        For knowledge graphs, the flipped edges can either have the original relation or an inverse relation.
+        The inverse relation for relation :math:`r` is defined as :math:`|R| + r`.
+
+        Parameters:
+            add_inverse (bool, optional): whether to use inverse relations for flipped edges
+        """
+        edge_list = self.edge_list.clone()
+        edge_list[:, :2] = edge_list[:, :2].flip(1)
+        num_relation = self.num_relation
+        if num_relation and add_inverse:
+            edge_list[:, 2] += num_relation
+            num_relation = num_relation * 2
+        edge_list = torch.stack([self.edge_list, edge_list], dim=1).flatten(0, 1)
+
+        index = torch.arange(self.num_edge, device=self.device).unsqueeze(-1).expand(-1, 2).flatten()
+        data_dict, meta_dict = self.data_mask(edge_index=index)
+
+        return type(self)(edge_list, edge_weight=self.edge_weight[index], num_node=self.num_node,
+                          num_relation=num_relation, meta_dict=meta_dict, **data_dict)
+
+    @utils.cached_property
+    def adjacency(self):
+        """
+        Adjacency matrix of this graph.
+
+        If :attr:`num_relation` is specified, a sparse tensor of shape :math:`(|V|, |V|, num\_relation)` will be
+        returned.
+        Otherwise, a sparse tensor of shape :math:`(|V|, |V|)` will be returned.
+        """
+        return utils.sparse_coo_tensor(self.edge_list.t(), self.edge_weight, self.shape)
+
+    _tensor_names = ["edge_list", "edge_weight", "num_node", "num_relation", "edge_feature"]
+
+    def to_tensors(self):
+        edge_feature = getattr(self, "edge_feature", torch.tensor(0, device=self.device))
+        return self.edge_list, self.edge_weight, self.num_node, self.num_relation, edge_feature
+
+    @classmethod
+    def from_tensors(cls, tensors):
+        edge_list, edge_weight, num_node, num_relation, edge_feature = tensors
+        if edge_feature.ndim == 0:
+            edge_feature = None
+        return cls(edge_list, edge_weight, num_node, num_relation, edge_feature=edge_feature)
+
+    @property
+    def node2graph(self):
+        """Node id to graph id mapping."""
+        return torch.zeros(self.num_node, dtype=torch.long, device=self.device)
+
+    @property
+    def edge2graph(self):
+        """Edge id to graph id mapping."""
+        return torch.zeros(self.num_edge, dtype=torch.long, device=self.device)
+
+    @utils.cached_property
+    def degree_out(self):
+        """
+        Weighted number of edges containing each node as output.
+
+        Note this is the **in-degree** in graph theory.
+        """
+        return scatter_add(self.edge_weight, self.edge_list[:, 1], dim_size=self.num_node)
+
+    @utils.cached_property
+    def degree_in(self):
+        """
+        Weighted number of edges containing each node as input.
+
+        Note this is the **out-degree** in graph theory.
+        """
+        return scatter_add(self.edge_weight, self.edge_list[:, 0], dim_size=self.num_node)
+
+    @property
+    def edge_list(self):
+        """List of edges."""
+        return self._edge_list
+
+    @property
+    def edge_weight(self):
+        """Edge weights."""
+        return self._edge_weight
+
+    @property
+    def device(self):
+        """Device."""
+        return self.edge_list.device
+
+    @property
+    def requires_grad(self):
+        return self.edge_weight.requires_grad
+
+    @property
+    def grad(self):
+        return self.edge_weight.grad
+
+    @property
+    def data(self):
+        return self
+
+    def requires_grad_(self):
+        self.edge_weight.requires_grad_()
+        return self
+
+    def size(self, dim=None):
+        if self.num_relation:
+            size = torch.Size((self.num_node, self.num_node, self.num_relation))
+        else:
+            size = torch.Size((self.num_node, self.num_node))
+        if dim is None:
+            return size
+        return size[dim]
+
+    @property
+    def shape(self):
+        return self.size()
+
+    def copy_(self, src):
+        """
+        Copy data from ``src`` into ``self`` and return ``self``.
+
+        The ``src`` graph must have the same set of attributes as ``self``.
+        """
+        self.edge_list.copy_(src.edge_list)
+        self.edge_weight.copy_(src.edge_weight)
+        self.num_node.copy_(src.num_node)
+        self.num_edge.copy_(src.num_edge)
+        if self.num_relation is not None:
+            self.num_relation.copy_(src.num_relation)
+
+        keys = set(self.data_dict.keys())
+        src_keys = set(src.data_dict.keys())
+        if keys != src_keys:
+            raise RuntimeError("Attributes mismatch. Trying to assign attributes %s, "
+                               "but current graph has attributes %s" % (src_keys, keys))
+        for k, v in self.data_dict.items():
+            v.copy_(src.data_dict[k])
+
+        return self
+
+    def detach(self):
+        """
+        Detach this graph.
+        """
+        return type(self)(self.edge_list.detach(), edge_weight=self.edge_weight.detach(),
+                          num_node=self.num_node, num_relation=self.num_relation,
+                          meta_dict=self.meta_dict, **utils.detach(self.data_dict))
+
+    def clone(self):
+        """
+        Clone this graph.
+        """
+        return type(self)(self.edge_list.clone(), edge_weight=self.edge_weight.clone(),
+                          num_node=self.num_node, num_relation=self.num_relation,
+                          meta_dict=self.meta_dict, **utils.clone(self.data_dict))
+
+    def cuda(self, *args, **kwargs):
+        """
+        Return a copy of this graph in CUDA memory.
+
+        This is a non-op if the graph is already on the correct device.
+        """
+        edge_list = self.edge_list.cuda(*args, **kwargs)
+
+        if edge_list is self.edge_list:
+            return self
+        else:
+            return type(self)(edge_list, edge_weight=self.edge_weight,
+                              num_node=self.num_node, num_relation=self.num_relation,
+                              meta_dict=self.meta_dict, **utils.cuda(self.data_dict, *args, **kwargs))
+
+    def cpu(self):
+        """
+        Return a copy of this graph in CPU memory.
+
+        This is a non-op if the graph is already in CPU memory.
+        """
+        edge_list = self.edge_list.cpu()
+
+        if edge_list is self.edge_list:
+            return self
+        else:
+            return type(self)(edge_list, edge_weight=self.edge_weight, num_node=self.num_node,
+                              num_relation=self.num_relation, meta_dict=self.meta_dict, **utils.cpu(self.data_dict))
+
+    def to(self, device, *args, **kwargs):
+        """
+        Return a copy of this graph on the given device.
+        """
+        device = torch.device(device)
+        if device.type == "cpu":
+            return self.cpu(*args, **kwargs)
+        else:
+            return self.cuda(device, *args, **kwargs)
+
+    def __repr__(self):
+        fields = ["num_node=%d" % self.num_node, "num_edge=%d" % self.num_edge]
+        if self.num_relation is not None:
+            fields.append("num_relation=%d" % self.num_relation)
+        if self.device.type != "cpu":
+            fields.append("device='%s'" % self.device)
+        return "%s(%s)" % (self.__class__.__name__, ", ".join(fields))
+
+    def visualize(self, title=None, save_file=None, figure_size=(3, 3), ax=None, layout="spring"):
+        """
+        Visualize this graph with matplotlib.
+
+        Parameters:
+            title (str, optional): title for this graph
+            save_file (str, optional): ``png`` or ``pdf`` file to save visualization.
+                If not provided, show the figure in window.
+            figure_size (tuple of int, optional): width and height of the figure
+            ax (matplotlib.axes.Axes, optional): axis to plot the figure
+            layout (str, optional): graph layout
+
+        See also:
+            `NetworkX graph layout`_
+
+            .. _NetworkX graph layout:
+                https://networkx.github.io/documentation/stable/reference/drawing.html#module-networkx.drawing.layout
+        """
+        is_root = ax is None
+        if ax is None:
+            fig = plt.figure(figsize=figure_size)
+            if title is not None:
+                ax = plt.gca()
+            else:
+                ax = fig.add_axes([0, 0, 1, 1])
+        if title is not None:
+            ax.set_title(title)
+
+        edge_list = self.edge_list[:, :2].tolist()
+        G = nx.DiGraph(edge_list)
+        G.add_nodes_from(range(self.num_node))
+        if hasattr(nx, "%s_layout" % layout):
+            func = getattr(nx, "%s_layout" % layout)
+        else:
+            raise ValueError("Unknown networkx layout `%s`" % layout)
+        if layout == "spring" or layout == "random":
+            pos = func(G, seed=0)
+        else:
+            pos = func(G)
+        nx.draw_networkx(G, pos, ax=ax)
+        if self.num_relation:
+            edge_labels = self.edge_list[:, 2].tolist()
+            edge_labels = {tuple(e): l for e, l in zip(edge_list, edge_labels)}
+            nx.draw_networkx_edge_labels(G, pos, edge_labels, ax=ax)
+        ax.set_frame_on(False)
+
+        if is_root:
+            if save_file:
+                fig.savefig(save_file)
+            else:
+                fig.show()
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        return NotImplemented
+
+    def __getstate__(self):
+        state = {}
+        cls = self.__class__
+        for k, v in self.__dict__.items():
+            # do not pickle property / cached property
+            if hasattr(cls, k) and isinstance(getattr(cls, k), property):
+                continue
+            state[k] = v
+        return state
+
+
+class PackedGraph(Graph):
+    """
+    Container for sparse graphs with variadic sizes.
+
+    To create a PackedGraph from Graph objects
+
+        >>> batch = data.Graph.pack(graphs)
+
+    To retrieve Graph objects from a PackedGraph
+
+        >>> graphs = batch.unpack()
+
+    .. warning::
+        
+        Edges of the same graph are guaranteed to be consecutive in the edge list.
+        However, this class doesn't enforce any order on the edges.
+
+    Parameters:
+        edge_list (array_like, optional): list of edges of shape :math:`(|E|, 2)` or :math:`(|E|, 3)`.
+            Each tuple is (node_in, node_out) or (node_in, node_out, relation).
+        edge_weight (array_like, optional): edge weights of shape :math:`(|E|,)`
+        num_nodes (array_like, optional): number of nodes in each graph
+            By default, it will be inferred from the largest id in `edge_list`
+        num_edges (array_like, optional): number of edges in each graph
+        num_relation (int, optional): number of relations
+        node_feature (array_like, optional): node features of shape :math:`(|V|, ...)`
+        edge_feature (array_like, optional): edge features of shape :math:`(|E|, ...)`
+        offsets (array_like, optional): node id offsets of shape :math:`(|E|,)`.
+            If not provided, nodes in `edge_list` should be relative index, i.e., the index in each graph.
+            If provided, nodes in `edge_list` should be absolute index, i.e., the index in the packed graph.
+    """
+
+    unpacked_type = Graph
+
+    def __init__(self, edge_list=None, edge_weight=None, num_nodes=None, num_edges=None, num_relation=None,
+                 offsets=None, **kwargs):
+        edge_list, num_nodes, num_edges, num_cum_nodes, num_cum_edges, offsets = \
+            self._get_cumulative(edge_list, num_nodes, num_edges, offsets)
+
+        if offsets is None:
+            offsets = self._get_offsets(num_nodes, num_edges, num_cum_nodes)
+            edge_list = edge_list.clone()
+            edge_list[:, :2] += offsets.unsqueeze(-1)
+
+        num_node = num_nodes.sum()
+        if (edge_list[:, :2] >= num_node).any():
+            raise ValueError("Sum of `num_nodes` is %d, but found %d in `edge_list`" %
+                             (num_node, edge_list[:, :2].max()))
+
+        self._offsets = offsets
+        self.num_nodes = num_nodes
+        self.num_edges = num_edges
+        self.num_cum_nodes = num_cum_nodes
+        self.num_cum_edges = num_cum_edges
+
+        super(PackedGraph, self).__init__(edge_list, edge_weight=edge_weight, num_node=num_node,
+                                          num_relation=num_relation, **kwargs)
+
+    def _get_offsets(self, num_nodes=None, num_edges=None, num_cum_nodes=None, num_cum_edges=None):
+        if num_nodes is None:
+            prepend = torch.tensor([0], device=self.device)
+            num_nodes = torch.diff(num_cum_nodes, prepend=prepend)
+        if num_edges is None:
+            prepend = torch.tensor([0], device=self.device)
+            num_edges = torch.diff(num_cum_edges, prepend=prepend)
+        if num_cum_nodes is None:
+            num_cum_nodes = num_nodes.cumsum(0)
+        return (num_cum_nodes - num_nodes).repeat_interleave(num_edges)
+
+    def merge(self, graph2graph):
+        """
+        Merge multiple graphs into a single graph.
+
+        Parameters:
+            graph2graph (array_like): ID of the new graph each graph belongs to
+        """
+        graph2graph = torch.as_tensor(graph2graph, dtype=torch.long, device=self.device)
+        # coalesce arbitrary graph IDs to [0, n)
+        _, graph2graph = torch.unique(graph2graph, return_inverse=True)
+
+        graph_key = graph2graph * self.batch_size + torch.arange(self.batch_size, device=self.device)
+        graph_index = graph_key.argsort()
+        graph = self.subbatch(graph_index)
+        graph2graph = graph2graph[graph_index]
+
+        num_graph = graph2graph[-1] + 1
+        num_nodes = scatter_add(graph.num_nodes, graph2graph, dim_size=num_graph)
+        num_edges = scatter_add(graph.num_edges, graph2graph, dim_size=num_graph)
+        offsets = self._get_offsets(num_nodes, num_edges)
+
+        data_dict, meta_dict = graph.data_mask(exclude="graph")
+
+        return type(self)(graph.edge_list, edge_weight=graph.edge_weight, num_nodes=num_nodes,
+                          num_edges=num_edges, num_relation=graph.num_relation, offsets=offsets,
+                          meta_dict=meta_dict, **data_dict)
+
+    def unpack(self):
+        """
+        Unpack this packed graph into a list of graphs.
+
+        Returns:
+            list of Graph
+        """
+        graphs = []
+        for i in range(self.batch_size):
+            graphs.append(self.get_item(i))
+        return graphs
+
+    def __iter__(self):
+        self._iter_index = 0
+        return self
+
+    def __next__(self):
+        if self._iter_index < self.batch_size:
+            item = self[self._iter_index]
+            self._iter_index += 1
+            return item
+        raise StopIteration
+
+    def _check_attribute(self, key, value):
+        for type in self._meta_contexts:
+            if "reference" in type:
+                if value.dtype != torch.long:
+                    raise TypeError("Tensors used as reference must be long tensors")
+            if type == "node":
+                if len(value) != self.num_node:
+                    raise ValueError("Expect node attribute `%s` to have shape (%d, *), but found %s" %
+                                     (key, self.num_node, value.shape))
+            elif type == "edge":
+                if len(value) != self.num_edge:
+                    raise ValueError("Expect edge attribute `%s` to have shape (%d, *), but found %s" %
+                                     (key, self.num_edge, value.shape))
+            elif type == "graph":
+                if len(value) != self.batch_size:
+                    raise ValueError("Expect graph attribute `%s` to have shape (%d, *), but found %s" %
+                                     (key, self.batch_size, value.shape))
+            elif type == "node reference":
+                is_valid = (value >= -1) & (value < self.num_node)
+                if not is_valid.all():
+                    error_value = value[~is_valid]
+                    raise ValueError("Expect node reference in [-1, %d), but found %d" %
+                                     (self.num_node, error_value[0]))
+            elif type == "edge reference":
+                is_valid = (value >= -1) & (value < self.num_edge)
+                if not is_valid.all():
+                    error_value = value[~is_valid]
+                    raise ValueError("Expect edge reference in [-1, %d), but found %d" %
+                                     (self.num_edge, error_value[0]))
+            elif type == "graph reference":
+                is_valid = (value >= -1) & (value < self.batch_size)
+                if not is_valid.all():
+                    error_value = value[~is_valid]
+                    raise ValueError("Expect graph reference in [-1, %d), but found %d" %
+                                     (self.batch_size, error_value[0]))
+
+    def unpack_data(self, data, type="auto"):
+        """
+        Unpack node or edge data according to the packed graph.
+
+        Parameters:
+            data (Tensor): data to unpack
+            type (str, optional): data type. Can be ``auto``, ``node``, or ``edge``.
+
+        Returns:
+            list of Tensor
+        """
+        if type == "auto":
+            if self.num_node == self.num_edge:
+                raise ValueError("Ambiguous type. Please specify either `node` or `edge`")
+            if len(data) == self.num_node:
+                type = "node"
+            elif len(data) == self.num_edge:
+                type = "edge"
+            else:
+                raise ValueError("Graph has %d nodes and %d edges, but data has %d entries" %
+                                 (self.num_node, self.num_edge, len(data)))
+        data_list = []
+        if type == "node":
+            for i in range(self.batch_size):
+                data_list.append(data[self.num_cum_nodes[i] - self.num_nodes[i]: self.num_cum_nodes[i]])
+        elif type == "edge":
+            for i in range(self.batch_size):
+                data_list.append(data[self.num_cum_edges[i] - self.num_edges[i]: self.num_cum_edges[i]])
+
+        return data_list
+
+    def repeat(self, count):
+        """
+        Repeat this packed graph. This function behaves similarly to `torch.Tensor.repeat`_.
+
+        .. _torch.Tensor.repeat:
+            https://pytorch.org/docs/stable/generated/torch.Tensor.repeat.html
+
+        Parameters:
+            count (int): number of repetitions
+
+        Returns:
+            PackedGraph
+        """
+        num_nodes = self.num_nodes.repeat(count)
+        num_edges = self.num_edges.repeat(count)
+        offsets = self._get_offsets(num_nodes, num_edges)
+        edge_list = self.edge_list.repeat(count, 1)
+        edge_list[:, :2] += (offsets - self._offsets.repeat(count)).unsqueeze(-1)
+
+        data_dict = {}
+        for k, v in self.data_dict.items():
+            shape = [1] * v.ndim
+            shape[0] = count
+            length = len(v)
+            v = v.repeat(shape)
+            for _type in self.meta_dict[k]:
+                if _type == "node reference":
+                    pack_offsets = torch.arange(count, device=self.device) * self.num_node
+                    v = v + pack_offsets.repeat_interleave(length)
+                elif _type == "edge reference":
+                    pack_offsets = torch.arange(count, device=self.device) * self.num_edge
+                    v = v + pack_offsets.repeat_interleave(length)
+                elif _type == "graph reference":
+                    pack_offsets = torch.arange(count, device=self.device) * self.batch_size
+                    v = v + pack_offsets.repeat_interleave(length)
+            data_dict[k] = v
+
+        return type(self)(edge_list, edge_weight=self.edge_weight.repeat(count),
+                          num_nodes=num_nodes, num_edges=num_edges, num_relation=self.num_relation,
+                          offsets=offsets, meta_dict=self.meta_dict, **data_dict)
+
+    def repeat_interleave(self, repeats):
+        """
+        Repeat this packed graph. This function behaves similarly to `torch.repeat_interleave`_.
+
+        .. _torch.repeat_interleave:
+            https://pytorch.org/docs/stable/generated/torch.repeat_interleave.html
+
+        Parameters:
+            repeats (Tensor or int): number of repetitions for each graph
+
+        Returns:
+            PackedGraph
+        """
+        repeats = torch.as_tensor(repeats, dtype=torch.long, device=self.device)
+        if repeats.numel() == 1:
+            repeats = repeats * torch.ones(self.batch_size, dtype=torch.long, device=self.device)
+        num_nodes = self.num_nodes.repeat_interleave(repeats)
+        num_edges = self.num_edges.repeat_interleave(repeats)
+        num_cum_nodes = num_nodes.cumsum(0)
+        num_cum_edges = num_edges.cumsum(0)
+        num_node = num_nodes.sum()
+        num_edge = num_edges.sum()
+        batch_size = repeats.sum()
+        num_graphs = torch.ones(batch_size, device=self.device)
+
+        # special case 1: graphs[i] may have no node or no edge
+        # special case 2: repeats[i] may be 0
+        cum_repeats_shifted = repeats.cumsum(0) - repeats
+        graph_mask = cum_repeats_shifted < batch_size
+        cum_repeats_shifted = cum_repeats_shifted[graph_mask]
+
+        index = num_cum_nodes - num_nodes
+        index = torch.cat([index, index[cum_repeats_shifted]])
+        value = torch.cat([-num_nodes, self.num_nodes[graph_mask]])
+        mask = index < num_node
+        node_index = scatter_add(value[mask], index[mask], dim_size=num_node)
+        node_index = (node_index + 1).cumsum(0) - 1
+
+        index = num_cum_edges - num_edges
+        index = torch.cat([index, index[cum_repeats_shifted]])
+        value = torch.cat([-num_edges, self.num_edges[graph_mask]])
+        mask = index < num_edge
+        edge_index = scatter_add(value[mask], index[mask], dim_size=num_edge)
+        edge_index = (edge_index + 1).cumsum(0) - 1
+
+        graph_index = torch.repeat_interleave(repeats)
+
+        offsets = self._get_offsets(num_nodes, num_edges)
+        edge_list = self.edge_list[edge_index]
+        edge_list[:, :2] += (offsets - self._offsets[edge_index]).unsqueeze(-1)
+
+        node_offsets = None
+        edge_offsets = None
+        graph_offsets = None
+        data_dict = {}
+        for k, v in self.data_dict.items():
+            num_xs = None
+            pack_offsets = None
+            for _type in self.meta_dict[k]:
+                if _type == "node":
+                    v = v[node_index]
+                    num_xs = num_nodes
+                elif _type == "edge":
+                    v = v[edge_index]
+                    num_xs = num_edges
+                elif _type == "graph":
+                    v = v[graph_index]
+                    num_xs = num_graphs
+                elif _type == "node reference":
+                    if node_offsets is None:
+                        node_offsets = self._get_repeat_pack_offsets(self.num_nodes, repeats)
+                    pack_offsets = node_offsets
+                elif _type == "edge reference":
+                    if edge_offsets is None:
+                        edge_offsets = self._get_repeat_pack_offsets(self.num_edges, repeats)
+                    pack_offsets = edge_offsets
+                elif _type == "graph reference":
+                    if graph_offsets is None:
+                        graph_offsets = self._get_repeat_pack_offsets(num_graphs, repeats)
+                    pack_offsets = graph_offsets
+            # add offsets to make references point to indexes in their own graph
+            if num_xs is not None and pack_offsets is not None:
+                v = v + pack_offsets.repeat_interleave(num_xs)
+            data_dict[k] = v
+
+        return type(self)(edge_list, edge_weight=self.edge_weight[edge_index],
+                          num_nodes=num_nodes, num_edges=num_edges, num_relation=self.num_relation,
+                          offsets=offsets, meta_dict=self.meta_dict, **data_dict)
+
+    def get_item(self, index):
+        """
+        Get the i-th graph from this packed graph.
+
+        Parameters:
+            index (int): graph index
+
+        Returns:
+            Graph
+        """
+        node_index = torch.arange(self.num_cum_nodes[index] - self.num_nodes[index], self.num_cum_nodes[index],
+                                  device=self.device)
+        edge_index = torch.arange(self.num_cum_edges[index] - self.num_edges[index], self.num_cum_edges[index],
+                                  device=self.device)
+        graph_index = index
+        edge_list = self.edge_list[edge_index].clone()
+        edge_list[:, :2] -= self._offsets[edge_index].unsqueeze(-1)
+        data_dict, meta_dict = self.data_mask(node_index, edge_index, graph_index=graph_index)
+
+        return self.unpacked_type(edge_list, edge_weight=self.edge_weight[edge_index], num_node=self.num_nodes[index],
+                                  num_relation=self.num_relation, meta_dict=meta_dict, **data_dict)
+
+    def _get_cumulative(self, edge_list, num_nodes, num_edges, offsets):
+        if edge_list is None:
+            raise ValueError("`edge_list` should be provided")
+        if num_edges is None:
+            raise ValueError("`num_edges` should be provided")
+
+        edge_list = torch.as_tensor(edge_list)
+        num_edges = torch.as_tensor(num_edges, device=edge_list.device)
+        num_edge = num_edges.sum()
+        if num_edge != len(edge_list):
+            raise ValueError("Sum of `num_edges` is %d, but found %d edges in `edge_list`" % (num_edge, len(edge_list)))
+        num_cum_edges = num_edges.cumsum(0)
+
+        if offsets is None:
+            _edge_list = edge_list
+        else:
+            offsets = torch.as_tensor(offsets, device=edge_list.device)
+            _edge_list = edge_list.clone()
+            _edge_list[:, :2] -= offsets.unsqueeze(-1)
+        if num_nodes is None:
+            num_nodes = []
+            for num_edge, num_cum_edge in zip(num_edges, num_cum_edges):
+                num_nodes.append(self._maybe_num_node(_edge_list[num_cum_edge - num_edge: num_cum_edge]))
+        num_nodes = torch.as_tensor(num_nodes, device=edge_list.device)
+        num_cum_nodes = num_nodes.cumsum(0)
+
+        return edge_list, num_nodes, num_edges, num_cum_nodes, num_cum_edges, offsets
+
+    def _get_num_xs(self, index, num_cum_xs):
+        x = torch.zeros(num_cum_xs[-1], dtype=torch.long, device=self.device)
+        x[index] = 1
+        num_cum_indexes = x.cumsum(0)
+        num_cum_indexes = torch.cat([torch.zeros(1, dtype=torch.long, device=self.device), num_cum_indexes])
+        new_num_cum_xs = num_cum_indexes[num_cum_xs]
+        prepend = torch.zeros(1, dtype=torch.long, device=self.device)
+        new_num_xs = torch.diff(new_num_cum_xs, prepend=prepend)
+        return new_num_xs
+
+    def data_mask(self, node_index=None, edge_index=None, graph_index=None, include=None, exclude=None):
+        data_dict, meta_dict = self.data_by_meta(include, exclude)
+        node_mapping = None
+        edge_mapping = None
+        graph_mapping = None
+        for k, v in data_dict.items():
+            for type in meta_dict[k]:
+                if type == "node" and node_index is not None:
+                    v = v[node_index]
+                elif type == "edge" and edge_index is not None:
+                    v = v[edge_index]
+                elif type == "graph" and graph_index is not None:
+                    v = v[graph_index]
+                elif type == "node reference" and node_index is not None:
+                    if node_mapping is None:
+                        node_mapping = self._get_mapping(node_index, self.num_node)
+                    v = node_mapping[v]
+                elif type == "edge reference" and edge_index is not None:
+                    if edge_mapping is None:
+                        edge_mapping = self._get_mapping(edge_index, self.num_edge)
+                    v = edge_mapping[v]
+                elif type == "graph reference" and graph_index is not None:
+                    if graph_mapping is None:
+                        graph_mapping = self._get_mapping(graph_index, self.batch_size)
+                    v = graph_mapping[v]
+            data_dict[k] = v
+
+        return data_dict, meta_dict
+
+    def __getitem__(self, index):
+        # why do we check tuple?
+        # case 1: x[0, 1] is parsed as (0, 1)
+        # case 2: x[[0, 1]] is parsed as [0, 1]
+        if not isinstance(index, tuple):
+            index = (index,)
+
+        if isinstance(index[0], int):
+            item = self.get_item(index[0])
+            if len(index) > 1:
+                item = item[index[1:]]
+            return item
+        if len(index) > 1:
+            raise ValueError("Complex indexing is not supported for PackedGraph")
+
+        index = self._standarize_index(index[0], self.batch_size)
+        count = index.bincount(minlength=self.batch_size)
+        if self.batch_size > 0 and count.max() > 1:
+            graph = self.repeat_interleave(count)
+            index_order = index.argsort()
+            order = torch.zeros_like(index)
+            order[index_order] = torch.arange(len(index), dtype=torch.long, device=self.device)
+            return graph.subbatch(order)
+
+        return self.subbatch(index)
+
+    def __len__(self):
+        return len(self.num_nodes)
+
+    def full(self):
+        """
+        Return a pack of fully connected graphs.
+
+        This is useful for computing node-pair-wise features.
+        The computation can be implemented as message passing over a fully connected graph.
+
+        Returns:
+            PackedGraph
+        """
+        # TODO: more efficient implementation?
+        graphs = self.unpack()
+        graphs = [graph.full() for graph in graphs]
+        return graphs[0].pack(graphs)
+
+    @utils.cached_property
+    def node2graph(self):
+        """Node id to graph id mapping."""
+        node2graph = torch.repeat_interleave(self.num_nodes)
+        return node2graph
+
+    @utils.cached_property
+    def edge2graph(self):
+        """Edge id to graph id mapping."""
+        edge2graph = torch.repeat_interleave(self.num_edges)
+        return edge2graph
+
+    @property
+    def batch_size(self):
+        """Batch size."""
+        return len(self.num_nodes)
+
+    def node_mask(self, index, compact=False):
+        """
+        Return a masked packed graph based on the specified nodes.
+
+        Note the compact option is only applied to node ids but not graph ids.
+        To generate compact graph ids, use :meth:`subbatch`.
+
+        Parameters:
+            index (array_like): node index
+            compact (bool, optional): compact node ids or not
+
+        Returns:
+            PackedGraph
+        """
+        index = self._standarize_index(index, self.num_node)
+        mapping = -torch.ones(self.num_node, dtype=torch.long, device=self.device)
+        if compact:
+            mapping[index] = torch.arange(len(index), device=self.device)
+            num_nodes = self._get_num_xs(index, self.num_cum_nodes)
+            offsets = self._get_offsets(num_nodes, self.num_edges)
+        else:
+            mapping[index] = index
+            num_nodes = self.num_nodes
+            offsets = self._offsets
+
+        edge_list = self.edge_list.clone()
+        edge_list[:, :2] = mapping[edge_list[:, :2]]
+        edge_index = (edge_list[:, :2] >= 0).all(dim=-1)
+        num_edges = self._get_num_xs(edge_index, self.num_cum_edges)
+
+        if compact:
+            data_dict, meta_dict = self.data_mask(index, edge_index)
+        else:
+            data_dict, meta_dict = self.data_mask(edge_index=edge_index)
+
+        return type(self)(edge_list[edge_index], edge_weight=self.edge_weight[edge_index], num_nodes=num_nodes,
+                          num_edges=num_edges, num_relation=self.num_relation, offsets=offsets[edge_index],
+                          meta_dict=meta_dict, **data_dict)
+
+    def edge_mask(self, index):
+        """
+        Return a masked packed graph based on the specified edges.
+
+        Parameters:
+            index (array_like): edge index
+
+        Returns:
+            PackedGraph
+        """
+        index = self._standarize_index(index, self.num_edge)
+        data_dict, meta_dict = self.data_mask(edge_index=index)
+        num_edges = self._get_num_xs(index, self.num_cum_edges)
+
+        return type(self)(self.edge_list[index], edge_weight=self.edge_weight[index], num_nodes=self.num_nodes,
+                          num_edges=num_edges, num_relation=self.num_relation, offsets=self._offsets[index],
+                          meta_dict=meta_dict, **data_dict)
+
+    def graph_mask(self, index, compact=False):
+        """
+        Return a masked packed graph based on the specified graphs.
+
+        This function can also be used to re-order the graphs.
+
+        Parameters:
+            index (array_like): graph index
+            compact (bool, optional): compact graph ids or not
+
+        Returns:
+            PackedGraph
+        """
+        index = self._standarize_index(index, self.batch_size)
+        graph_mapping = -torch.ones(self.batch_size, dtype=torch.long, device=self.device)
+        graph_mapping[index] = torch.arange(len(index), device=self.device)
+
+        node_index = graph_mapping[self.node2graph] >= 0
+        node_index = self._standarize_index(node_index, self.num_node)
+        mapping = -torch.ones(self.num_node, dtype=torch.long, device=self.device)
+        if compact:
+            key = graph_mapping[self.node2graph[node_index]] * self.num_node + node_index
+            order = key.argsort()
+            node_index = node_index[order]
+            mapping[node_index] = torch.arange(len(node_index), device=self.device)
+            num_nodes = self.num_nodes[index]
+        else:
+            mapping[node_index] = node_index
+            num_nodes = torch.zeros_like(self.num_nodes)
+            num_nodes[index] = self.num_nodes[index]
+
+        edge_list = self.edge_list.clone()
+        edge_list[:, :2] = mapping[edge_list[:, :2]]
+        edge_index = (edge_list[:, :2] >= 0).all(dim=-1)
+        edge_index = self._standarize_index(edge_index, self.num_edge)
+        if compact:
+            key = graph_mapping[self.edge2graph[edge_index]] * self.num_edge + edge_index
+            order = key.argsort()
+            edge_index = edge_index[order]
+            num_edges = self.num_edges[index]
+        else:
+            num_edges = torch.zeros_like(self.num_edges)
+            num_edges[index] = self.num_edges[index]
+        offsets = self._get_offsets(num_nodes, num_edges)
+
+        if compact:
+            data_dict, meta_dict = self.data_mask(node_index, edge_index, graph_index=index)
+        else:
+            data_dict, meta_dict = self.data_mask(edge_index=edge_index)
+
+        return type(self)(edge_list[edge_index], edge_weight=self.edge_weight[edge_index], num_nodes=num_nodes,
+                          num_edges=num_edges, num_relation=self.num_relation, offsets=offsets,
+                          meta_dict=meta_dict, **data_dict)
+
+    def subbatch(self, index):
+        """
+        Return a subbatch based on the specified graphs.
+        Equivalent to :meth:`graph_mask(index, compact=True) <graph_mask>`.
+
+        Parameters:
+            index (array_like): graph index
+
+        Returns:
+            PackedGraph
+
+        See also:
+            :meth:`PackedGraph.graph_mask`
+        """
+        return self.graph_mask(index, compact=True)
+
+    def line_graph(self):
+        """
+        Construct a packed line graph of this packed graph.
+        The node features of the line graphs are inherited from the edge features of the original graphs.
+
+        In the line graph, each node corresponds to an edge in the original graph.
+        For a pair of edges (a, b) and (b, c) that share the same intermediate node in the original graph,
+        there is a directed edge (a, b) -> (b, c) in the line graph.
+
+        Returns:
+            PackedGraph
+        """
+        node_in, node_out = self.edge_list.t()[:2]
+        edge_index = torch.arange(self.num_edge, device=self.device)
+        edge_in = edge_index[node_out.argsort()]
+        edge_out = edge_index[node_in.argsort()]
+
+        degree_in = node_in.bincount(minlength=self.num_node)
+        degree_out = node_out.bincount(minlength=self.num_node)
+        size = degree_out * degree_in
+        starts = (size.cumsum(0) - size).repeat_interleave(size)
+        range = torch.arange(size.sum(), device=self.device)
+        # each node u has degree_out[u] * degree_in[u] local edges
+        local_index = range - starts
+        local_inner_size = degree_in.repeat_interleave(size)
+        edge_in_offset = (degree_out.cumsum(0) - degree_out).repeat_interleave(size)
+        edge_out_offset = (degree_in.cumsum(0) - degree_in).repeat_interleave(size)
+        edge_in_index = torch.div(local_index, local_inner_size, rounding_mode="floor") + edge_in_offset
+        edge_out_index = local_index % local_inner_size + edge_out_offset
+
+        edge_in = edge_in[edge_in_index]
+        edge_out = edge_out[edge_out_index]
+        edge_list = torch.stack([edge_in, edge_out], dim=-1)
+        node_feature = getattr(self, "edge_feature", None)
+        num_nodes = self.num_edges
+        num_edges = scatter_add(size, self.node2graph, dim=0, dim_size=self.batch_size)
+        offsets = self._get_offsets(num_nodes, num_edges)
+
+        return PackedGraph(edge_list, num_nodes=num_nodes, num_edges=num_edges, offsets=offsets,
+                           node_feature=node_feature)
+
+    def undirected(self, add_inverse=False):
+        """
+        Flip all the edges to create undirected graphs.
+
+        For knowledge graphs, the flipped edges can either have the original relation or an inverse relation.
+        The inverse relation for relation :math:`r` is defined as :math:`|R| + r`.
+
+        Parameters:
+            add_inverse (bool, optional): whether to use inverse relations for flipped edges
+        """
+        edge_list = self.edge_list.clone()
+        edge_list[:, :2] = edge_list[:, :2].flip(1)
+        num_relation = self.num_relation
+        if num_relation and add_inverse:
+            edge_list[:, 2] += num_relation
+            num_relation = num_relation * 2
+        edge_list = torch.stack([self.edge_list, edge_list], dim=1).flatten(0, 1)
+        offsets = self._offsets.unsqueeze(-1).expand(-1, 2).flatten()
+
+        index = torch.arange(self.num_edge, device=self.device).unsqueeze(-1).expand(-1, 2).flatten()
+        data_dict, meta_dict = self.data_mask(edge_index=index, exclude="edge reference")
+
+        return type(self)(edge_list, edge_weight=self.edge_weight[index], num_nodes=self.num_nodes,
+                          num_edges=self.num_edges * 2, num_relation=num_relation, offsets=offsets,
+                          meta_dict=meta_dict, **data_dict)
+
+    def detach(self):
+        """
+        Detach this packed graph.
+        """
+        return type(self)(self.edge_list.detach(), edge_weight=self.edge_weight.detach(),
+                          num_nodes=self.num_nodes, num_edges=self.num_edges, num_relation=self.num_relation,
+                          offsets=self._offsets, meta_dict=self.meta_dict, **utils.detach(self.data_dict))
+
+    def clone(self):
+        """
+        Clone this packed graph.
+        """
+        return type(self)(self.edge_list.clone(), edge_weight=self.edge_weight.clone(),
+                          num_nodes=self.num_nodes, num_edges=self.num_edges, num_relation=self.num_relation,
+                          offsets=self._offsets, meta_dict=self.meta_dict, **utils.clone(self.data_dict))
+
+    def cuda(self, *args, **kwargs):
+        """
+        Return a copy of this packed graph in CUDA memory.
+
+        This is a non-op if the graph is already on the correct device.
+        """
+        edge_list = self.edge_list.cuda(*args, **kwargs)
+
+        if edge_list is self.edge_list:
+            return self
+        else:
+            return type(self)(edge_list, edge_weight=self.edge_weight,
+                              num_nodes=self.num_nodes, num_edges=self.num_edges, num_relation=self.num_relation,
+                              offsets=self._offsets, meta_dict=self.meta_dict,
+                              **utils.cuda(self.data_dict, *args, **kwargs))
+
+    def cpu(self):
+        """
+        Return a copy of this packed graph in CPU memory.
+
+        This is a non-op if the graph is already in CPU memory.
+        """
+        edge_list = self.edge_list.cpu()
+
+        if edge_list is self.edge_list:
+            return self
+        else:
+            return type(self)(edge_list, edge_weight=self.edge_weight,
+                              num_nodes=self.num_nodes, num_edges=self.num_edges, num_relation=self.num_relation,
+                              offsets=self._offsets, meta_dict=self.meta_dict, **utils.cpu(self.data_dict))
+
+    def __repr__(self):
+        fields = ["batch_size=%d" % self.batch_size,
+                  "num_nodes=%s" % pretty.long_array(self.num_nodes.tolist()),
+                  "num_edges=%s" % pretty.long_array(self.num_edges.tolist())]
+        if self.num_relation is not None:
+            fields.append("num_relation=%d" % self.num_relation)
+        if self.device.type != "cpu":
+            fields.append("device='%s'" % self.device)
+        return "%s(%s)" % (self.__class__.__name__, ", ".join(fields))
+
+    def visualize(self, titles=None, save_file=None, figure_size=(3, 3), layout="spring", num_row=None, num_col=None):
+        """
+        Visualize the packed graphs with matplotlib.
+
+        Parameters:
+            titles (list of str, optional): title for each graph. Default is the ID of each graph.
+            save_file (str, optional): ``png`` or ``pdf`` file to save visualization.
+                If not provided, show the figure in window.
+            figure_size (tuple of int, optional): width and height of the figure
+            layout (str, optional): graph layout
+            num_row (int, optional): number of rows in the figure
+            num_col (int, optional): number of columns in the figure
+
+        See also:
+            `NetworkX graph layout`_
+
+            .. _NetworkX graph layout:
+                https://networkx.github.io/documentation/stable/reference/drawing.html#module-networkx.drawing.layout
+        """
+        if titles is None:
+            graph = self.get_item(0)
+            titles = ["%s %d" % (type(graph).__name__, i) for i in range(self.batch_size)]
+        if num_col is None:
+            if num_row is None:
+                num_col = math.ceil(self.batch_size ** 0.5)
+            else:
+                num_col = math.ceil(self.batch_size / num_row)
+        if num_row is None:
+            num_row = math.ceil(self.batch_size / num_col)
+
+        figure_size = (num_col * figure_size[0], num_row * figure_size[1])
+        fig = plt.figure(figsize=figure_size)
+
+        for i in range(self.batch_size):
+            graph = self.get_item(i)
+            ax = fig.add_subplot(num_row, num_col, i + 1)
+            graph.visualize(title=titles[i], ax=ax, layout=layout)
+        # remove the space of axis labels
+        fig.tight_layout()
+
+        if save_file:
+            fig.savefig(save_file)
+        else:
+            fig.show()
+
+
+Graph.packed_type = PackedGraph
+
+
+def cat(graphs):
+    for i, graph in enumerate(graphs):
+        if not isinstance(graph, PackedGraph):
+            graphs[i] = graph.pack([graph])
+
+    edge_list = torch.cat([graph.edge_list for graph in graphs])
+    pack_num_nodes = torch.stack([graph.num_node for graph in graphs])
+    pack_num_edges = torch.stack([graph.num_edge for graph in graphs])
+    pack_num_cum_edges = pack_num_edges.cumsum(0)
+    graph_index = pack_num_cum_edges < len(edge_list)
+    pack_offsets = scatter_add(pack_num_nodes[graph_index], pack_num_cum_edges[graph_index],
+                               dim_size=len(edge_list))
+    pack_offsets = pack_offsets.cumsum(0)
+
+    edge_list[:, :2] += pack_offsets.unsqueeze(-1)
+    offsets = torch.cat([graph._offsets for graph in graphs]) + pack_offsets
+
+    edge_weight = torch.cat([graph.edge_weight for graph in graphs])
+    num_nodes = torch.cat([graph.num_nodes for graph in graphs])
+    num_edges = torch.cat([graph.num_edges for graph in graphs])
+    num_relation = graphs[0].num_relation
+    assert all(graph.num_relation == num_relation for graph in graphs)
+
+    # only keep attributes that exist in all graphs
+    # TODO: this interface is not safe. re-design the interface
+    keys = set(graphs[0].meta_dict.keys())
+    for graph in graphs:
+        keys = keys.intersection(graph.meta_dict.keys())
+
+    meta_dict = {k: graphs[0].meta_dict[k] for k in keys}
+    data_dict = {}
+    for k in keys:
+        data_dict[k] = torch.cat([graph.data_dict[k] for graph in graphs])
+
+    return type(graphs[0])(edge_list, edge_weight=edge_weight,
+                           num_nodes=num_nodes, num_edges=num_edges, num_relation=num_relation, offsets=offsets,
+                           meta_dict=meta_dict, **data_dict)
\ No newline at end of file