graph index

python/dgl/backend/pytorch.py

@@ -34,6 +34,9 @@ def astype(a, ty):
 def asnumpy(a):
     return a.cpu().numpy()
 
+def from_numpy(np_data):
+    return th.from_numpy(np_data)
+
 def pack(tensors):
     return th.cat(tensors)
 
@@ -43,8 +46,8 @@ def unpack(x, indices_or_sections=1):
 def shape(x):
     return x.shape
 
-def isinteger(x):
-    return x.dtype in [th.int, th.int8, th.int16, th.int32, th.int64]
+def dtype(x):
+    return x.dtype
 
 unique = th.unique
 
@@ -107,11 +110,16 @@ def _typestr(arr_dtype):
     else:
         raise RuntimeError('Unsupported data type:', arr_dtype)
 
-def asdglarray(arr):
-    """The data is copied to the new array."""
-    assert arr.is_contiguous()
-    rst = nd.empty(tuple(arr.shape), _typestr(arr.dtype), get_context(arr))
-    data = ctypes.cast(arr.data_ptr(), ctypes.c_void_p)
-    nbytes = ctypes.c_size_t(arr.numel() * arr.element_size())
-    check_call(_LIB.TVMArrayCopyFromBytes(rst.handle, data, nbytes))
-    return rst
+def astvmarray(arr_data):
+    """Return a TVMArray representation of the underlying data."""
+    data = arr_data
+    assert data.is_contiguous()
+    arr = TVMArray()
+    shape = c_array(tvm_shape_index_t, tuple(data.shape))
+    arr.data = ctypes.cast(data.data_ptr(), ctypes.c_void_p)
+    arr.shape = shape
+    arr.strides = None
+    arr.dtype = TVMType(_typestr(data.dtype))
+    arr.ndim = len(shape)
+    arr.ctx = get_context(data)
+    return arr

python/dgl/graph.py

@@ -3,19 +3,17 @@
 from __future__ import absolute_import
 
 import networkx as nx
-from networkx.classes.digraph import DiGraph
 
 import dgl
 from .base import ALL, is_all, __MSG__, __REPR__
 from . import backend as F
 from .backend import Tensor
-from .cached_graph import CachedGraph, create_cached_graph
+from .graph_index import GraphIndex
 from .frame import FrameRef, merge_frames
-from .nx_adapt import nx_init
 from . import scheduler
 from . import utils
 
-class DGLGraph(DiGraph):
+class DGLGraph(object):
     """Base graph class specialized for neural networks on graphs.
 
     TODO(minjie): document of batching semantics
@@ -38,20 +36,20 @@ class DGLGraph(DiGraph):
                  edge_frame=None,
                  **attr):
         # TODO(minjie): maintaining node/edge list is costly when graph is large.
-        self._edge_list = []
-        nx_init(self,
-                self._add_node_callback,
-                self._add_edge_callback,
-                self._del_node_callback,
-                self._del_edge_callback,
-                graph_data,
-                **attr)
-        # cached graph and storage
-        self._cached_graph = None
+        #nx_init(self,
+        #        self._add_node_callback,
+        #        self._add_edge_callback,
+        #        self._del_node_callback,
+        #        self._del_edge_callback,
+        #        graph_data,
+        #        **attr)
+        # graph
+        self._graph = GraphIndex(graph_data)
+        # frame
         self._node_frame = node_frame if node_frame is not None else FrameRef()
         self._edge_frame = edge_frame if edge_frame is not None else FrameRef()
         # other class members
-        self._msg_graph = None
+        self._msg_graph = GraphIndex()
         self._msg_frame = FrameRef()
         self._message_func = (None, None)
         self._reduce_func = (None, None)
@@ -919,13 +917,6 @@ class DGLGraph(DiGraph):
         pos = graphviz_layout(self, prog='dot')
         nx.draw(self, pos, with_labels=True)
 
-    @property
-    def cached_graph(self):
-        # TODO: dirty flag when mutated
-        if self._cached_graph is None:
-            self._cached_graph = create_cached_graph(self)
-        return self._cached_graph
-
     @property
     def msg_graph(self):
         # TODO: dirty flag when mutated

python/dgl/cgraph.py → python/dgl/graph_index.py

 from __future__ import absolute_import
 
+import numpy as np
+import networkx as nx
+
 from ._ffi.function import _init_api
 from . import backend as F
 from . import utils
@@ -9,12 +12,15 @@ class GraphIndex(object):
 
     Parameters
     ----------
-    graph_data : graph data
+    graph_data : graph data, optional
         Data to initialize graph. Same as networkx's semantics.
     """
     def __init__(self, graph_data=None):
-        # TODO: convert from graph data
         self._handle = _CAPI_DGLGraphCreate()
+        if isinstance(graph_data, nx.DiGraph):
+            self.from_networkx(graph_data)
+        elif graph_data is not None:
+            self.from_networkx(nx.DiGraph(graph_data))
 
     def __del__(self):
         """Free this graph index object."""
@@ -52,10 +58,6 @@ class GraphIndex(object):
         v : utils.Index
             The dst nodes.
         """
-        #u = utils.Index(u)
-        #v = utils.Index(v)
-        #u_array = F.asdglarray(u.totensor())
-        #v_array = F.asdglarray(v.totensor())
         u_array = u.todgltensor()
         v_array = v.todgltensor()
         _CAPI_DGLGraphAddEdges(self._handle, u_array, v_array)
@@ -113,7 +115,7 @@ class GraphIndex(object):
             0-1 array indicating existence
         """
         vid_array = vids.todgltensor()
-        return utils.Index(_CAPI_DGLGraphHasVertices(self._handle, vid_array))
+        return utils.toindex(_CAPI_DGLGraphHasVertices(self._handle, vid_array))
 
     def has_edge(self, u, v):
         """Return true if the edge exists.
@@ -149,7 +151,7 @@ class GraphIndex(object):
         """
         u_array = u.todgltensor()
         v_array = v.todgltensor()
-        return utils.Index(_CAPI_DGLGraphHasEdges(self._handle, u_array, v_array))
+        return utils.toindex(_CAPI_DGLGraphHasEdges(self._handle, u_array, v_array))
 
     def predecessors(self, v, radius=1):
         """Return the predecessors of the node.
@@ -166,7 +168,7 @@ class GraphIndex(object):
         utils.Index
             Array of predecessors
         """
-        return utils.Index(_CAPI_DGLGraphPredecessors(self._handle, v, radius))
+        return utils.toindex(_CAPI_DGLGraphPredecessors(self._handle, v, radius))
 
     def successors(self, v, radius=1):
         """Return the successors of the node.
@@ -183,7 +185,7 @@ class GraphIndex(object):
         utils.Index
             Array of successors
         """
-        return utils.Index(_CAPI_DGLGraphSuccessors(self._handle, v, radius))
+        return utils.toindex(_CAPI_DGLGraphSuccessors(self._handle, v, radius))
 
     def edge_id(self, u, v):
         """Return the id of the edge.
@@ -219,7 +221,7 @@ class GraphIndex(object):
         """
         u_array = u.todgltensor()
         v_array = v.todgltensor()
-        return utils.Index(_CAPI_DGLGraphEdgeIds(self._handle, u_array, v_array))
+        return utils.toindex(_CAPI_DGLGraphEdgeIds(self._handle, u_array, v_array))
 
     def in_edges(self, v):
         """Return the in edges of the node(s).
@@ -243,9 +245,9 @@ class GraphIndex(object):
         else:
             v_array = v.todgltensor()
             edge_array = _CAPI_DGLGraphInEdges_2(self._handle, v_array)
-        src = utils.Index(edge_array(0))
-        dst = utils.Index(edge_array(1))
-        eid = utils.Index(edge_array(2))
+        src = utils.toindex(edge_array(0))
+        dst = utils.toindex(edge_array(1))
+        eid = utils.toindex(edge_array(2))
         return src, dst, eid
 
     def out_edges(self, v):
@@ -270,9 +272,9 @@ class GraphIndex(object):
         else:
             v_array = v.todgltensor()
             edge_array = _CAPI_DGLGraphOutEdges_2(self._handle, v_array)
-        src = utils.Index(edge_array(0))
-        dst = utils.Index(edge_array(1))
-        eid = utils.Index(edge_array(2))
+        src = utils.toindex(edge_array(0))
+        dst = utils.toindex(edge_array(1))
+        eid = utils.toindex(edge_array(2))
         return src, dst, eid
 
     def edges(self, sorted=False):
@@ -293,9 +295,9 @@ class GraphIndex(object):
             The edge ids.
         """
         edge_array = _CAPI_DGLGraphEdges(self._handle, sorted)
-        src = edge_array(0)
-        dst = edge_array(1)
-        eid = edge_array(2)
+        src = utils.toindex(edge_array(0))
+        dst = utils.toindex(edge_array(1))
+        eid = utils.toindex(edge_array(2))
         return src, dst, eid
 
     def in_degree(self, v):
@@ -327,7 +329,7 @@ class GraphIndex(object):
             The in degree array.
         """
         v_array = v.todgltensor()
-        return utils.Index(_CAPI_DGLGraphInDegrees(self._handle, v_array))
+        return utils.toindex(_CAPI_DGLGraphInDegrees(self._handle, v_array))
 
     def out_degree(self, v):
         """Return the out degree of the node.
@@ -358,17 +360,56 @@ class GraphIndex(object):
             The out degree array.
         """
         v_array = v.todgltensor()
-        return utils.Index(_CAPI_DGLGraphOutDegrees(self._handle, v_array))
+        return utils.toindex(_CAPI_DGLGraphOutDegrees(self._handle, v_array))
 
-    def asnetworkx(self):
+    def to_networkx(self):
         """Convert to networkx graph.
 
+        The edge id will be saved as the 'id' edge attribute.
+
         Returns
         -------
         networkx.DiGraph
             The nx graph
         """
-        # TODO
-        return None
+        src, dst, eid = self.edges()
+        ret = nx.DiGraph()
+        for u, v, id in zip(src, dst, eid):
+            ret.add_edge(u, v, id=id)
+        return ret
+
+    def from_networkx(self, nx_graph):
+        """Convert from networkx graph.
 
-_init_api("dgl.graph")
+        If 'id' edge attribute exists, the edge will be added follows
+        the edge id order. Otherwise, order is undefined.
+        
+        Parameters
+        ----------
+        nx_graph : networkx.DiGraph
+            The nx graph
+        """
+        self.clear()
+        num_nodes = nx_graph.number_of_nodes()
+        self.add_nodes(num_nodes)
+        has_edge_id = 'id' in next(iter(nx_graph.edges))
+        if has_edge_id:
+            num_edges = nx_graph.number_of_edges()
+            src = np.zeros((num_edges,), dtype=np.int64)
+            dst = np.zeros((num_edges,), dtype=np.int64)
+            for e, attr in nx_graph.edges.items:
+                u, v = e
+                eid = attr['id']
+                src[eid] = u
+                dst[eid] = v
+        else:
+            src = []
+            dst = []
+            for u, v in nx_graph.edges:
+                src.append(u)
+                dst.append(v)
+        src = utils.toindex(src)
+        dst = utils.toindex(dst)
+        self.add_edges(src, dst)
+
+_init_api("dgl.graph_index")

python/dgl/ndarray.py

 """DGL Runtime NDArray API.
 
-dgl.ndarray provides a minimum runtime array API to unify
-different array libraries used as backend.
+dgl.ndarray provides a minimum runtime array structure to be
+used with C++ library.
 """
 # pylint: disable=invalid-name,unused-import
 from __future__ import absolute_import as _abs
+
+import ctypes
+import operator
 import numpy as _np
 
 from ._ffi.ndarray import TVMContext, TVMType, NDArrayBase
-from ._ffi.ndarray import context, empty, from_dlpack
+from ._ffi.ndarray import context, empty, from_dlpack, numpyasarray
 from ._ffi.ndarray import _set_class_ndarray
+from . import backend as F
 
 class NDArray(NDArrayBase):
     """Lightweight NDArray class for DGL framework."""
-    pass
+    def __len__(self):
+        return reduce(operator.mul, self.shape, 1)
 
 def cpu(dev_id=0):
     """Construct a CPU device
@@ -66,4 +71,38 @@ def array(arr, ctx=cpu(0)):
         arr = _np.array(arr)
     return empty(arr.shape, arr.dtype, ctx).copyfrom(arr)
 
+def from_numpy(np_data):
+    """Create an array that shares the given numpy data.
+
+    Parameters
+    ----------
+    np_data : numpy.ndarray
+        The numpy data
+
+    Returns
+    -------
+    NDArray
+        The array
+    """
+    arr, _ = numpyasarray(np_data)
+    handle = ctypes.pointer(arr)
+    return NDArray(handle, is_view=True)
+
+def from_user_tensor(data):
+    """Create an array that shares the given user tensor data.
+
+    Parameters
+    ----------
+    data : F.Tensor
+        The user tensor data.
+
+    Returns
+    -------
+    NDArray
+        The array
+    """
+    arr = F.astvmarray(data)
+    handle = ctypes.pointer(arr)
+    return NDArray(handle, is_view=True)
+
 _set_class_ndarray(NDArray)

python/dgl/utils.py

@@ -7,48 +7,68 @@ import numpy as np
 
 from . import backend as F
 from .backend import Tensor, SparseTensor
-
-def is_id_tensor(u):
-    """Return whether the input is a supported id tensor."""
-    return isinstance(u, Tensor) and F.isinteger(u) and len(F.shape(u)) == 1
-
-def is_id_container(u):
-    """Return whether the input is a supported id container."""
-    return (getattr(u, '__iter__', None) is not None
-            and getattr(u, '__len__', None) is not None)
+from . import ndarray as nd
 
 class Index(object):
     """Index class that can be easily converted to list/tensor."""
     def __init__(self, data):
-        self._list_data = None
-        self._tensor_data = None
-        self._ctx_data = dict()
+        self._list_data = None  # a numpy type data
+        self._user_tensor_data = dict()  # dictionary of user tensors
+        self._dgl_tensor_data = None  # a dgl ndarray
         self._dispatch(data)
 
     def _dispatch(self, data):
-        if is_id_tensor(data):
-            self._tensor_data = data
-        elif is_id_container(data):
-            self._list_data = data
+        """Store data based on its type."""
+        if isinstance(data, Tensor):
+            if not (F.dtype(data) == F.int64 and len(F.shape(data)) == 1):
+                raise ValueError('Index data must be 1D int64 vector, but got: %s' % str(data))
+            self._user_tensor_data[F.get_context(data)] = data
+        elif isinstance(data, nd.NDArray): 
+            if not (data.dtype == 'int64' and len(data.shape) == 1):
+                raise ValueError('Index data must be 1D int64 vector, but got: %s' % str(data))
+            self._dgl_tensor_data = data
         else:
             try:
-                self._list_data = [int(data)]
+                self._list_data = np.array([int(data)]).astype(np.int64)
             except:
-                raise TypeError('Error index data: %s' % str(x))
+                try:
+                    self._list_data = np.array(data).astype(np.int64)
+                except:
+                    raise ValueError('Error index data: %s' % str(data))
 
     def tolist(self):
+        """Convert to a python-list compatible object."""
         if self._list_data is None:
-            self._list_data = list(F.asnumpy(self._tensor_data))
+            if self._dgl_tensor_data is not None:
+                self._list_data = self._dgl_tensor_data.asnumpy()
+            else:
+                assert len(self._user_tensor_data) > 0
+                data = next(iter(self._user_tensor_data.values()))
+                self._list_data = F.asnumpy(data)
         return self._list_data
 
-    def totensor(self, ctx=None):
-        if self._tensor_data is None:
-            self._tensor_data = F.tensor(self._list_data, dtype=F.int64)
+    def tousertensor(self, ctx=None):
+        """Convert to user tensor (defined in `backend`)."""
+        if len(self._user_tensor_data) == 0:
+            self._user_tensor_data[nd.cpu()] = F.from_numpy(self.tolist())
         if ctx is None:
-            return self._tensor_data
-        if ctx not in self._ctx_data:
-            self._ctx_data[ctx] = F.to_context(self._tensor_data, ctx)
-        return self._ctx_data[ctx]
+            ctx = nd.cpu()
+        if ctx not in self._user_tensor_data:
+            data = next(iter(self._user_tensor_data.values()))
+            self._user_tensor_data[ctx] = F.to_context(data, ctx)
+        return self._user_tensor_data[ctx]
+
+    def todgltensor(self):
+        """Convert to dgl.NDArray."""
+        if self._dgl_tensor_data is None:
+            if self._list_data is not None:
+                # create a view ndarray from numpy
+                self._dgl_tensor_data = nd.from_numpy(self._list_data)
+            else:
+                # create a view ndarray from user tensor
+                self._dgl_tensor_data = nd.from_user_tensor(
+                        self.tousertensor(ctx=nd.cpu()))
+        return self._dgl_tensor_data
 
     def __iter__(self):
         return iter(self.tolist())
@@ -56,8 +76,11 @@ class Index(object):
     def __len__(self):
         if self._list_data is not None:
             return len(self._list_data)
+        elif len(self._user_tensor_data) > 0:
+            data = next(iter(self._user_tensor_data.values()))
+            return len(data)
         else:
-            return len(self._tensor_data)
+            return len(self._dgl_tensor_data)
 
     def __getitem__(self, i):
         return self.tolist()[i]
@@ -125,33 +148,6 @@ def edge_broadcasting(u, v):
         assert len(u) == len(v)
     return u, v
 
-'''
-def convert_to_id_container(x):
-    if is_id_container(x):
-        return x
-    elif is_id_tensor(x):
-        return F.asnumpy(x)
-    else:
-        try:
-            return [int(x)]
-        except:
-            raise TypeError('Error node: %s' % str(x))
-    return None
-
-def convert_to_id_tensor(x, ctx=None):
-    if is_id_container(x):
-        ret = F.tensor(x, dtype=F.int64)
-    elif is_id_tensor(x):
-        ret = x
-    else:
-        try:
-            ret = F.tensor([int(x)], dtype=F.int64)
-        except:
-            raise TypeError('Error node: %s' % str(x))
-    ret = F.to_context(ret, ctx)
-    return ret
-'''
-
 class LazyDict(Mapping):
     """A readonly dictionary that does not materialize the storage."""
     def __init__(self, fn, keys):

src/graph/graph_apis.cc

@@ -9,7 +9,7 @@ using tvm::runtime::PackedFunc;
 
 namespace dgl {
 namespace {
-
+/*!\brief Convert EdgeArray structure to PackedFunc */
 PackedFunc ConvertEdgeArrayToPackedFunc(const Graph::EdgeArray& ea) {
   auto body = [ea] (TVMArgs args, TVMRetValue* rv) {
       int which = args[0];
@@ -26,25 +26,34 @@ PackedFunc ConvertEdgeArrayToPackedFunc(const Graph::EdgeArray& ea) {
   return PackedFunc(body);
 }
 
+DLManagedTensor* CreateTmpDLManagedTensor(const TVMArgValue& arg) {
+  const DLTensor* dl_tensor = arg;
+  DLManagedTensor* ret = new DLManagedTensor();
+  ret->deleter = [] (DLManagedTensor* self) { delete self; };
+  ret->manager_ctx = nullptr;
+  ret->dl_tensor = *dl_tensor;
+  return ret;
+}
+
 }  // namespace
 
 // Graph handler type
 typedef void* GraphHandle;
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphCreate")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreate")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = new Graph();
     *rv = ghandle;
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphFree")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFree")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
     delete gptr;
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphAddVertices")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddVertices")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
@@ -52,7 +61,7 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphAddVertices")
     gptr->AddVertices(num_vertices);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphAddEdge")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
@@ -61,37 +70,37 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphAddEdge")
     gptr->AddEdge(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphAddEdges")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
-    const IdArray src = args[1];
-    const IdArray dst = args[2];
+    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
+    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
     gptr->AddEdges(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphClear")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphClear")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     Graph* gptr = static_cast<Graph*>(ghandle);
     gptr->Clear();
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphNumVertices")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumVertices")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
     *rv = static_cast<int64_t>(gptr->NumVertices());
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphNumEdges")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumEdges")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
     *rv = static_cast<int64_t>(gptr->NumEdges());
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasVertex")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertex")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -99,15 +108,15 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasVertex")
     *rv = gptr->HasVertex(vid);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasVertices")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertices")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
-    const IdArray vids = args[1];
+    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
     *rv = gptr->HasVertices(vids);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasEdge")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdge")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -116,16 +125,16 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasEdge")
     *rv = gptr->HasEdge(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphHasEdges")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdges")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
-    const IdArray src = args[1];
-    const IdArray dst = args[2];
+    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
+    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
     *rv = gptr->HasEdges(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphPredecessors")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphPredecessors")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -134,7 +143,7 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphPredecessors")
     *rv = gptr->Predecessors(vid, radius);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphSuccessors")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -143,7 +152,7 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphSuccessors")
     *rv = gptr->Successors(vid, radius);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphEdgeId")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -152,16 +161,16 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphEdgeId")
     *rv = static_cast<int64_t>(gptr->EdgeId(src, dst));
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphEdgeIds")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
-    const IdArray src = args[1];
-    const IdArray dst = args[2];
+    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
+    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
     *rv = gptr->EdgeIds(src, dst);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInEdges_1")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_1")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -169,15 +178,15 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInEdges_1")
     *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vid));
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInEdges_2")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_2")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
-    const IdArray vids = args[1];
+    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
     *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vids));
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutEdges_1")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_1")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -185,15 +194,15 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutEdges_1")
     *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vid));
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutEdges_2")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_2")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
-    const IdArray vids = args[1];
+    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
     *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vids));
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphEdges")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdges")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -201,7 +210,7 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphEdges")
     *rv = ConvertEdgeArrayToPackedFunc(gptr->Edges(sorted));
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInDegree")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegree")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -209,15 +218,15 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInDegree")
     *rv = static_cast<int64_t>(gptr->InDegree(vid));
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphInDegrees")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegrees")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
-    const IdArray vids = args[1];
+    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
     *rv = gptr->InDegrees(vids);
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutDegree")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegree")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
@@ -225,11 +234,11 @@ TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutDegree")
     *rv = static_cast<int64_t>(gptr->OutDegree(vid));
   });
 
-TVM_REGISTER_GLOBAL("graph._CAPI_DGLGraphOutDegrees")
+TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegrees")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
-    const IdArray vids = args[1];
+    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
     *rv = gptr->OutDegrees(vids);
   });