chore: import upstream snapshot with attribution

python/dgl/__init__.py

@@ -0,0 +1,64 @@
+"""
+The ``dgl`` package contains data structure for storing structural and feature data
+(i.e., the :class:`DGLGraph` class) and also utilities for generating, manipulating
+and transforming graphs.
+"""
+
+
+# Windows compatibility
+# This initializes Winsock and performs cleanup at termination as required
+import socket
+
+# Backend and logging should be imported before other modules.
+from .logging import enable_verbose_logging  # usort: skip
+from .backend import backend_name, load_backend  # usort: skip
+
+from . import (
+    container,
+    cuda,
+    dataloading,
+    function,
+    ops,
+    random,
+    sampling,
+    storages,
+)
+from ._ffi.base import __version__, DGLError
+from ._ffi.function import (
+    extract_ext_funcs,
+    get_global_func,
+    list_global_func_names,
+    register_func,
+)
+
+from ._ffi.runtime_ctypes import TypeCode
+
+from .base import ALL, EID, ETYPE, NID, NTYPE
+from .readout import *
+from .batch import *
+from .convert import *
+from .generators import *
+from .dataloading import (
+    set_dst_lazy_features,
+    set_edge_lazy_features,
+    set_node_lazy_features,
+    set_src_lazy_features,
+)
+from .heterograph import (  # pylint: disable=reimported
+    DGLGraph,
+    DGLGraph as DGLHeteroGraph,
+)
+from .merge import *
+from .subgraph import *
+from .traversal import *
+from .transforms import *
+from .propagate import *
+from .random import *
+from . import optim
+from .data.utils import load_graphs, save_graphs
+from .frame import LazyFeature
+from .global_config import is_libxsmm_enabled, use_libxsmm
+from .utils import apply_each
+from .mpops import *
+from .homophily import *
+from .label_informativeness import *

python/dgl/_api_internal.py

@@ -0,0 +1 @@
+"""Namespace for internal apis."""

python/dgl/_ffi/README.md

@@ -0,0 +1,3 @@
+# C API and runtime
+
+Borrowed and adapted from TVM project. (commit: 2ce5277)

python/dgl/_ffi/__init__.py

@@ -0,0 +1,10 @@
+"""C interfacing code.
+
+This namespace contains everything that interacts with C code.
+Most C related object are ctypes compatible, which means
+they contains a handle field that is ctypes.c_void_p and can
+be used via ctypes function calls.
+
+Some performance critical functions are implemented by cython
+and have a ctypes fallback implementation.
+"""

python/dgl/_ffi/_ctypes/__init__.py

@@ -0,0 +1 @@
+"""ctypes specific implementation of FFI"""

python/dgl/_ffi/_ctypes/function.py

@@ -0,0 +1,298 @@
+# coding: utf-8
+# pylint: disable=invalid-name, protected-access, too-many-branches, global-statement
+"""Function configuration API."""
+from __future__ import absolute_import
+
+import ctypes
+import traceback
+from numbers import Integral, Number
+
+from ..base import _LIB, c_str, check_call, string_types
+from ..object_generic import convert_to_object, ObjectGeneric
+from ..runtime_ctypes import DGLByteArray, DGLContext, DGLDataType
+from . import ndarray as _nd, object as _object
+from .ndarray import _make_array, NDArrayBase
+from .object import ObjectBase
+from .types import (
+    _wrap_arg_func,
+    C_TO_PY_ARG_SWITCH,
+    DGLCFuncFinalizer,
+    DGLPackedCFunc,
+    DGLValue,
+    RETURN_SWITCH,
+    TypeCode,
+)
+
+FunctionHandle = ctypes.c_void_p
+ModuleHandle = ctypes.c_void_p
+DGLRetValueHandle = ctypes.c_void_p
+
+
+def _ctypes_free_resource(rhandle):
+    """callback to free resources when it it not needed."""
+    pyobj = ctypes.cast(rhandle, ctypes.py_object)
+    ctypes.pythonapi.Py_DecRef(pyobj)
+
+
+# Global callback that is always alive
+DGL_FREE_PYOBJ = DGLCFuncFinalizer(_ctypes_free_resource)
+ctypes.pythonapi.Py_IncRef(ctypes.py_object(DGL_FREE_PYOBJ))
+
+
+def convert_to_dgl_func(pyfunc):
+    """Convert a python function to DGL function
+
+    Parameters
+    ----------
+    pyfunc : python function
+        The python function to be converted.
+
+    Returns
+    -------
+    dglfunc: dgl.nd.Function
+        The converted dgl function.
+    """
+    local_pyfunc = pyfunc
+
+    def cfun(args, type_codes, num_args, ret, _):
+        """ctypes function"""
+        num_args = (
+            num_args.value if isinstance(num_args, ctypes.c_int) else num_args
+        )
+        pyargs = (
+            C_TO_PY_ARG_SWITCH[type_codes[i]](args[i]) for i in range(num_args)
+        )
+        # pylint: disable=broad-except
+        try:
+            rv = local_pyfunc(*pyargs)
+        except Exception:
+            msg = traceback.format_exc()
+            _LIB.DGLAPISetLastError(c_str(msg))
+            return -1
+
+        if rv is not None:
+            if isinstance(rv, tuple):
+                raise ValueError(
+                    "PackedFunction can only support one return value"
+                )
+            temp_args = []
+            values, tcodes, _ = _make_dgl_args((rv,), temp_args)
+            if not isinstance(ret, DGLRetValueHandle):
+                ret = DGLRetValueHandle(ret)
+            check_call(
+                _LIB.DGLCFuncSetReturn(ret, values, tcodes, ctypes.c_int(1))
+            )
+            _ = temp_args
+            _ = rv
+        return 0
+
+    handle = FunctionHandle()
+    f = DGLPackedCFunc(cfun)
+    # NOTE: We will need to use python-api to increase ref count of the f
+    # DGL_FREE_PYOBJ will be called after it is no longer needed.
+    pyobj = ctypes.py_object(f)
+    ctypes.pythonapi.Py_IncRef(pyobj)
+    check_call(
+        _LIB.DGLFuncCreateFromCFunc(
+            f, pyobj, DGL_FREE_PYOBJ, ctypes.byref(handle)
+        )
+    )
+    return _CLASS_FUNCTION(handle, False)
+
+
+def _make_dgl_args(args, temp_args):
+    """Pack arguments into c args dgl call accept.
+
+    temp_args is used to temporarily save the arguments so they will not be
+    freed during C API function call.
+    """
+    num_args = len(args)
+    values = (DGLValue * num_args)()
+    type_codes = (ctypes.c_int * num_args)()
+    for i, arg in enumerate(args):
+        if arg is None:
+            values[i].v_handle = None
+            type_codes[i] = TypeCode.NULL
+        elif isinstance(arg, ObjectBase):
+            values[i].v_handle = arg.handle
+            type_codes[i] = TypeCode.OBJECT_HANDLE
+        elif isinstance(arg, (list, tuple, dict, ObjectGeneric)):
+            arg = convert_to_object(arg)
+            values[i].v_handle = arg.handle
+            type_codes[i] = TypeCode.OBJECT_HANDLE
+            temp_args.append(arg)
+        elif isinstance(arg, NDArrayBase):
+            values[i].v_handle = ctypes.cast(arg.handle, ctypes.c_void_p)
+            type_codes[i] = (
+                TypeCode.NDARRAY_CONTAINER
+                if not arg.is_view
+                else TypeCode.ARRAY_HANDLE
+            )
+        elif isinstance(arg, _nd._DGL_COMPATS):
+            values[i].v_handle = ctypes.c_void_p(arg._dgl_handle)
+            type_codes[i] = arg.__class__._dgl_tcode
+        elif isinstance(arg, Integral):
+            values[i].v_int64 = arg
+            type_codes[i] = TypeCode.INT
+        elif isinstance(arg, Number):
+            values[i].v_float64 = arg
+            type_codes[i] = TypeCode.FLOAT
+        elif isinstance(arg, DGLDataType):
+            values[i].v_str = c_str(str(arg))
+            type_codes[i] = TypeCode.STR
+        elif isinstance(arg, DGLContext):
+            values[i].v_ctx = arg
+            type_codes[i] = TypeCode.DGL_CONTEXT
+        elif isinstance(arg, bytearray):
+            arr = DGLByteArray()
+            arr.data = ctypes.cast(
+                (ctypes.c_byte * len(arg)).from_buffer(arg),
+                ctypes.POINTER(ctypes.c_byte),
+            )
+            arr.size = len(arg)
+            values[i].v_handle = ctypes.c_void_p(ctypes.addressof(arr))
+            temp_args.append(arr)
+            type_codes[i] = TypeCode.BYTES
+        elif isinstance(arg, string_types):
+            values[i].v_str = c_str(arg)
+            type_codes[i] = TypeCode.STR
+        # NOTE(minjie): module is not used in DGL
+        # elif isinstance(arg, _CLASS_MODULE):
+        #    values[i].v_handle = arg.handle
+        #    type_codes[i] = TypeCode.MODULE_HANDLE
+        elif isinstance(arg, FunctionBase):
+            values[i].v_handle = arg.handle
+            type_codes[i] = TypeCode.FUNC_HANDLE
+        elif isinstance(arg, ctypes.c_void_p):
+            values[i].v_handle = arg
+            type_codes[i] = TypeCode.HANDLE
+        elif callable(arg):
+            arg = convert_to_dgl_func(arg)
+            values[i].v_handle = arg.handle
+            type_codes[i] = TypeCode.FUNC_HANDLE
+            temp_args.append(arg)
+        else:
+            raise TypeError("Don't know how to handle type %s" % type(arg))
+    return values, type_codes, num_args
+
+
+class FunctionBase(object):
+    """Function base."""
+
+    __slots__ = ["handle", "is_global"]
+    # pylint: disable=no-member
+    def __init__(self, handle, is_global):
+        """Initialize the function with handle
+
+        Parameters
+        ----------
+        handle : FunctionHandle
+            the handle to the underlying function.
+
+        is_global : bool
+            Whether this is a global function in python
+        """
+        self.handle = handle
+        self.is_global = is_global
+
+    def __del__(self):
+        if not self.is_global and _LIB is not None:
+            check_call(_LIB.DGLFuncFree(self.handle))
+
+    def __call__(self, *args):
+        """Call the function with positional arguments
+
+        args : list
+           The positional arguments to the function call.
+        """
+        temp_args = []
+        values, tcodes, num_args = _make_dgl_args(args, temp_args)
+        ret_val = DGLValue()
+        ret_tcode = ctypes.c_int()
+        check_call(
+            _LIB.DGLFuncCall(
+                self.handle,
+                values,
+                tcodes,
+                ctypes.c_int(num_args),
+                ctypes.byref(ret_val),
+                ctypes.byref(ret_tcode),
+            )
+        )
+        _ = temp_args
+        _ = args
+        return RETURN_SWITCH[ret_tcode.value](ret_val)
+
+
+def __init_handle_by_constructor__(fconstructor, args):
+    """Initialize handle by constructor"""
+    temp_args = []
+    values, tcodes, num_args = _make_dgl_args(args, temp_args)
+    ret_val = DGLValue()
+    ret_tcode = ctypes.c_int()
+    check_call(
+        _LIB.DGLFuncCall(
+            fconstructor.handle,
+            values,
+            tcodes,
+            ctypes.c_int(num_args),
+            ctypes.byref(ret_val),
+            ctypes.byref(ret_tcode),
+        )
+    )
+    _ = temp_args
+    _ = args
+    assert ret_tcode.value == TypeCode.OBJECT_HANDLE
+    handle = ret_val.v_handle
+    return handle
+
+
+def _return_module(x):
+    """Return function"""
+    handle = x.v_handle
+    if not isinstance(handle, ModuleHandle):
+        handle = ModuleHandle(handle)
+    return _CLASS_MODULE(handle)
+
+
+def _handle_return_func(x):
+    """Return function"""
+    handle = x.v_handle
+    if not isinstance(handle, FunctionHandle):
+        handle = FunctionHandle(handle)
+    return _CLASS_FUNCTION(handle, False)
+
+
+# setup return handle for function type
+_object.__init_by_constructor__ = __init_handle_by_constructor__
+RETURN_SWITCH[TypeCode.FUNC_HANDLE] = _handle_return_func
+RETURN_SWITCH[TypeCode.MODULE_HANDLE] = _return_module
+RETURN_SWITCH[TypeCode.NDARRAY_CONTAINER] = lambda x: _make_array(
+    x.v_handle, False
+)
+C_TO_PY_ARG_SWITCH[TypeCode.FUNC_HANDLE] = _wrap_arg_func(
+    _handle_return_func, TypeCode.FUNC_HANDLE
+)
+C_TO_PY_ARG_SWITCH[TypeCode.MODULE_HANDLE] = _wrap_arg_func(
+    _return_module, TypeCode.MODULE_HANDLE
+)
+C_TO_PY_ARG_SWITCH[TypeCode.ARRAY_HANDLE] = lambda x: _make_array(
+    x.v_handle, True
+)
+C_TO_PY_ARG_SWITCH[TypeCode.NDARRAY_CONTAINER] = lambda x: _make_array(
+    x.v_handle, False
+)
+
+_CLASS_MODULE = None
+_CLASS_FUNCTION = None
+
+
+def _set_class_module(module_class):
+    """Initialize the module."""
+    global _CLASS_MODULE
+    _CLASS_MODULE = module_class
+
+
+def _set_class_function(func_class):
+    global _CLASS_FUNCTION
+    _CLASS_FUNCTION = func_class

python/dgl/_ffi/_ctypes/ndarray.py

@@ -0,0 +1,137 @@
+# pylint: disable=invalid-name
+"""Runtime NDArray api"""
+from __future__ import absolute_import
+
+import ctypes
+
+from ..base import _LIB, c_str, check_call
+from ..runtime_ctypes import DGLArrayHandle
+from .types import (
+    _return_handle,
+    _wrap_arg_func,
+    C_TO_PY_ARG_SWITCH,
+    RETURN_SWITCH,
+)
+
+DGLPyCapsuleDestructor = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
+_c_str_dltensor = c_str("dltensor")
+_c_str_used_dltensor = c_str("used_dltensor")
+
+
+# used for PyCapsule manipulation
+if hasattr(ctypes, "pythonapi"):
+    ctypes.pythonapi.PyCapsule_GetName.restype = ctypes.c_char_p
+    ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p
+    ctypes.pythonapi.PyCapsule_New.restype = ctypes.py_object
+
+
+def _from_dlpack(dltensor):
+    dltensor = ctypes.py_object(dltensor)
+    if ctypes.pythonapi.PyCapsule_IsValid(dltensor, _c_str_dltensor):
+        ptr = ctypes.pythonapi.PyCapsule_GetPointer(dltensor, _c_str_dltensor)
+        # XXX(minjie): The below cast should be unnecessary given the code to
+        #   set restype of PyCapsule calls. But weirdly, this does not
+        #   work out always.
+        ptr = ctypes.cast(ptr, ctypes.c_void_p)
+        handle = DGLArrayHandle()
+        check_call(_LIB.DGLArrayFromDLPack(ptr, ctypes.byref(handle)))
+        ctypes.pythonapi.PyCapsule_SetName(dltensor, _c_str_used_dltensor)
+        ctypes.pythonapi.PyCapsule_SetDestructor(
+            dltensor, DGLPyCapsuleDestructor(0)
+        )
+        return _make_array(handle, False)
+    raise ValueError(
+        "Expect a dltensor field, PyCapsule can only be consumed once"
+    )
+
+
+def _dlpack_deleter(pycapsule):
+    pycapsule = ctypes.cast(pycapsule, ctypes.py_object)
+    if ctypes.pythonapi.PyCapsule_IsValid(pycapsule, _c_str_dltensor):
+        ptr = ctypes.pythonapi.PyCapsule_GetPointer(pycapsule, _c_str_dltensor)
+        # XXX(minjie): The below cast should be unnecessary given the code to
+        #   set restype of PyCapsule calls. But weirdly, this does not
+        #   work out always.
+        ptr = ctypes.cast(ptr, ctypes.c_void_p)
+        _LIB.DGLDLManagedTensorCallDeleter(ptr)
+        ctypes.pythonapi.PyCapsule_SetDestructor(
+            pycapsule, DGLPyCapsuleDestructor(0)
+        )
+
+
+_c_dlpack_deleter = DGLPyCapsuleDestructor(_dlpack_deleter)
+
+
+class NDArrayBase(object):
+    """A simple Device/CPU Array object in runtime."""
+
+    __slots__ = ["handle", "is_view"]
+    # pylint: disable=no-member
+    def __init__(self, handle, is_view=False):
+        """Initialize the function with handle
+
+        Parameters
+        ----------
+        handle : DGLArrayHandle
+            the handle to the underlying C++ DGLArray
+        """
+        self.handle = handle
+        self.is_view = is_view
+
+    def __del__(self):
+        if not self.is_view and _LIB:
+            check_call(_LIB.DGLArrayFree(self.handle))
+
+    @property
+    def _dgl_handle(self):
+        return ctypes.cast(self.handle, ctypes.c_void_p).value
+
+    def to_dlpack(self, alignment=0):
+        """Produce an array from a DLPack Tensor without copying memory
+
+        Args
+        -------
+        alignment: int, default to be 0
+        Indicates the alignment requirement when converting to dlpack. Will copy to a
+        new tensor if the alignment requirement is not satisfied.
+        0 means no alignment requirement.
+
+
+        Returns
+        -------
+        dlpack : DLPack tensor view of the array data
+        """
+        ptr = ctypes.c_void_p()
+        check_call(
+            _LIB.DGLArrayToDLPack(self.handle, ctypes.byref(ptr), alignment)
+        )
+        return ctypes.pythonapi.PyCapsule_New(
+            ptr, _c_str_dltensor, _c_dlpack_deleter
+        )
+
+
+def _make_array(handle, is_view):
+    handle = ctypes.cast(handle, DGLArrayHandle)
+    return _CLASS_NDARRAY(handle, is_view)
+
+
+_DGL_COMPATS = ()
+
+
+def _reg_extension(cls, fcreate):
+    global _DGL_COMPATS
+    _DGL_COMPATS += (cls,)
+    if fcreate:
+        fret = lambda x: fcreate(_return_handle(x))
+        RETURN_SWITCH[cls._dgl_tcode] = fret
+        C_TO_PY_ARG_SWITCH[cls._dgl_tcode] = _wrap_arg_func(
+            fret, cls._dgl_tcode
+        )
+
+
+_CLASS_NDARRAY = None
+
+
+def _set_class_ndarray(cls):
+    global _CLASS_NDARRAY
+    _CLASS_NDARRAY = cls

python/dgl/_ffi/_ctypes/object.py

@@ -0,0 +1,109 @@
+"""ctypes object API."""
+from __future__ import absolute_import
+
+import ctypes
+
+from ..base import _LIB, c_str, check_call
+from ..object_generic import _set_class_object_base
+from .types import (
+    _wrap_arg_func,
+    C_TO_PY_ARG_SWITCH,
+    DGLValue,
+    RETURN_SWITCH,
+    TypeCode,
+)
+
+ObjectHandle = ctypes.c_void_p
+__init_by_constructor__ = None
+
+"""Maps object type to its constructor"""
+OBJECT_TYPE = {}
+
+
+def _register_object(index, cls):
+    """register object class in python"""
+    OBJECT_TYPE[index] = cls
+
+
+def _return_object(x):
+    """Construct a object object from the given DGLValue object"""
+    handle = x.v_handle
+    if not isinstance(handle, ObjectHandle):
+        handle = ObjectHandle(handle)
+    tindex = ctypes.c_int()
+    check_call(_LIB.DGLObjectGetTypeIndex(handle, ctypes.byref(tindex)))
+    cls = OBJECT_TYPE.get(tindex.value, ObjectBase)
+    # Avoid calling __init__ of cls, instead directly call __new__
+    # This allows child class to implement their own __init__
+    obj = cls.__new__(cls)
+    obj.handle = handle
+    return obj
+
+
+RETURN_SWITCH[TypeCode.OBJECT_HANDLE] = _return_object
+C_TO_PY_ARG_SWITCH[TypeCode.OBJECT_HANDLE] = _wrap_arg_func(
+    _return_object, TypeCode.OBJECT_HANDLE
+)
+
+
+class ObjectBase(object):
+    """Object base class"""
+
+    __slots__ = ["handle"]
+
+    # pylint: disable=no-member
+    def __del__(self):
+        if _LIB is not None and hasattr(self, "handle"):
+            check_call(_LIB.DGLObjectFree(self.handle))
+
+    def __getattr__(self, name):
+        if name == "handle":
+            raise AttributeError(
+                "'handle' is a reserved attribute name that should not be used"
+            )
+        ret_val = DGLValue()
+        ret_type_code = ctypes.c_int()
+        ret_success = ctypes.c_int()
+        check_call(
+            _LIB.DGLObjectGetAttr(
+                self.handle,
+                c_str(name),
+                ctypes.byref(ret_val),
+                ctypes.byref(ret_type_code),
+                ctypes.byref(ret_success),
+            )
+        )
+        if not ret_success.value:
+            raise AttributeError(
+                "'%s' object has no attribute '%s'" % (str(type(self)), name)
+            )
+        return RETURN_SWITCH[ret_type_code.value](ret_val)
+
+    def __init_handle_by_constructor__(self, fconstructor, *args):
+        """Initialize the handle by calling constructor function.
+
+        Parameters
+        ----------
+        fconstructor : Function
+            Constructor function.
+
+        args: list of objects
+            The arguments to the constructor
+
+        Note
+        ----
+        We have a special calling convention to call constructor functions.
+        So the return handle is directly set into the Object object
+        instead of creating a new Object.
+        """
+        # assign handle first to avoid error raising
+        self.handle = None
+        handle = __init_by_constructor__(
+            fconstructor, args
+        )  # pylint: disable=not-callable
+        if not isinstance(handle, ObjectHandle):
+            handle = ObjectHandle(handle)
+        self.handle = handle
+
+
+_set_class_object_base(ObjectBase)

python/dgl/_ffi/_ctypes/types.py

@@ -0,0 +1,91 @@
+"""The C Types used in API."""
+# pylint: disable=invalid-name
+from __future__ import absolute_import as _abs
+
+import ctypes
+
+from ..base import _LIB, check_call, py_str
+from ..runtime_ctypes import DGLByteArray, DGLContext, DGLDataType, TypeCode
+
+
+class DGLValue(ctypes.Union):
+    """DGLValue in C API"""
+
+    _fields_ = [
+        ("v_int64", ctypes.c_int64),
+        ("v_float64", ctypes.c_double),
+        ("v_handle", ctypes.c_void_p),
+        ("v_str", ctypes.c_char_p),
+        ("v_type", DGLDataType),
+        ("v_ctx", DGLContext),
+    ]
+
+
+DGLPackedCFunc = ctypes.CFUNCTYPE(
+    ctypes.c_int,
+    ctypes.POINTER(DGLValue),
+    ctypes.POINTER(ctypes.c_int),
+    ctypes.c_int,
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+)
+
+
+DGLCFuncFinalizer = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
+
+
+def _return_handle(x):
+    """return handle"""
+    handle = x.v_handle
+    if not isinstance(handle, ctypes.c_void_p):
+        handle = ctypes.c_void_p(handle)
+    return handle
+
+
+def _return_bytes(x):
+    """return handle"""
+    handle = x.v_handle
+    if not isinstance(handle, ctypes.c_void_p):
+        handle = ctypes.c_void_p(handle)
+    arr = ctypes.cast(handle, ctypes.POINTER(DGLByteArray))[0]
+    size = arr.size
+    res = bytearray(size)
+    rptr = (ctypes.c_byte * size).from_buffer(res)
+    if not ctypes.memmove(rptr, arr.data, size):
+        raise RuntimeError("memmove failed")
+    return res
+
+
+def _wrap_arg_func(return_f, type_code):
+    tcode = ctypes.c_int(type_code)
+
+    def _wrap_func(x):
+        check_call(_LIB.DGLCbArgToReturn(ctypes.byref(x), tcode))
+        return return_f(x)
+
+    return _wrap_func
+
+
+RETURN_SWITCH = {
+    TypeCode.INT: lambda x: x.v_int64,
+    TypeCode.FLOAT: lambda x: x.v_float64,
+    TypeCode.HANDLE: _return_handle,
+    TypeCode.NULL: lambda x: None,
+    TypeCode.STR: lambda x: py_str(x.v_str),
+    TypeCode.BYTES: _return_bytes,
+    TypeCode.DGL_CONTEXT: lambda x: DGLContext(
+        x.v_ctx.device_type, x.v_ctx.device_id
+    ),
+}
+
+C_TO_PY_ARG_SWITCH = {
+    TypeCode.INT: lambda x: x.v_int64,
+    TypeCode.FLOAT: lambda x: x.v_float64,
+    TypeCode.HANDLE: _return_handle,
+    TypeCode.NULL: lambda x: None,
+    TypeCode.STR: lambda x: py_str(x.v_str),
+    TypeCode.BYTES: _return_bytes,
+    TypeCode.DGL_CONTEXT: lambda x: DGLContext(
+        x.v_ctx.device_type, x.v_ctx.device_id
+    ),
+}

python/dgl/_ffi/_cy2/__init__.py

@@ -0,0 +1 @@
+"""cython2 namespace"""

python/dgl/_ffi/_cy3/__init__.py

@@ -0,0 +1 @@
+"""cython3 namespace"""

python/dgl/_ffi/_cython/.gitignore

@@ -0,0 +1 @@
+*.cpp

python/dgl/_ffi/_cython/base.pxi

@@ -0,0 +1,169 @@
+from ..base import DGLError
+from libcpp.vector cimport vector
+from libcpp cimport bool
+from cpython.version cimport PY_MAJOR_VERSION
+from cpython cimport pycapsule
+from libc.stdint cimport int32_t, int64_t, uint64_t, uint8_t, uint16_t
+import ctypes
+
+cdef enum DGLObjectTypeCode:
+    kObjectInt = 0
+    kObjectUInt = 1
+    kObjectFloat = 2
+    kHandle = 3
+    kNull = 4
+    kDGLDataType = 5
+    kDGLContext = 6
+    kArrayHandle = 7
+    kObjectHandle = 8
+    kModuleHandle = 9
+    kFuncHandle = 10
+    kStr = 11
+    kBytes = 12
+    kNDArrayContainer = 13
+    kExtBegin = 15
+
+cdef extern from "dgl/runtime/c_runtime_api.h":
+    ctypedef struct DGLDataType:
+        uint8_t code
+        uint8_t bits
+        uint16_t lanes
+
+    ctypedef struct DGLContext:
+        int32_t device_type
+        int32_t device_id
+
+    ctypedef struct DGLArray:
+        void* data
+        DGLContext ctx
+        int32_t ndim
+        DGLDataType dtype
+        int64_t* shape
+        int64_t* strides
+        uint64_t byte_offset
+
+    ctypedef struct DLManagedTensor:
+        DGLArray dl_tensor
+        void* manager_ctx
+        void (*deleter)(DLManagedTensor* self)
+
+    ctypedef struct DGLValue:
+        int64_t v_int64
+        double v_float64
+        void* v_handle
+        const char* v_str
+        DGLDataType v_type
+        DGLContext v_ctx
+
+ctypedef int64_t dgl_index_t
+ctypedef DGLArray* DGLArrayHandle
+ctypedef void* DGLStreamHandle
+ctypedef void* DGLRetValueHandle
+ctypedef void* DGLFunctionHandle
+ctypedef void* ObjectHandle
+
+ctypedef int (*DGLPackedCFunc)(
+    DGLValue* args,
+    int* type_codes,
+    int num_args,
+    DGLRetValueHandle ret,
+    void* resource_handle)
+
+ctypedef void (*DGLPackedCFuncFinalizer)(void* resource_handle)
+
+cdef extern from "dgl/runtime/c_runtime_api.h":
+    void DGLAPISetLastError(const char* msg)
+    const char *DGLGetLastError()
+    int DGLFuncCall(DGLFunctionHandle func,
+                    DGLValue* arg_values,
+                    int* type_codes,
+                    int num_args,
+                    DGLValue* ret_val,
+                    int* ret_type_code) nogil
+    int DGLFuncFree(DGLFunctionHandle func)
+    int DGLCFuncSetReturn(DGLRetValueHandle ret,
+                          DGLValue* value,
+                          int* type_code,
+                          int num_ret)
+    int DGLFuncCreateFromCFunc(DGLPackedCFunc func,
+                               void* resource_handle,
+                               DGLPackedCFuncFinalizer fin,
+                               DGLFunctionHandle *out)
+    int DGLCbArgToReturn(DGLValue* value, int code)
+    int DGLArrayAlloc(dgl_index_t* shape,
+                      dgl_index_t ndim,
+                      DGLDataType dtype,
+                      DGLContext ctx,
+                      DGLArrayHandle* out)
+    int DGLArrayAllocSharedMem(const char *mem_name,
+                               const dgl_index_t *shape,
+                               int ndim,
+                               int dtype_code,
+                               int dtype_bits,
+                               int dtype_lanes,
+                               bool is_create,
+                               DGLArrayHandle* out)
+    int DGLArrayFree(DGLArrayHandle handle)
+    int DGLArrayCopyFromTo(DGLArrayHandle src,
+                           DGLArrayHandle to)
+
+cdef extern from "dgl/runtime/c_object_api.h":
+    int DGLObjectFree(ObjectHandle handle)
+    int DGLObjectTypeKey2Index(const char* type_key,
+                               int* out_index)
+    int DGLObjectGetTypeIndex(ObjectHandle handle,
+                              int* out_index)
+    int DGLObjectGetAttr(ObjectHandle handle,
+                         const char* key,
+                         DGLValue* out_value,
+                         int* out_type_code,
+                         int* out_success)
+
+cdef extern from "dgl/runtime/dlpack_convert.h":
+    int DGLArrayFromDLPack(DLManagedTensor* arr_from,
+                           DGLArrayHandle* out)
+    int DGLArrayToDLPack(DGLArrayHandle arr_from,
+                         DLManagedTensor** out,
+                         int alignment)
+    void DGLDLManagedTensorCallDeleter(DLManagedTensor* dltensor)
+
+cdef inline py_str(const char* x):
+    if PY_MAJOR_VERSION < 3:
+        return x
+    else:
+        return x.decode("utf-8")
+
+
+cdef inline c_str(pystr):
+    """Create ctypes char * from a python string
+    Parameters
+    ----------
+    string : string type
+        python string
+
+    Returns
+    -------
+    str : c_char_p
+        A char pointer that can be passed to C API
+    """
+    return pystr.encode("utf-8")
+
+
+cdef inline CALL(int ret):
+    if ret != 0:
+        raise DGLError(py_str(DGLGetLastError()))
+
+
+cdef inline object ctypes_handle(void* chandle):
+    """Cast C handle to ctypes handle."""
+    return ctypes.cast(<unsigned long long>chandle, ctypes.c_void_p)
+
+
+cdef inline void* c_handle(object handle):
+    """Cast C types handle to c handle."""
+    cdef unsigned long long v_ptr
+    if handle.value is None:
+        return NULL
+    else:
+        v_ptr = handle.value
+        return <void*>(v_ptr)

python/dgl/_ffi/_cython/core.pyx

@@ -0,0 +1,4 @@
+include "./base.pxi"
+include "./object.pxi"
+include "./function.pxi"
+include "./ndarray.pxi"

python/dgl/_ffi/_cython/function.pxi

@@ -0,0 +1,308 @@
+import ctypes
+import traceback
+from cpython cimport Py_INCREF, Py_DECREF
+from numbers import Number, Integral
+from ..base import string_types
+from ..object_generic import convert_to_object, ObjectGeneric
+from ..runtime_ctypes import DGLDataType as CTypesDGLDataType, \
+                             DGLContext as CTypesDGLContext, \
+                             DGLByteArray
+
+
+cdef void dgl_callback_finalize(void* fhandle):
+    local_pyfunc = <object>(fhandle)
+    Py_DECREF(local_pyfunc)
+
+cdef int dgl_callback(DGLValue* args,
+                      int* type_codes,
+                      int num_args,
+                      DGLRetValueHandle ret,
+                      void* fhandle) with gil:
+    cdef list pyargs
+    cdef DGLValue value
+    cdef int tcode
+    local_pyfunc = <object>(fhandle)
+    pyargs = []
+    for i in range(num_args):
+        value = args[i]
+        tcode = type_codes[i]
+        if (tcode == kObjectHandle or
+            tcode == kFuncHandle or
+            tcode == kModuleHandle or
+            tcode > kExtBegin):
+            CALL(DGLCbArgToReturn(&value, tcode))
+
+        if tcode != kArrayHandle:
+            pyargs.append(make_ret(value, tcode))
+        else:
+            pyargs.append(c_make_array(value.v_handle, True))
+    try:
+        rv = local_pyfunc(*pyargs)
+    except Exception:
+        msg = traceback.format_exc()
+        DGLAPISetLastError(c_str(msg))
+        return -1
+    if rv is not None:
+        if isinstance(rv, tuple):
+            raise ValueError("PackedFunction can only support one return value")
+        temp_args = []
+        make_arg(rv, &value, &tcode, temp_args)
+        CALL(DGLCFuncSetReturn(ret, &value, &tcode, 1))
+    return 0
+
+
+def convert_to_dgl_func(object pyfunc):
+    """Convert a python function to DGL function
+
+    Parameters
+    ----------
+    pyfunc : python function
+        The python function to be converted.
+
+    Returns
+    -------
+    dglfunc: dgl.Function
+        The converted dgl function.
+    """
+    cdef DGLFunctionHandle chandle
+    Py_INCREF(pyfunc)
+    CALL(DGLFuncCreateFromCFunc(dgl_callback,
+                                <void*>(pyfunc),
+                                dgl_callback_finalize,
+                                &chandle))
+    ret = _CLASS_FUNCTION(None, False)
+    (<FunctionBase>ret).chandle = chandle
+    return ret
+
+
+cdef inline int make_arg(object arg,
+                         DGLValue* value,
+                         int* tcode,
+                         list temp_args) except -1:
+    """Pack arguments into c args dgl call accept"""
+    cdef unsigned long long ptr
+    if isinstance(arg, ObjectBase):
+        value[0].v_handle = (<ObjectBase>arg).chandle
+        tcode[0] = kObjectHandle
+    elif isinstance(arg, NDArrayBase):
+        value[0].v_handle = (<NDArrayBase>arg).chandle
+        tcode[0] = (kNDArrayContainer if
+                    not (<NDArrayBase>arg).c_is_view else kArrayHandle)
+    elif isinstance(arg, _DGL_COMPATS):
+        ptr = arg._dgl_handle
+        value[0].v_handle = (<void*>ptr)
+        tcode[0] = arg.__class__._dgl_tcode
+    elif isinstance(arg, (int, long)):
+        value[0].v_int64 = arg
+        tcode[0] = kObjectInt
+    elif isinstance(arg, float):
+        value[0].v_float64 = arg
+        tcode[0] = kObjectFloat
+    elif isinstance(arg, str):
+        tstr = c_str(arg)
+        value[0].v_str = tstr
+        tcode[0] = kStr
+        temp_args.append(tstr)
+    elif arg is None:
+        value[0].v_handle = NULL
+        tcode[0] = kNull
+    elif isinstance(arg, Number):
+        value[0].v_float64 = arg
+        tcode[0] = kObjectFloat
+    elif isinstance(arg, CTypesDGLDataType):
+        tstr = c_str(str(arg))
+        value[0].v_str = tstr
+        tcode[0] = kStr
+        temp_args.append(tstr)
+    elif isinstance(arg, CTypesDGLContext):
+        value[0].v_ctx = (<DGLContext*>(
+            <unsigned long long>ctypes.addressof(arg)))[0]
+        tcode[0] = kDGLContext
+    elif isinstance(arg, bytearray):
+        arr = DGLByteArray()
+        arr.data = ctypes.cast(
+            (ctypes.c_byte * len(arg)).from_buffer(arg),
+            ctypes.POINTER(ctypes.c_byte))
+        arr.size = len(arg)
+        value[0].v_handle = <void*>(
+            <unsigned long long>ctypes.addressof(arr))
+        tcode[0] = kBytes
+        temp_args.append(arr)
+    elif isinstance(arg, string_types):
+        tstr = c_str(arg)
+        value[0].v_str = tstr
+        tcode[0] = kStr
+        temp_args.append(tstr)
+    elif isinstance(arg, (list, tuple, dict, ObjectGeneric)):
+        arg = convert_to_object(arg)
+        value[0].v_handle = (<ObjectBase>arg).chandle
+        tcode[0] = kObjectHandle
+        temp_args.append(arg)
+    #elif isinstance(arg, _CLASS_MODULE):
+    #    value[0].v_handle = c_handle(arg.handle)
+    #    tcode[0] = kModuleHandle
+    elif isinstance(arg, FunctionBase):
+        value[0].v_handle = (<FunctionBase>arg).chandle
+        tcode[0] = kFuncHandle
+    elif isinstance(arg, ctypes.c_void_p):
+        value[0].v_handle = c_handle(arg)
+        tcode[0] = kHandle
+    elif callable(arg):
+        arg = convert_to_dgl_func(arg)
+        value[0].v_handle = (<FunctionBase>arg).chandle
+        tcode[0] = kFuncHandle
+        temp_args.append(arg)
+    else:
+        raise TypeError("Don't know how to handle type %s" % type(arg))
+    return 0
+
+cdef inline bytearray make_ret_bytes(void* chandle):
+    handle = ctypes_handle(chandle)
+    arr = ctypes.cast(handle, ctypes.POINTER(DGLByteArray))[0]
+    size = arr.size
+    res = bytearray(size)
+    rptr = (ctypes.c_byte * size).from_buffer(res)
+    if not ctypes.memmove(rptr, arr.data, size):
+        raise RuntimeError('memmove failed')
+    return res
+
+cdef inline object make_ret(DGLValue value, int tcode):
+    """convert result to return value."""
+    if tcode == kObjectHandle:
+        return make_ret_object(value.v_handle)
+    elif tcode == kNull:
+        return None
+    elif tcode == kObjectInt:
+        return value.v_int64
+    elif tcode == kObjectFloat:
+        return value.v_float64
+    elif tcode == kNDArrayContainer:
+        return c_make_array(value.v_handle, False)
+    elif tcode == kStr:
+        return py_str(value.v_str)
+    elif tcode == kBytes:
+        return make_ret_bytes(value.v_handle)
+    elif tcode == kHandle:
+        return ctypes_handle(value.v_handle)
+    elif tcode == kDGLContext:
+        return CTypesDGLContext(value.v_ctx.device_type, value.v_ctx.device_id)
+    # (minjie): class module are not used in DGL.
+    #elif tcode == kModuleHandle:
+    #    return _CLASS_MODULE(ctypes_handle(value.v_handle))
+    elif tcode == kFuncHandle:
+        fobj = _CLASS_FUNCTION(None, False)
+        (<FunctionBase>fobj).chandle = value.v_handle
+        return fobj
+    elif tcode in _DGL_EXT_RET:
+        return _DGL_EXT_RET[tcode](ctypes_handle(value.v_handle))
+
+    raise ValueError("Unhandled type code %d" % tcode)
+
+
+cdef inline int FuncCall3(void* chandle,
+                          tuple args,
+                          int nargs,
+                          DGLValue* ret_val,
+                          int* ret_tcode) except -1:
+    cdef DGLValue[3] values
+    cdef int[3] tcodes
+    nargs = len(args)
+    temp_args = []
+    for i in range(nargs):
+        make_arg(args[i], &values[i], &tcodes[i], temp_args)
+    with nogil:
+        ret = DGLFuncCall(chandle, &values[0], &tcodes[0],
+                          nargs, ret_val, ret_tcode)
+    if ret != 0:
+        raise DGLError(py_str(DGLGetLastError()))
+    return 0
+
+cdef inline int FuncCall(void* chandle,
+                         tuple args,
+                         DGLValue* ret_val,
+                         int* ret_tcode) except -1:
+    cdef int nargs
+    nargs = len(args)
+    if nargs <= 3:
+        FuncCall3(chandle, args, nargs, ret_val, ret_tcode)
+        return 0
+
+    cdef vector[DGLValue] values
+    cdef vector[int] tcodes
+    values.resize(max(nargs, 1))
+    tcodes.resize(max(nargs, 1))
+    temp_args = []
+    for i in range(nargs):
+        make_arg(args[i], &values[i], &tcodes[i], temp_args)
+    with nogil:
+        ret = DGLFuncCall(chandle, &values[0], &tcodes[0],
+                          nargs, ret_val, ret_tcode)
+    if ret != 0:
+        raise DGLError(py_str(DGLGetLastError()))
+    return 0
+
+
+cdef inline int ConstructorCall(void* constructor_handle,
+                                int type_code,
+                                tuple args,
+                                void** handle) except -1:
+    """Call contructor of a handle function"""
+    cdef DGLValue ret_val
+    cdef int ret_tcode
+    FuncCall(constructor_handle, args, &ret_val, &ret_tcode)
+    assert ret_tcode == type_code
+    handle[0] = ret_val.v_handle
+    return 0
+
+
+cdef class FunctionBase:
+    cdef DGLFunctionHandle chandle
+    cdef int is_global
+
+    cdef inline _set_handle(self, handle):
+        if handle is None:
+            self.chandle = NULL
+        else:
+            self.chandle = c_handle(handle)
+
+    property is_global:
+        def __get__(self):
+            return self.c_is_global != 0
+
+        def __set__(self, value):
+            self.c_is_global = value
+
+    property handle:
+        def __get__(self):
+            if self.chandle == NULL:
+                return None
+            else:
+                return ctypes.cast(<unsigned long long>self.chandle, ctypes.c_void_p)
+        def __set__(self, value):
+            self._set_handle(value)
+
+    def __init__(self, handle, is_global):
+        self._set_handle(handle)
+        self.c_is_global = is_global
+
+    def __dealloc__(self):
+        if self.is_global == 0:
+            CALL(DGLFuncFree(self.chandle))
+
+    def __call__(self, *args):
+        cdef DGLValue ret_val
+        cdef int ret_tcode
+        FuncCall(self.chandle, args, &ret_val, &ret_tcode)
+        return make_ret(ret_val, ret_tcode)
+
+_CLASS_FUNCTION = None
+_CLASS_MODULE = None
+
+def _set_class_module(module_class):
+    """Initialize the module."""
+    global _CLASS_MODULE
+    _CLASS_MODULE = module_class
+
+def _set_class_function(func_class):
+    global _CLASS_FUNCTION
+    _CLASS_FUNCTION = func_class

python/dgl/_ffi/_cython/ndarray.pxi

@@ -0,0 +1,110 @@
+from ..runtime_ctypes import DGLArrayHandle as PyDGLArrayHandle
+from cpython cimport PyCapsule_Destructor
+
+cdef const char* _c_str_dltensor = "dltensor"
+cdef const char* _c_str_used_dltensor = "used_dltensor"
+
+
+cdef _c_dlpack_deleter(object pycaps):
+    cdef DLManagedTensor* dltensor
+    if pycapsule.PyCapsule_IsValid(pycaps, _c_str_dltensor):
+        dltensor = <DLManagedTensor*>pycapsule.PyCapsule_GetPointer(pycaps, _c_str_dltensor)
+        DGLDLManagedTensorCallDeleter(dltensor)
+
+
+def _from_dlpack(object dltensor):
+    cdef DLManagedTensor* ptr
+    cdef DGLArrayHandle chandle
+    if pycapsule.PyCapsule_IsValid(dltensor, _c_str_dltensor):
+        ptr = <DLManagedTensor*>pycapsule.PyCapsule_GetPointer(dltensor, _c_str_dltensor)
+        CALL(DGLArrayFromDLPack(ptr, &chandle))
+        # set name and destructor to be empty
+        pycapsule.PyCapsule_SetDestructor(dltensor, NULL)
+        pycapsule.PyCapsule_SetName(dltensor, _c_str_used_dltensor)
+        return c_make_array(chandle, 0)
+    raise ValueError("Expect a dltensor field, pycapsule.PyCapsule can only be consumed once")
+
+
+cdef class NDArrayBase:
+    cdef DGLArray* chandle
+    cdef int c_is_view
+
+    cdef inline _set_handle(self, handle):
+        cdef unsigned long long ptr
+        if handle is None:
+            self.chandle = NULL
+        else:
+            ptr = ctypes.cast(handle, ctypes.c_void_p).value
+            self.chandle = <DGLArray*>(ptr)
+
+    property _dgl_handle:
+        def __get__(self):
+            return <unsigned long long>self.chandle
+
+    property handle:
+        def __get__(self):
+            if self.chandle == NULL:
+                return None
+            else:
+                return ctypes.cast(
+                    <unsigned long long>self.chandle, PyDGLArrayHandle)
+
+        def __set__(self, value):
+            self._set_handle(value)
+
+    def __init__(self, handle, is_view):
+        self._set_handle(handle)
+        self.c_is_view = is_view
+
+    def __dealloc__(self):
+        if self.c_is_view == 0:
+            CALL(DGLArrayFree(self.chandle))
+
+    def to_dlpack(self, alignment=0):
+        """Produce an array from a DLPack Tensor without copying memory
+
+        Args
+        -------
+        alignment: int, default to be 0
+        Indicates the alignment requirement when converting to dlpack. Will copy to a 
+        new tensor if the alignment requirement is not satisfied. 
+        0 means no alignment requirement.
+        
+        Returns
+        -------
+        dlpack : DLPack tensor view of the array data
+        """
+        cdef DLManagedTensor* dltensor
+        if self.c_is_view != 0:
+            raise ValueError("to_dlpack do not work with memory views")
+        CALL(DGLArrayToDLPack(self.chandle, &dltensor, alignment))
+        return pycapsule.PyCapsule_New(dltensor, _c_str_dltensor, <PyCapsule_Destructor>_c_dlpack_deleter)
+
+
+cdef c_make_array(void* chandle, is_view):
+    ret = _CLASS_NDARRAY(None, is_view)
+    (<NDArrayBase>ret).chandle = <DGLArray*>chandle
+    return ret
+
+
+cdef _DGL_COMPATS = ()
+
+cdef _DGL_EXT_RET = {}
+
+def _reg_extension(cls, fcreate):
+    global _DGL_COMPATS
+    _DGL_COMPATS += (cls,)
+    if fcreate:
+        _DGL_EXT_RET[cls._dgl_tcode] = fcreate
+
+
+def _make_array(handle, is_view):
+    cdef unsigned long long ptr
+    ptr = ctypes.cast(handle, ctypes.c_void_p).value
+    return c_make_array(<void*>ptr, is_view)
+
+cdef object _CLASS_NDARRAY = None
+
+def _set_class_ndarray(cls):
+    global _CLASS_NDARRAY
+    _CLASS_NDARRAY = cls

python/dgl/_ffi/_cython/object.pxi

@@ -0,0 +1,91 @@
+from ... import _api_internal
+from ..base import string_types
+from ..object_generic import _set_class_object_base
+
+"""Maps object type to its constructor"""
+OBJECT_TYPE = []
+
+def _register_object(int index, object cls):
+    """register object class"""
+    while len(OBJECT_TYPE) <= index:
+        OBJECT_TYPE.append(None)
+    OBJECT_TYPE[index] = cls
+
+
+cdef inline object make_ret_object(void* chandle):
+    global OBJECT_TYPE
+    cdef int tindex
+    cdef list object_type
+    cdef object cls
+    object_type = OBJECT_TYPE
+    CALL(DGLObjectGetTypeIndex(chandle, &tindex))
+    if tindex < len(object_type):
+        cls = object_type[tindex]
+        if cls is not None:
+            obj = cls.__new__(cls)
+        else:
+            obj = ObjectBase.__new__(ObjectBase)
+    else:
+        obj = ObjectBase.__new__(ObjectBase)
+    (<ObjectBase>obj).chandle = chandle
+    return obj
+
+
+cdef class ObjectBase:
+    cdef void* chandle
+
+    cdef _set_handle(self, handle):
+        cdef unsigned long long ptr
+        if handle is None:
+            self.chandle = NULL
+        else:
+            ptr = handle.value
+            self.chandle = <void*>(ptr)
+
+    property handle:
+        def __get__(self):
+            if self.chandle == NULL:
+                return None
+            else:
+                return ctypes_handle(self.chandle)
+
+        def __set__(self, value):
+            self._set_handle(value)
+
+    def __dealloc__(self):
+        CALL(DGLObjectFree(self.chandle))
+
+    def __getattr__(self, name):
+        cdef DGLValue ret_val
+        cdef int ret_type_code, ret_succ
+        CALL(DGLObjectGetAttr(self.chandle, c_str(name),
+                            &ret_val, &ret_type_code, &ret_succ))
+        if ret_succ == 0:
+            raise AttributeError(
+                "'%s' object has no attribute '%s'" % (type(self), name))
+        return make_ret(ret_val, ret_type_code)
+
+    def __init_handle_by_constructor__(self, fconstructor, *args):
+        """Initialize the handle by calling constructor function.
+
+        Parameters
+        ----------
+        fconstructor : Function
+            Constructor function.
+
+        args: list of objects
+            The arguments to the constructor
+
+        Note
+        ----
+        We have a special calling convention to call constructor functions.
+        So the return handle is directly set into the Object object
+        instead of creating a new Object.
+        """
+        cdef void* chandle
+        ConstructorCall(
+            (<FunctionBase>fconstructor).chandle,
+            kObjectHandle, args, &chandle)
+        self.chandle = chandle
+
+_set_class_object_base(ObjectBase)

python/dgl/_ffi/base.py

@@ -0,0 +1,155 @@
+# coding: utf-8
+# pylint: disable=invalid-name
+"""ctypes library and helper functions """
+from __future__ import absolute_import
+
+import ctypes
+import logging
+import os
+import sys
+
+import numpy as np
+
+from . import libinfo
+
+# ----------------------------
+# library loading
+# ----------------------------
+if sys.version_info[0] == 3:
+    string_types = (str,)
+    numeric_types = (float, int, np.float32, np.int32)
+    # this function is needed for python3
+    # to convert ctypes.char_p .value back to python str
+    py_str = lambda x: x.decode("utf-8")
+else:
+    string_types = (basestring,)
+    numeric_types = (float, int, long, np.float32, np.int32)
+    py_str = lambda x: x
+
+
+class DGLError(Exception):
+    """Error thrown by DGL function"""
+
+    pass  # pylint: disable=unnecessary-pass
+
+
+def _load_lib():
+    """Load libary by searching possible path."""
+    lib_path = libinfo.find_lib_path()
+    lib = ctypes.CDLL(lib_path[0])
+    dirname = os.path.dirname(lib_path[0])
+    basename = os.path.basename(lib_path[0])
+    # DMatrix functions
+    lib.DGLGetLastError.restype = ctypes.c_char_p
+    return lib, basename, dirname
+
+
+# version number
+__version__ = libinfo.__version__
+# library instance of nnvm
+_LIB, _LIB_NAME, _DIR_NAME = _load_lib()
+
+# The FFI mode of DGL
+_FFI_MODE = os.environ.get("DGL_FFI", "auto")
+
+# ----------------------------
+# helper function in ctypes.
+# ----------------------------
+def check_call(ret):
+    """Check the return value of C API call
+
+    This function will raise exception when error occurs.
+    Wrap every API call with this function
+
+    Parameters
+    ----------
+    ret : int
+        return value from API calls
+    """
+    if ret != 0:
+        raise DGLError(py_str(_LIB.DGLGetLastError()))
+
+
+def c_str(string):
+    """Create ctypes char * from a python string
+    Parameters
+    ----------
+    string : string type
+        python string
+
+    Returns
+    -------
+    str : c_char_p
+        A char pointer that can be passed to C API
+    """
+    return ctypes.c_char_p(string.encode("utf-8"))
+
+
+def c_array(ctype, values):
+    """Create ctypes array from a python array
+
+    Parameters
+    ----------
+    ctype : ctypes data type
+        data type of the array we want to convert to
+
+    values : tuple or list
+        data content
+
+    Returns
+    -------
+    out : ctypes array
+        Created ctypes array
+    """
+    return (ctype * len(values))(*values)
+
+
+def decorate(func, fwrapped):
+    """A wrapper call of decorator package, differs to call time
+
+    Parameters
+    ----------
+    func : function
+        The original function
+
+    fwrapped : function
+        The wrapped function
+    """
+    import decorator
+
+    return decorator.decorate(func, fwrapped)
+
+
+tensor_adapter_loaded = False
+
+
+def load_tensor_adapter(backend, version):
+    """Tell DGL to load a tensoradapter library for given backend and version.
+
+    Parameters
+    ----------
+    backend : str
+        The backend (currently ``pytorch``, ``mxnet`` or ``tensorflow``).
+    version : str
+        The version number of the backend.
+    """
+    global tensor_adapter_loaded
+    version = version.split("+")[0]
+    if sys.platform.startswith("linux"):
+        basename = "libtensoradapter_%s_%s.so" % (backend, version)
+    elif sys.platform.startswith("darwin"):
+        basename = "libtensoradapter_%s_%s.dylib" % (backend, version)
+    elif sys.platform.startswith("win"):
+        basename = "tensoradapter_%s_%s.dll" % (backend, version)
+    else:
+        raise NotImplementedError("Unsupported system: %s" % sys.platform)
+    path = os.path.join(_DIR_NAME, "tensoradapter", backend, basename)
+    tensor_adapter_loaded = _LIB.DGLLoadTensorAdapter(path.encode("utf-8")) == 0
+    if not tensor_adapter_loaded:
+        logger = logging.getLogger("dgl-core")
+        logger.debug("Memory optimization with PyTorch is not enabled.")
+
+
+def is_tensor_adaptor_enabled() -> bool:
+    """Check whether TensorAdaptor is enabled."""
+    return tensor_adapter_loaded

python/dgl/_ffi/capi.py

@@ -0,0 +1,4 @@
+"""Init all C APIs in the default namespace."""
+from .function import _init_api
+
+__all__ = _init_api("dgl.capi", __name__)

python/dgl/_ffi/function.py

@@ -0,0 +1,350 @@
+# pylint: disable=invalid-name, unused-import
+"""Function namespace."""
+from __future__ import absolute_import
+
+import ctypes
+import sys
+
+from .base import _FFI_MODE, _LIB, c_str, check_call, py_str, string_types
+
+IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError
+
+try:
+    # pylint: disable=wrong-import-position
+    if _FFI_MODE == "ctypes":
+        raise ImportError()
+    if sys.version_info >= (3, 0):
+        from ._cy3.core import (
+            _set_class_function,
+            _set_class_module,
+            convert_to_dgl_func,
+            FunctionBase as _FunctionBase,
+        )
+    else:
+        from ._cy2.core import (
+            _set_class_function,
+            _set_class_module,
+            convert_to_dgl_func,
+            FunctionBase as _FunctionBase,
+        )
+except IMPORT_EXCEPT:
+    # pylint: disable=wrong-import-position
+    from ._ctypes.function import (
+        _set_class_function,
+        _set_class_module,
+        convert_to_dgl_func,
+        FunctionBase as _FunctionBase,
+    )
+
+FunctionHandle = ctypes.c_void_p
+
+
+class Function(_FunctionBase):
+    """The PackedFunc object.
+
+    Function plays an key role to bridge front and backend in DGL.
+    Function provide a type-erased interface, you can call function with positional arguments.
+
+    The compiled module returns Function.
+    DGL backend also registers and exposes its API as Functions.
+    For example, the developer function exposed in dgl.ir_pass are actually
+    C++ functions that are registered as PackedFunc
+
+    The following are list of common usage scenario of dgl.Function.
+
+    - Automatic exposure of C++ API into python
+    - To call PackedFunc from python side
+    - To call python callbacks to inspect results in generated code
+    - Bring python hook into C++ backend
+
+    See Also
+    --------
+    dgl.register_func: How to register global function.
+    dgl.get_global_func: How to get global function.
+    """
+
+    pass  # pylint: disable=unnecessary-pass
+
+
+class ModuleBase(object):
+    """Base class for module"""
+
+    __slots__ = ["handle", "_entry", "entry_name"]
+
+    def __init__(self, handle):
+        self.handle = handle
+        self._entry = None
+        self.entry_name = "__dgl_main__"
+
+    def __del__(self):
+        check_call(_LIB.DGLModFree(self.handle))
+
+    @property
+    def entry_func(self):
+        """Get the entry function
+
+        Returns
+        -------
+        f : Function
+            The entry function if exist
+        """
+        if self._entry:
+            return self._entry
+        self._entry = self.get_function(self.entry_name)
+        return self._entry
+
+    def get_function(self, name, query_imports=False):
+        """Get function from the module.
+
+        Parameters
+        ----------
+        name : str
+            The name of the function
+
+        query_imports : bool
+            Whether also query modules imported by this module.
+
+        Returns
+        -------
+        f : Function
+            The result function.
+        """
+        ret_handle = FunctionHandle()
+        check_call(
+            _LIB.DGLModGetFunction(
+                self.handle,
+                c_str(name),
+                ctypes.c_int(query_imports),
+                ctypes.byref(ret_handle),
+            )
+        )
+        if not ret_handle.value:
+            raise AttributeError("Module has no function '%s'" % name)
+        return Function(ret_handle, False)
+
+    def import_module(self, module):
+        """Add module to the import list of current one.
+
+        Parameters
+        ----------
+        module : Module
+            The other module.
+        """
+        check_call(_LIB.DGLModImport(self.handle, module.handle))
+
+    def __getitem__(self, name):
+        if not isinstance(name, string_types):
+            raise ValueError("Can only take string as function name")
+        return self.get_function(name)
+
+    def __call__(self, *args):
+        if self._entry:
+            return self._entry(*args)
+        f = self.entry_func
+        return f(*args)
+
+
+def register_func(func_name, f=None, override=False):
+    """Register global function
+
+    Parameters
+    ----------
+    func_name : str or function
+        The function name
+
+    f : function, optional
+        The function to be registered.
+
+    override: boolean optional
+        Whether override existing entry.
+
+    Returns
+    -------
+    fregister : function
+        Register function if f is not specified.
+
+    Examples
+    --------
+    The following code registers my_packed_func as global function.
+    Note that we simply get it back from global function table to invoke
+    it from python side. However, we can also invoke the same function
+    from C++ backend, or in the compiled DGL code.
+
+    .. code-block:: python
+
+      targs = (10, 10.0, "hello")
+      @dgl.register_func
+      def my_packed_func(*args):
+          assert(tuple(args) == targs)
+          return 10
+      # Get it out from global function table
+      f = dgl.get_global_func("my_packed_func")
+      assert isinstance(f, dgl.nd.Function)
+      y = f(*targs)
+      assert y == 10
+    """
+    if callable(func_name):
+        f = func_name
+        func_name = f.__name__
+
+    if not isinstance(func_name, str):
+        raise ValueError("expect string function name")
+
+    ioverride = ctypes.c_int(override)
+
+    def register(myf):
+        """internal register function"""
+        if not isinstance(myf, Function):
+            myf = convert_to_dgl_func(myf)
+        check_call(
+            _LIB.DGLFuncRegisterGlobal(c_str(func_name), myf.handle, ioverride)
+        )
+        return myf
+
+    if f:
+        return register(f)
+    return register
+
+
+def get_global_func(name, allow_missing=False):
+    """Get a global function by name
+
+    Parameters
+    ----------
+    name : str
+        The name of the global function
+
+    allow_missing : bool
+        Whether allow missing function or raise an error.
+
+    Returns
+    -------
+    func : dgl.Function
+        The function to be returned, None if function is missing.
+    """
+    handle = FunctionHandle()
+    check_call(_LIB.DGLFuncGetGlobal(c_str(name), ctypes.byref(handle)))
+    if handle.value:
+        return Function(handle, False)
+    else:
+        if allow_missing:
+            return None
+        else:
+            raise ValueError("Cannot find global function %s" % name)
+
+
+def list_global_func_names():
+    """Get list of global functions registered.
+
+    Returns
+    -------
+    names : list
+       List of global functions names.
+    """
+    plist = ctypes.POINTER(ctypes.c_char_p)()
+    size = ctypes.c_uint()
+
+    check_call(
+        _LIB.DGLFuncListGlobalNames(ctypes.byref(size), ctypes.byref(plist))
+    )
+    fnames = []
+    for i in range(size.value):
+        fnames.append(py_str(plist[i]))
+    return fnames
+
+
+def extract_ext_funcs(finit):
+    """
+    Extract the extension PackedFuncs from a C module.
+
+    Parameters
+    ----------
+    finit : ctypes function
+        a ctypes that takes signature of DGLExtensionDeclarer
+
+    Returns
+    -------
+    fdict : dict of str to Function
+        The extracted functions
+    """
+    fdict = {}
+
+    def _list(name, func):
+        fdict[name] = func
+
+    myf = convert_to_dgl_func(_list)
+    ret = finit(myf.handle)
+    _ = myf
+    if ret != 0:
+        raise RuntimeError("cannot initialize with %s" % finit)
+    return fdict
+
+
+def _get_api(f):
+    flocal = f
+    flocal.is_global = True
+    return flocal
+
+
+def _init_api(namespace, target_module_name=None):
+    """Initialize api for a given module name
+
+    namespace : str
+       The namespace of the source registry
+
+    target_module_name : str
+       The target module name if different from namespace
+    """
+    target_module_name = target_module_name if target_module_name else namespace
+    if namespace.startswith("dgl."):
+        return _init_api_prefix(target_module_name, namespace[4:])
+    else:
+        return _init_api_prefix(target_module_name, namespace)
+
+
+def _init_api_prefix(module_name, prefix):
+    module = sys.modules[module_name]
+    name_list = []
+
+    for name in list_global_func_names():
+        if name.startswith("_") and not name.startswith("_deprecate"):
+            # internal APIs are ignored
+            continue
+        name_split = name.rsplit(".", 1)
+        if name_split[0] != prefix:
+            continue
+
+        if len(name_split) == 1:
+            print('Warning: invalid API name "%s".' % name)
+            continue
+        fname = name_split[1]
+        target_module = module
+
+        f = get_global_func(name)
+        ff = _get_api(f)
+        ff.__name__ = fname
+        ff.__doc__ = "DGL PackedFunc %s. " % fname
+        setattr(target_module, ff.__name__, ff)
+        name_list.append(fname)
+
+    return name_list
+
+
+def _init_internal_api():
+    for name in list_global_func_names():
+        if not name.startswith("_") or name.startswith("_deprecate"):
+            # normal APIs are ignored
+            continue
+        target_module = sys.modules["dgl._api_internal"]
+        fname = name
+        if fname.find(".") != -1:
+            print('Warning: invalid API name "%s".' % fname)
+            continue
+        f = get_global_func(name)
+        ff = _get_api(f)
+        ff.__name__ = fname
+        ff.__doc__ = "DGL PackedFunc %s. " % fname
+        setattr(target_module, ff.__name__, ff)
+
+
+_set_class_function(Function)

python/dgl/_ffi/libinfo.py

@@ -0,0 +1,108 @@
+"""Library information."""
+from __future__ import absolute_import
+
+import os
+import pathlib
+import sys
+
+
+def find_lib_path(name=None, search_path=None, optional=False):
+    """Find dynamic library files.
+
+    Parameters
+    ----------
+    name : list of str
+        List of names to be found.
+
+    Returns
+    -------
+    lib_path : list(string)
+        List of all found path to the libraries
+    """
+    # See https://github.com/dmlc/tvm/issues/281 for some background.
+
+    # NB: This will either be the source directory (if DGL is run
+    # inplace) or the install directory (if DGL is installed).
+    # An installed DGL's curr_path will look something like:
+    #   $PREFIX/lib/python3.6/site-packages/dgl/_ffi
+    ffi_dir = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
+    source_dir = os.path.join(ffi_dir, "..", "..", "..")
+    install_lib_dir = os.path.join(ffi_dir, "..", "..", "..", "..")
+
+    dll_path = []
+
+    if os.environ.get("DGL_LIBRARY_PATH", None):
+        dll_path.append(os.environ["DGL_LIBRARY_PATH"])
+
+    if sys.platform.startswith("linux") and os.environ.get(
+        "LD_LIBRARY_PATH", None
+    ):
+        dll_path.extend(
+            [p.strip() for p in os.environ["LD_LIBRARY_PATH"].split(":")]
+        )
+    elif sys.platform.startswith("darwin") and os.environ.get(
+        "DYLD_LIBRARY_PATH", None
+    ):
+        dll_path.extend(
+            [p.strip() for p in os.environ["DYLD_LIBRARY_PATH"].split(":")]
+        )
+
+    # Pip lib directory
+    dll_path.append(os.path.join(ffi_dir, ".."))
+    # Default cmake build directory
+    dll_path.append(os.path.join(source_dir, "build"))
+    dll_path.append(os.path.join(source_dir, "build", "Release"))
+    # Default make build directory
+    dll_path.append(os.path.join(source_dir, "lib"))
+
+    dll_path.append(install_lib_dir)
+
+    if search_path is not None:
+        if isinstance(search_path, (list, tuple, set)):
+            dll_path = dll_path + list(search_path)
+        elif isinstance(search_path, str):
+            dll_path.append(search_path)
+        else:
+            raise ValueError(
+                "type(search_path)={} is invalid".format(type(search_path))
+            )
+    dll_path = [
+        str(x.absolute()) if isinstance(x, pathlib.Path) else os.path.abspath(x)
+        for x in dll_path
+    ]
+
+    if name is None:
+        if sys.platform.startswith("win32"):
+            name = ["libdgl.dll", "dgl.dll"]
+        elif sys.platform.startswith("darwin"):
+            name = "libdgl.dylib"
+        else:
+            name = "libdgl.so"
+
+    if isinstance(name, str):
+        name = [name]
+
+    lib_dll_path = []
+    for n in name:
+        lib_dll_path += [os.path.join(p, n) for p in dll_path]
+
+    lib_found = [p for p in lib_dll_path if os.path.isfile(p)]
+
+    if not lib_found:
+        message = (
+            "Cannot find the files.\n"
+            + "List of candidates:\n"
+            + str("\n".join(lib_dll_path))
+        )
+        if not optional:
+            raise RuntimeError(message)
+        return None
+
+    return lib_found
+
+
+# current version
+# We use the version of the incoming release for code
+# that is under development.
+# The following line is set by dgl/python/update_version.py
+__version__ = "2.5"

python/dgl/_ffi/ndarray.py

@@ -0,0 +1,448 @@
+# pylint: disable=invalid-name, unused-import
+"""Runtime NDArray api"""
+from __future__ import absolute_import
+
+import ctypes
+import sys
+
+import numpy as np
+
+from .base import _FFI_MODE, _LIB, c_array, c_str, check_call, string_types
+from .runtime_ctypes import (
+    dgl_shape_index_t,
+    DGLArray,
+    DGLArrayHandle,
+    DGLContext,
+    DGLDataType,
+    TypeCode,
+)
+
+IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError
+
+try:
+    # pylint: disable=wrong-import-position
+    if _FFI_MODE == "ctypes":
+        raise ImportError()
+    if sys.version_info >= (3, 0):
+        from ._cy3.core import (
+            _from_dlpack,
+            _make_array,
+            _reg_extension,
+            _set_class_ndarray,
+            NDArrayBase as _NDArrayBase,
+        )
+    else:
+        from ._cy2.core import (
+            _from_dlpack,
+            _make_array,
+            _reg_extension,
+            _set_class_ndarray,
+            NDArrayBase as _NDArrayBase,
+        )
+except IMPORT_EXCEPT:
+    # pylint: disable=wrong-import-position
+    from ._ctypes.ndarray import (
+        _from_dlpack,
+        _make_array,
+        _reg_extension,
+        _set_class_ndarray,
+        NDArrayBase as _NDArrayBase,
+    )
+
+
+def context(dev_type, dev_id=0):
+    """Construct a DGL context with given device type and id.
+
+    Parameters
+    ----------
+    dev_type: int or str
+        The device type mask or name of the device.
+
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    ctx: DGLContext
+        The corresponding context.
+
+    Examples
+    --------
+    Context can be used to create reflection of context by
+    string representation of the device type.
+
+    .. code-block:: python
+
+      assert dgl.context("cpu", 1) == dgl.cpu(1)
+      assert dgl.context("gpu", 0) == dgl.gpu(0)
+      assert dgl.context("cuda", 0) == dgl.gpu(0)
+    """
+    if isinstance(dev_type, string_types):
+        dev_type = dev_type.split()[0]
+        if dev_type not in DGLContext.STR2MASK:
+            raise ValueError("Unknown device type %s" % dev_type)
+        dev_type = DGLContext.STR2MASK[dev_type]
+    return DGLContext(dev_type, dev_id)
+
+
+def numpyasarray(np_data):
+    """Return a DGLArray representation of a numpy array."""
+    data = np_data
+    assert data.flags["C_CONTIGUOUS"]
+    arr = DGLArray()
+    shape = c_array(dgl_shape_index_t, data.shape)
+    arr.data = data.ctypes.data_as(ctypes.c_void_p)
+    arr.shape = shape
+    arr.strides = None
+    arr.dtype = DGLDataType(np.dtype(data.dtype).name)
+    arr.ndim = data.ndim
+    # CPU device
+    arr.ctx = context(1, 0)
+    return arr, shape
+
+
+def empty(shape, dtype="float32", ctx=context(1, 0)):
+    """Create an empty array given shape and device
+
+    Parameters
+    ----------
+    shape : tuple of int
+        The shape of the array
+
+    dtype : type or str
+        The data type of the array.
+
+    ctx : DGLContext
+        The context of the array
+
+    Returns
+    -------
+    arr : dgl.nd.NDArray
+        The array dgl supported.
+    """
+    shape = c_array(dgl_shape_index_t, shape)
+    ndim = ctypes.c_int(len(shape))
+    handle = DGLArrayHandle()
+    dtype = DGLDataType(dtype)
+    check_call(
+        _LIB.DGLArrayAlloc(
+            shape,
+            ndim,
+            ctypes.c_int(dtype.type_code),
+            ctypes.c_int(dtype.bits),
+            ctypes.c_int(dtype.lanes),
+            ctx.device_type,
+            ctx.device_id,
+            ctypes.byref(handle),
+        )
+    )
+    return _make_array(handle, False)
+
+
+def empty_shared_mem(name, is_create, shape, dtype="float32"):
+    """Create an empty array with shared memory given shape and dtype
+
+    Parameters
+    ----------
+    name : string
+        The name of the shared memory. It's a file name in Unix.
+
+    is_create : bool
+        Whether to create the shared memory or use the one created by somewhere else.
+
+    shape : tuple of int
+        The shape of the array
+
+    dtype : type or str
+        The data type of the array.
+
+    Returns
+    -------
+    arr : dgl.nd.NDArray
+        The array dgl supported.
+    """
+    name = ctypes.c_char_p(name.encode("utf-8"))
+    shape = c_array(dgl_shape_index_t, shape)
+    ndim = ctypes.c_int(len(shape))
+    handle = DGLArrayHandle()
+    dtype = DGLDataType(dtype)
+    check_call(
+        _LIB.DGLArrayAllocSharedMem(
+            name,
+            shape,
+            ndim,
+            ctypes.c_int(dtype.type_code),
+            ctypes.c_int(dtype.bits),
+            ctypes.c_int(dtype.lanes),
+            is_create,
+            ctypes.byref(handle),
+        )
+    )
+    return _make_array(handle, False)
+
+
+def from_dlpack(dltensor):
+    """Produce an array from a DLPack tensor without memory copy.
+    Retrieves the underlying DLPack tensor's pointer to create an array from the
+    data. Removes the original DLPack tensor's destructor as now the array is
+    responsible for destruction.
+
+    Parameters
+    ----------
+    dltensor : DLPack tensor
+        Input DLManagedTensor, can only be consumed once.
+
+    Returns
+    -------
+    arr: dgl.nd.NDArray
+        The array view of the tensor data.
+    """
+    return _from_dlpack(dltensor)
+
+
+class NDArrayBase(_NDArrayBase):
+    """A simple Device/CPU Array object in runtime."""
+
+    @property
+    def shape(self):
+        """Shape of this array"""
+        return tuple(
+            self.handle.contents.shape[i]
+            for i in range(self.handle.contents.ndim)
+        )
+
+    @property
+    def dtype(self):
+        """Type of this array"""
+        return str(self.handle.contents.dtype)
+
+    @property
+    def ctx(self):
+        """context of this array"""
+        return self.handle.contents.ctx
+
+    @property
+    def context(self):
+        """context of this array"""
+        return self.ctx
+
+    def __hash__(self):
+        return ctypes.cast(self.handle, ctypes.c_void_p).value
+
+    def __eq__(self, other):
+        return self.same_as(other)
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def same_as(self, other):
+        """Check object identity equality
+
+        Parameters
+        ----------
+        other : object
+            The other object to compare to
+
+        Returns
+        -------
+        same : bool
+            Whether other is same as self.
+        """
+        if not isinstance(other, NDArrayBase):
+            return False
+        return self.__hash__() == other.__hash__()
+
+    def __setitem__(self, in_slice, value):
+        """Set ndarray value"""
+        if (
+            not isinstance(in_slice, slice)
+            or in_slice.start is not None
+            or in_slice.stop is not None
+        ):
+            raise ValueError("Array only support set from numpy array")
+        if isinstance(value, NDArrayBase):
+            if value.handle is not self.handle:
+                value.copyto(self)
+        elif isinstance(value, (np.ndarray, np.generic)):
+            self.copyfrom(value)
+        else:
+            raise TypeError("type %s not supported" % str(type(value)))
+
+    def copyfrom(self, source_array):
+        """Perform a synchronized copy from the array.
+
+        Parameters
+        ----------
+        source_array : array_like
+            The data source we should like to copy from.
+
+        Returns
+        -------
+        arr : NDArray
+            Reference to self.
+        """
+        if isinstance(source_array, NDArrayBase):
+            source_array.copyto(self)
+            return self
+
+        if not isinstance(source_array, np.ndarray):
+            try:
+                source_array = np.asarray(source_array, dtype=self.dtype)
+            except:
+                raise TypeError(
+                    "array must be an array_like data,"
+                    + "type %s is not supported" % str(type(source_array))
+                )
+        t = DGLDataType(self.dtype)
+        shape, dtype = self.shape, self.dtype
+        if t.lanes > 1:
+            shape = shape + (t.lanes,)
+            t.lanes = 1
+            dtype = str(t)
+
+        if source_array.shape != shape:
+            raise ValueError(
+                "array shape do not match the shape of NDArray {0} vs {1}".format(
+                    source_array.shape, shape
+                )
+            )
+        source_array = np.ascontiguousarray(source_array, dtype=dtype)
+        assert source_array.flags["C_CONTIGUOUS"]
+        data = source_array.ctypes.data_as(ctypes.c_void_p)
+        nbytes = ctypes.c_size_t(
+            source_array.size * source_array.dtype.itemsize
+        )
+        check_call(_LIB.DGLArrayCopyFromBytes(self.handle, data, nbytes))
+        return self
+
+    def __repr__(self):
+        res = "dgl.{0}@{1}".format(self.asnumpy().__repr__(), self.context)
+        return res
+
+    def __str__(self):
+        return str(self.asnumpy())
+
+    def asnumpy(self):
+        """Convert this array to numpy array
+
+        Returns
+        -------
+        np_arr : numpy.ndarray
+            The corresponding numpy array.
+        """
+        t = DGLDataType(self.dtype)
+        shape, dtype = self.shape, self.dtype
+        if t.lanes > 1:
+            shape = shape + (t.lanes,)
+            t.lanes = 1
+            dtype = str(t)
+        np_arr = np.empty(shape, dtype=dtype)
+        assert np_arr.flags["C_CONTIGUOUS"]
+        data = np_arr.ctypes.data_as(ctypes.c_void_p)
+        nbytes = ctypes.c_size_t(np_arr.size * np_arr.dtype.itemsize)
+        check_call(_LIB.DGLArrayCopyToBytes(self.handle, data, nbytes))
+        return np_arr
+
+    def copyto(self, target):
+        """Copy array to target
+
+        Parameters
+        ----------
+        target : NDArray
+            The target array to be copied, must have same shape as this array.
+        """
+        if isinstance(target, DGLContext):
+            target = empty(self.shape, self.dtype, target)
+        if isinstance(target, NDArrayBase):
+            check_call(_LIB.DGLArrayCopyFromTo(self.handle, target.handle))
+        else:
+            raise ValueError("Unsupported target type %s" % str(type(target)))
+        return target
+
+    def pin_memory_(self):
+        """Pin host memory and map into GPU address space (in-place)"""
+        check_call(_LIB.DGLArrayPinData(self.handle))
+
+    def unpin_memory_(self):
+        """Unpin host memory pinned by pin_memory_()"""
+        check_call(_LIB.DGLArrayUnpinData(self.handle))
+
+    def record_stream(self, stream):
+        """Record the stream that is using this tensor.
+
+        Note
+        ----
+        This API is more for testing. Users should call ``record_stream``
+        on torch.Tensor or dgl.graph directly.
+
+        Parameters
+        ----------
+        stream : DGLStreamHandle
+        """
+        check_call(_LIB.DGLArrayRecordStream(self.handle, stream))
+
+
+def free_extension_handle(handle, type_code):
+    """Free c++ extension type handle
+
+    Parameters
+    ----------
+    handle : ctypes.c_void_p
+        The handle to the extension type.
+
+    type_code : int
+         The tyoe code
+    """
+    check_call(_LIB.DGLExtTypeFree(handle, ctypes.c_int(type_code)))
+
+
+def register_extension(cls, fcreate=None):
+    """Register a extension class to DGL.
+
+    After the class is registered, the class will be able
+    to directly pass as Function argument generated by DGL.
+
+    Parameters
+    ----------
+    cls : class
+        The class object to be registered as extension.
+
+    Note
+    ----
+    The registered class is requires one property: _dgl_handle and a class attribute _dgl_tcode.
+
+    - ```_dgl_handle``` returns integer represents the address of the handle.
+    - ```_dgl_tcode``` gives integer represents type code of the class.
+
+    Returns
+    -------
+    cls : class
+        The class being registered.
+
+    fcreate : function, optional
+        The creation function to create a class object given handle value.
+
+    Example
+    -------
+    The following code registers user defined class
+    MyTensor to be DLTensor compatible.
+
+    .. code-block:: python
+
+       @dgl.register_extension
+       class MyTensor(object):
+           _dgl_tcode = dgl.TypeCode.ARRAY_HANDLE
+
+           def __init__(self):
+               self.handle = _LIB.NewDLTensor()
+
+           @property
+           def _dgl_handle(self):
+               return self.handle.value
+    """
+    if fcreate and cls._dgl_tcode < TypeCode.EXT_BEGIN:
+        raise ValueError(
+            "Cannot register create when extension tcode is same as buildin"
+        )
+    _reg_extension(cls, fcreate)
+    return cls

python/dgl/_ffi/object.py

@@ -0,0 +1,117 @@
+"""Object namespace"""
+# pylint: disable=unused-import
+from __future__ import absolute_import
+
+import ctypes
+import sys
+
+from .. import _api_internal
+from .base import _FFI_MODE, _LIB, c_str, check_call, py_str
+from .object_generic import convert_to_object, ObjectGeneric
+
+# pylint: disable=invalid-name
+IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError
+try:
+    # pylint: disable=wrong-import-position
+    if _FFI_MODE == "ctypes":
+        raise ImportError()
+    if sys.version_info >= (3, 0):
+        from ._cy3.core import _register_object, ObjectBase as _ObjectBase
+    else:
+        from ._cy2.core import _register_object, ObjectBase as _ObjectBase
+except IMPORT_EXCEPT:
+    # pylint: disable=wrong-import-position
+    from ._ctypes.object import _register_object, ObjectBase as _ObjectBase
+
+
+def _new_object(cls):
+    """Helper function for pickle"""
+    return cls.__new__(cls)
+
+
+class ObjectBase(_ObjectBase):
+    """ObjectBase is the base class of all DGL CAPI object.
+
+    The core attribute is ``handle``, which is a C raw pointer.  It must be initialized
+    via ``__init_handle_by_constructor__``.
+
+    Note that the same handle **CANNOT** be shared across multiple ObjectBase instances.
+    """
+
+    def __dir__(self):
+        plist = ctypes.POINTER(ctypes.c_char_p)()
+        size = ctypes.c_uint()
+        check_call(
+            _LIB.DGLObjectListAttrNames(
+                self.handle, ctypes.byref(size), ctypes.byref(plist)
+            )
+        )
+        names = []
+        for i in range(size.value):
+            names.append(py_str(plist[i]))
+        return names
+
+    def __hash__(self):
+        return _api_internal._raw_ptr(self)
+
+    def __eq__(self, other):
+        return self.same_as(other)
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __reduce__(self):
+        cls = type(self)
+        return (_new_object, (cls,), self.__getstate__())
+
+    def __getstate__(self):
+        # TODO(minjie): TVM assumes that a Node (Object in DGL) can be serialized
+        #   to json. However, this is not true in DGL because DGL Object is meant
+        #   for runtime API, so it could contain binary data such as NDArray.
+        #   If this feature is required, please raise a RFC to DGL issue.
+        raise RuntimeError("__getstate__ is not supported for object type")
+
+    def __setstate__(self, state):
+        # pylint: disable=assigning-non-slot
+        # TODO(minjie): TVM assumes that a Node (Object in DGL) can be serialized
+        #   to json. However, this is not true in DGL because DGL Object is meant
+        #   for runtime API, so it could contain binary data such as NDArray.
+        #   If this feature is required, please raise a RFC to DGL issue.
+        raise RuntimeError("__setstate__ is not supported for object type")
+
+    def same_as(self, other):
+        """check object identity equality"""
+        if not isinstance(other, ObjectBase):
+            return False
+        return self.__hash__() == other.__hash__()
+
+
+def register_object(type_key=None):
+    """Decorator used to register object type
+
+    Examples
+    --------
+    >>> @register_object
+    >>> class MyObject:
+    >>> ... pass
+
+    Parameters
+    ----------
+    type_key : str or cls
+        The type key of the object
+    """
+    object_name = type_key if isinstance(type_key, str) else type_key.__name__
+
+    def register(cls):
+        """internal register function"""
+        tindex = ctypes.c_int()
+        ret = _LIB.DGLObjectTypeKey2Index(
+            c_str(object_name), ctypes.byref(tindex)
+        )
+        if ret == 0:
+            _register_object(tindex.value, cls)
+        return cls
+
+    if isinstance(type_key, str):
+        return register
+    return register(type_key)

python/dgl/_ffi/object_generic.py

@@ -0,0 +1,59 @@
+"""Common implementation of Object generic related logic"""
+# pylint: disable=unused-import
+from __future__ import absolute_import
+
+from numbers import Integral, Number
+
+from .. import _api_internal
+from .base import string_types
+
+# Object base class
+_CLASS_OBJECT_BASE = None
+
+
+def _set_class_object_base(cls):
+    global _CLASS_OBJECT_BASE
+    _CLASS_OBJECT_BASE = cls
+
+
+class ObjectGeneric(object):
+    """Base class for all classes that can be converted to object."""
+
+    def asobject(self):
+        """Convert value to object"""
+        raise NotImplementedError()
+
+
+def convert_to_object(value):
+    """Convert a python value to corresponding object type.
+
+    Parameters
+    ----------
+    value : str
+        The value to be inspected.
+
+    Returns
+    -------
+    object : Object
+        The corresponding object value.
+    """
+    if isinstance(value, _CLASS_OBJECT_BASE):
+        return value
+    if isinstance(value, (list, tuple)):
+        value = [convert_to_object(x) for x in value]
+        return _api_internal._List(*value)
+    if isinstance(value, dict):
+        vlist = []
+        for item in value.items():
+            if not isinstance(item[0], _CLASS_OBJECT_BASE) and not isinstance(
+                item[0], string_types
+            ):
+                raise ValueError(
+                    "key of map must already been a container type"
+                )
+            vlist.append(item[0])
+            vlist.append(convert_to_object(item[1]))
+        return _api_internal._Map(*vlist)
+    if isinstance(value, ObjectGeneric):
+        return value.asobject()
+    return _api_internal._Value(value)

python/dgl/_ffi/runtime_ctypes.py

@@ -0,0 +1,278 @@
+"""Common runtime ctypes."""
+# pylint: disable=invalid-name, super-init-not-called
+from __future__ import absolute_import
+
+import ctypes
+import json
+
+import numpy as np
+
+from .. import _api_internal
+from .base import _LIB, check_call
+
+dgl_shape_index_t = ctypes.c_int64
+
+
+class TypeCode(object):
+    """Type code used in API calls"""
+
+    INT = 0
+    UINT = 1
+    FLOAT = 2
+    HANDLE = 3
+    NULL = 4
+    DGL_DATA_TYPE = 5
+    DGL_CONTEXT = 6
+    ARRAY_HANDLE = 7
+    OBJECT_HANDLE = 8
+    MODULE_HANDLE = 9
+    FUNC_HANDLE = 10
+    STR = 11
+    BYTES = 12
+    NDARRAY_CONTAINER = 13
+    EXT_BEGIN = 15
+
+
+class DGLByteArray(ctypes.Structure):
+    """Temp data structure for byte array."""
+
+    _fields_ = [
+        ("data", ctypes.POINTER(ctypes.c_byte)),
+        ("size", ctypes.c_size_t),
+    ]
+
+
+class DGLDataType(ctypes.Structure):
+    """DGL datatype structure"""
+
+    _fields_ = [
+        ("type_code", ctypes.c_uint8),
+        ("bits", ctypes.c_uint8),
+        ("lanes", ctypes.c_uint16),
+    ]
+    CODE2STR = {0: "int", 1: "uint", 2: "float", 4: "handle"}
+    _cache = {}
+
+    def __new__(cls, type_str):
+        if type_str in cls._cache:
+            return cls._cache[type_str]
+
+        inst = super(DGLDataType, cls).__new__(DGLDataType)
+
+        if isinstance(type_str, np.dtype):
+            type_str = str(type_str)
+        arr = type_str.split("x")
+        head = arr[0]
+        inst.lanes = int(arr[1]) if len(arr) > 1 else 1
+        bits = 32
+
+        if head.startswith("int"):
+            inst.type_code = 0
+            head = head[3:]
+        elif head.startswith("uint"):
+            inst.type_code = 1
+            head = head[4:]
+        elif head.startswith("float"):
+            inst.type_code = 2
+            head = head[5:]
+        elif head.startswith("handle"):
+            inst.type_code = 4
+            bits = 64
+            head = ""
+        else:
+            raise ValueError("Do not know how to handle type %s" % type_str)
+        bits = int(head) if head else bits
+        inst.bits = bits
+
+        cls._cache[type_str] = inst
+        return inst
+
+    def __init__(self, type_str):
+        pass
+
+    def __repr__(self):
+        x = "%s%d" % (DGLDataType.CODE2STR[self.type_code], self.bits)
+        if self.lanes != 1:
+            x += "x%d" % self.lanes
+        return x
+
+    def __eq__(self, other):
+        return (
+            self.bits == other.bits
+            and self.type_code == other.type_code
+            and self.lanes == other.lanes
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+
+RPC_SESS_MASK = 128
+
+
+class DGLContext(ctypes.Structure):
+    """DGL context strucure."""
+
+    _fields_ = [("device_type", ctypes.c_int), ("device_id", ctypes.c_int)]
+    MASK2STR = {
+        1: "cpu",
+        2: "gpu",
+        4: "opencl",
+        5: "aocl",
+        6: "sdaccel",
+        7: "vulkan",
+        8: "metal",
+        9: "vpi",
+        10: "rocm",
+        11: "opengl",
+        12: "ext_dev",
+    }
+    STR2MASK = {
+        "llvm": 1,
+        "stackvm": 1,
+        "cpu": 1,
+        "gpu": 2,
+        "cuda": 2,
+        "nvptx": 2,
+        "cl": 4,
+        "opencl": 4,
+        "aocl": 5,
+        "aocl_sw_emu": 5,
+        "sdaccel": 6,
+        "vulkan": 7,
+        "metal": 8,
+        "vpi": 9,
+        "rocm": 10,
+        "opengl": 11,
+        "ext_dev": 12,
+    }
+    _cache = {}
+
+    def __new__(cls, device_type, device_id):
+        if (device_type, device_id) in cls._cache:
+            return cls._cache[(device_type, device_id)]
+
+        inst = super(DGLContext, cls).__new__(DGLContext)
+
+        inst.device_type = device_type
+        inst.device_id = device_id
+
+        cls._cache[(device_type, device_id)] = inst
+        return inst
+
+    def __init__(self, device_type, device_id):
+        pass
+
+    @property
+    def exist(self):
+        """Whether this device exist."""
+        return (
+            _api_internal._GetDeviceAttr(self.device_type, self.device_id, 0)
+            != 0
+        )
+
+    @property
+    def max_threads_per_block(self):
+        """Maximum number of threads on each block."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 1)
+
+    @property
+    def warp_size(self):
+        """Number of threads that executes in concurrent."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 2)
+
+    @property
+    def max_shared_memory_per_block(self):
+        """Total amount of shared memory per block in bytes."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 3)
+
+    @property
+    def compute_version(self):
+        """Get compute verison number in string.
+
+        Currently used to get compute capability of CUDA device.
+
+        Returns
+        -------
+        version : str
+            The version string in `major.minor` format.
+        """
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 4)
+
+    @property
+    def device_name(self):
+        """Return the string name of device."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 5)
+
+    @property
+    def max_clock_rate(self):
+        """Return the max clock frequency of device."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 6)
+
+    @property
+    def multi_processor_count(self):
+        """Return the number of compute units of device."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 7)
+
+    @property
+    def max_thread_dimensions(self):
+        """Return the maximum size of each thread axis
+
+        Returns
+        -------
+        dims: List of int
+            The maximum length of threadIdx.x, threadIdx.y, threadIdx.z
+        """
+        return json.loads(
+            _api_internal._GetDeviceAttr(self.device_type, self.device_id, 8)
+        )
+
+    def sync(self):
+        """Synchronize until jobs finished at the context."""
+        check_call(_LIB.DGLSynchronize(self.device_type, self.device_id, None))
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, DGLContext)
+            and self.device_id == other.device_id
+            and self.device_type == other.device_type
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __repr__(self):
+        if self.device_type >= RPC_SESS_MASK:
+            tbl_id = self.device_type / RPC_SESS_MASK - 1
+            dev_type = self.device_type % RPC_SESS_MASK
+            return "remote[%d]:%s(%d)" % (
+                tbl_id,
+                DGLContext.MASK2STR[dev_type],
+                self.device_id,
+            )
+        return "%s(%d)" % (
+            DGLContext.MASK2STR[self.device_type],
+            self.device_id,
+        )
+
+    def __hash__(self):
+        return hash((self.device_type, self.device_id))
+
+
+class DGLArray(ctypes.Structure):
+    """DGLValue in C API"""
+
+    _fields_ = [
+        ("data", ctypes.c_void_p),
+        ("ctx", DGLContext),
+        ("ndim", ctypes.c_int),
+        ("dtype", DGLDataType),
+        ("shape", ctypes.POINTER(dgl_shape_index_t)),
+        ("strides", ctypes.POINTER(dgl_shape_index_t)),
+        ("byte_offset", ctypes.c_uint64),
+    ]
+
+
+DGLArrayHandle = ctypes.POINTER(DGLArray)
+
+DGLStreamHandle = ctypes.c_void_p

python/dgl/_ffi/streams.py

@@ -0,0 +1,46 @@
+# pylint: disable=invalid-name, unused-import
+"""Runtime stream APIs which are mainly for internal test use only.
+For applications, please use PyTorch's stream management, of which DGL is aware.
+"""
+from __future__ import absolute_import
+
+import ctypes
+
+from .base import _FFI_MODE, _LIB, check_call
+from .runtime_ctypes import DGLStreamHandle
+
+
+def to_dgl_stream_handle(cuda_stream):
+    """Convert torch.cuda.Stream to DGL stream handle
+
+    Parameters
+    ----------
+    cuda_stream : torch.cuda.Stream.
+
+    Returns
+    -------
+    DGLStreamHandle
+        DGLStreamHandle of the input ``cuda_stream``.
+    """
+    return ctypes.c_void_p(cuda_stream.cuda_stream)
+
+
+def _dgl_get_stream(ctx):
+    """Get the current CUDA stream of the given DGL context.
+
+    Parameters
+    ----------
+    ctx : DGL context.
+
+    Returns
+    -------
+    DGLStreamHandle
+        DGLStreamHandle of the current CUDA stream.
+    """
+    current_cuda_stream = DGLStreamHandle()
+    check_call(
+        _LIB.DGLGetStream(
+            ctx.device_type, ctx.device_id, ctypes.byref(current_cuda_stream)
+        )
+    )
+    return current_cuda_stream

python/dgl/backend/__init__.py

@@ -0,0 +1,146 @@
+from __future__ import absolute_import
+
+import importlib
+import json
+import logging
+import os
+import sys
+
+from . import backend
+from .set_default_backend import set_default_backend
+
+_enabled_apis = set()
+
+logger = logging.getLogger("dgl-core")
+
+
+def _gen_missing_api(api, mod_name):
+    def _missing_api(*args, **kwargs):
+        raise ImportError(
+            'API "%s" is not supported by backend "%s".'
+            " You can switch to other backends by setting"
+            " the DGLBACKEND environment." % (api, mod_name)
+        )
+
+    return _missing_api
+
+
+def load_backend(mod_name):
+    # Load backend does four things:
+    # (1) Import backend framework (PyTorch, MXNet, Tensorflow, etc.)
+    # (2) Import DGL C library.  DGL imports it *after* PyTorch/MXNet/Tensorflow.  Otherwise
+    #     DGL will crash with errors like `munmap_chunk(): invalid pointer`.
+    # (3) Sets up the tensoradapter library path.
+    # (4) Import the Python wrappers of the backend framework.  DGL does this last because
+    #     it already depends on both the backend framework and the DGL C library.
+    if mod_name == "pytorch":
+        import torch
+
+        mod = torch
+    elif mod_name == "mxnet":
+        import mxnet
+
+        mod = mxnet
+    elif mod_name == "tensorflow":
+        import tensorflow
+
+        mod = tensorflow
+    else:
+        raise NotImplementedError("Unsupported backend: %s" % mod_name)
+
+    from .._ffi.base import load_tensor_adapter  # imports DGL C library
+
+    version = mod.__version__
+    load_tensor_adapter(mod_name, version)
+
+    logger.debug("Using backend: %s" % mod_name)
+    mod = importlib.import_module(".%s" % mod_name, __name__)
+    thismod = sys.modules[__name__]
+    for api in backend.__dict__.keys():
+        if api.startswith("__"):
+            # ignore python builtin attributes
+            continue
+        if api == "data_type_dict":
+            # load data type
+            if api not in mod.__dict__:
+                raise ImportError(
+                    'API "data_type_dict" is required but missing for'
+                    ' backend "%s".' % (mod_name)
+                )
+            data_type_dict = mod.__dict__[api]()
+            for name, dtype in data_type_dict.items():
+                setattr(thismod, name, dtype)
+
+            # override data type dict function
+            setattr(thismod, "data_type_dict", data_type_dict)
+
+            # for data types with aliases, treat the first listed type as
+            # the true one
+            rev_data_type_dict = {}
+            for k, v in data_type_dict.items():
+                if not v in rev_data_type_dict.keys():
+                    rev_data_type_dict[v] = k
+            setattr(thismod, "reverse_data_type_dict", rev_data_type_dict)
+            # log backend name
+            setattr(thismod, "backend_name", mod_name)
+        else:
+            # load functions
+            if api in mod.__dict__:
+                _enabled_apis.add(api)
+                setattr(thismod, api, mod.__dict__[api])
+            else:
+                setattr(thismod, api, _gen_missing_api(api, mod_name))
+
+
+def get_preferred_backend():
+    default_dir = None
+    if "DGLDEFAULTDIR" in os.environ:
+        default_dir = os.getenv("DGLDEFAULTDIR")
+    else:
+        default_dir = os.path.join(os.path.expanduser("~"), ".dgl")
+    config_path = os.path.join(default_dir, "config.json")
+    backend_name = None
+    if "DGLBACKEND" in os.environ:
+        backend_name = os.getenv("DGLBACKEND")
+    elif os.path.exists(config_path):
+        with open(config_path, "r") as config_file:
+            config_dict = json.load(config_file)
+            backend_name = config_dict.get("backend", "").lower()
+
+    if backend_name in ["tensorflow", "mxnet", "pytorch"]:
+        return backend_name
+    else:
+        print(
+            "DGL backend not selected or invalid.  "
+            "Assuming PyTorch for now.",
+            file=sys.stderr,
+        )
+        set_default_backend(default_dir, "pytorch")
+        return "pytorch"
+
+
+load_backend(get_preferred_backend())
+
+
+def is_enabled(api):
+    """Return true if the api is enabled by the current backend.
+
+    Parameters
+    ----------
+    api : str
+        The api name.
+
+    Returns
+    -------
+    bool
+        True if the API is enabled by the current backend.
+    """
+    return api in _enabled_apis
+
+
+def to_dgl_nd(data):
+    return zerocopy_to_dgl_ndarray(data)
+
+
+def from_dgl_nd(data):
+    return zerocopy_from_dgl_ndarray(data)

python/dgl/backend/mxnet/__init__.py

@@ -0,0 +1,2 @@
+from .sparse import *
+from .tensor import *

python/dgl/backend/mxnet/sparse.py

@@ -0,0 +1,558 @@
+import mxnet as mx
+import numpy as np
+from mxnet import nd
+
+from ..._sparse_ops import (
+    _bwd_segment_cmp,
+    _csrmask,
+    _csrmm,
+    _csrsum,
+    _gsddmm,
+    _gspmm,
+    _scatter_add,
+    _segment_reduce,
+)
+
+from ...base import ALL, dgl_warning, is_all
+from ...heterograph_index import create_unitgraph_from_csr
+from .tensor import (
+    asnumpy,
+    context,
+    copy_to,
+    to_backend_ctx,
+    zerocopy_from_numpy,
+)
+
+__all__ = [
+    "gspmm",
+    "gsddmm",
+    "edge_softmax",
+    "segment_reduce",
+    "scatter_add",
+    "csrmm",
+    "csrsum",
+    "csrmask",
+]
+
+
+def _scatter_nd(index, src, n_rows):
+    """Similar to PyTorch's scatter nd on first dimension."""
+    assert index.shape == src.shape
+    dgl_warning("MXNet do not support scatter_add, fallback to numpy.")
+    ctx = context(src)
+    index = asnumpy(index)
+    src = asnumpy(src)
+    shp = index.shape
+    ndim = src.ndim
+    offsets = []
+    stride = 1
+    for i in reversed(range(1, ndim)):
+        di = shp[i]
+        offset_i = np.arange(di, dtype=index.dtype)
+        offsets.append(
+            (stride * offset_i).reshape(
+                (1,) * i + (di,) + (1,) * (ndim - 1 - i)
+            )
+        )
+        stride *= di
+    if ndim > 1:
+        new_idx = index * stride + sum(offsets)
+    else:
+        new_idx = index
+    src = src.reshape(-1)
+    new_idx = new_idx.reshape(-1)
+    rst = np.zeros((stride * n_rows,), dtype=src.dtype)
+    np.add.at(rst, new_idx, src)
+    rst = rst.reshape(n_rows, *shp[1:])
+    rst = copy_to(zerocopy_from_numpy(rst), ctx)
+    return rst
+
+
+def _gather_nd(index, src):
+    """Similar to PyTorch's gather nd on first dimension."""
+    ctx = context(src)
+    shp = index.shape
+    ndim = src.ndim
+    offsets = []
+    stride = 1
+    for i in reversed(range(1, ndim)):
+        di = shp[i]
+        offset_i = nd.arange(di, dtype=index.dtype)
+        offsets.append(
+            (stride * offset_i).reshape(
+                (1,) * i + (di,) + (1,) * (ndim - 1 - i)
+            )
+        )
+        stride *= di
+    if ndim > 1:
+        new_idx = index * stride + copy_to(sum(offsets), ctx)
+    else:
+        new_idx = index
+    src = src.reshape(-1)
+    new_idx = new_idx.reshape(-1)
+    rst = nd.take(src, new_idx).reshape(shp)
+    return rst
+
+
+def _reduce_grad(grad, shape):
+    """Reduce gradient on the broadcast dimension
+    If there is broadcast in forward pass, gradients need to be reduced on
+    broadcast dimension. This function checks the input tensor shape and
+    gradient shape and perform the reduction.
+
+    Parameters
+    ----------
+    grad: Tensor
+        Gradient tensor
+    shape: tuple
+        Shape of input tensor
+
+    Returns
+    -------
+    Tensor
+    """
+    grad_shape = grad.shape[1:]
+    in_shape = shape[1:]
+    if in_shape == grad_shape:
+        # no need to reduce
+        return grad
+    num_to_squeeze = len(grad_shape) - len(in_shape)
+    # pad inshape
+    in_shape = (1,) * num_to_squeeze + in_shape
+    # pad in_shape
+    in_shape = (1,) * num_to_squeeze + in_shape
+    reduce_idx = np.nonzero(np.asarray(grad_shape) - np.asarray(in_shape))[0]
+    reduce_idx += 1  # skip batch dim
+    grad = grad.sum(axis=tuple(reduce_idx), keepdims=True)
+    return grad.reshape(shape)
+
+
+def _need_reduce_last_dim(ufeat, efeat):
+    """Indicates whether to reduce the last dimension on edges
+    in the backward pass of spmm,
+    if so, use dot instead of mul."""
+    ushp = ufeat.shape
+    eshp = efeat.shape
+    return ushp[1:-1] == eshp[1:-1] and eshp[-1] == 1 and ushp[-1] > 1
+
+
+def _muldiv(op, x):
+    return 1.0 / x if op == "div" else x
+
+
+def _addsub(op, x):
+    return -x if op == "sub" else x
+
+
+def _expand(x, shape):
+    return x.broadcast_to((x.shape[0], *shape))
+
+
+class GSpMM(mx.autograd.Function):
+    def __init__(self, gidx, op, reduce_op):
+        super(GSpMM, self).__init__()
+        self.gidx = gidx
+        self.op = op
+        self.reduce_op = reduce_op
+
+    def forward(self, X, Y):
+        out, (argX, argY) = _gspmm(self.gidx, self.op, self.reduce_op, X, Y)
+        self.save_for_backward(X, Y, argX, argY)
+        return out
+
+    def backward(self, dZ):
+        ctx = context(dZ)
+        X, Y, argX, argY = self.saved_tensors
+        gidx, op, reduce_op = self.gidx, self.op, self.reduce_op
+        if op != "copy_rhs":
+            g_rev = gidx.reverse()
+            if reduce_op == "sum":
+                if op in ["mul", "div"]:
+                    dX = _gspmm(g_rev, "mul", "sum", dZ, _muldiv(op, Y))[0]
+                elif op in ["add", "sub"]:
+                    dX = _gspmm(g_rev, "copy_lhs", "sum", dZ, Y)[0]
+                elif op == "copy_lhs":
+                    dX = _gspmm(g_rev, "copy_lhs", "sum", dZ, None)[0]
+            else:
+                if op in ["mul", "div"]:
+                    dX = _scatter_nd(
+                        argX,
+                        _muldiv(op, _gather_nd(argY, _expand(Y, dZ.shape[1:])))
+                        * dZ,
+                        X.shape[0],
+                    )
+                elif op in ["add", "sub", "copy_lhs"]:
+                    dX = _scatter_nd(argX, dZ, X.shape[0])
+            dX = _reduce_grad(dX, X.shape)
+        else:
+            dX = nd.zeros_like(X)
+        if op != "copy_lhs":
+            if reduce_op == "sum":
+                if op == "mul" and _need_reduce_last_dim(X, Y):
+                    dY = _gsddmm(gidx, "dot", X, dZ)
+                elif op in ["mul", "div"]:
+                    dY = _gsddmm(gidx, "mul", X, dZ)
+                    if op == "div":
+                        dY = -dY / (Y**2)
+                elif op in ["add", "sub", "copy_rhs"]:
+                    dY = _gsddmm(gidx, "copy_rhs", X, _addsub(op, dZ))
+            else:
+                if op in ["mul", "div"]:
+                    dY = _scatter_nd(
+                        argY,
+                        _gather_nd(argX, _expand(X, dZ.shape[1:])) * dZ,
+                        Y.shape[0],
+                    )
+                    if op == "div":
+                        dY = -dY / (Y**2)
+                elif op in ["add", "sub", "copy_rhs"]:
+                    dY = _scatter_nd(argY, _addsub(op, dZ), Y.shape[0])
+            dY = _reduce_grad(dY, Y.shape)
+        else:
+            dY = nd.zeros_like(Y)
+        self.saved_tensors = None
+        return dX, dY
+
+
+def gspmm(gidx, op, reduce_op, lhs_data, rhs_data):
+    func = GSpMM(gidx, op, reduce_op)
+    ctx = to_backend_ctx(gidx.ctx)
+    # XXX(minjie): There is a bug in MXNet's autograd system when one of the inputs
+    #   does not require gradient. Although it still invokes the backward function,
+    #   it does not set the gradient value to the correct buffer, resulting all the
+    #   input gradients to be zero. Fix this by enforcing all the inputs to require
+    #   gradients.
+    if lhs_data is None:
+        lhs_data = nd.zeros((1,), ctx=ctx)
+        lhs_data.attach_grad()
+    if rhs_data is None:
+        rhs_data = nd.zeros((1,), ctx=ctx)
+        rhs_data.attach_grad()
+    return func(lhs_data, rhs_data)
+
+
+class GSDDMM(mx.autograd.Function):
+    def __init__(self, gidx, op, lhs_target, rhs_target):
+        super(GSDDMM, self).__init__()
+        self.gidx = gidx
+        self.op = op
+        self.lhs_target = lhs_target
+        self.rhs_target = rhs_target
+
+    def forward(self, X, Y):
+        out = _gsddmm(
+            self.gidx, self.op, X, Y, self.lhs_target, self.rhs_target
+        )
+        self.save_for_backward(X, Y)
+        return out
+
+    def backward(self, dZ):
+        ctx = context(dZ)
+        X, Y = self.saved_tensors
+        gidx, op = self.gidx, self.op
+        lhs_target, rhs_target = self.lhs_target, self.rhs_target
+        if op != "copy_rhs":
+            if lhs_target in ["u", "v"]:
+                _gidx = gidx if self.lhs_target == "v" else gidx.reverse()
+                if op in ["add", "sub", "copy_lhs"]:
+                    dX = _gspmm(_gidx, "copy_rhs", "sum", None, dZ)[0]
+                else:  # mul, div, dot
+                    if rhs_target == lhs_target:
+                        dX = _gspmm(_gidx, "copy_rhs", "sum", None, dZ)[
+                            0
+                        ] * _muldiv(op, Y)
+                    elif self.rhs_target == "e":
+                        dX = _gspmm(
+                            _gidx, "copy_rhs", "sum", None, dZ * _muldiv(op, Y)
+                        )[0]
+                    else:  # rhs_target = !lhs_target
+                        dX = _gspmm(_gidx, "mul", "sum", _muldiv(op, Y), dZ)[0]
+            else:  # lhs_target == 'e'
+                if op in ["add", "sub", "copy_lhs"]:
+                    dX = dZ
+                else:  # mul, div, dot
+                    dX = _gsddmm(
+                        gidx, "mul", dZ, _muldiv(op, Y), "e", rhs_target
+                    )
+            dX = _reduce_grad(dX, X.shape)
+        else:
+            dX = nd.zeros_like(X)
+        if op != "copy_lhs":
+            if self.rhs_target in ["u", "v"]:
+                _gidx = gidx if rhs_target == "v" else gidx.reverse()
+                if op in ["add", "sub", "copy_rhs"]:
+                    dY = _gspmm(
+                        _gidx, "copy_rhs", "sum", None, _addsub(op, dZ)
+                    )[0]
+                else:  # mul, div, dot
+                    if lhs_target == rhs_target:
+                        dY = _gspmm(_gidx, "copy_rhs", "sum", None, dZ)[0] * X
+                    elif self.lhs_target == "e":
+                        dY = _gspmm(_gidx, "copy_rhs", "sum", None, dZ * X)[0]
+                    else:  # rhs_target = !lhs_target
+                        dY = _gspmm(_gidx, "mul", "sum", X, dZ)[0]
+                    if op == "div":
+                        dY = -dY / (Y**2)
+            else:
+                if op in ["add", "sub", "copy_rhs"]:
+                    dY = _addsub(op, dZ)
+                else:  # mul, div, dot
+                    dY = _gsddmm(gidx, "mul", dZ, X, "e", lhs_target)
+                    if op == "div":
+                        dY = -dY / (Y**2)
+            dY = _reduce_grad(dY, Y.shape)
+        else:
+            dY = nd.zeros_like(Y)
+        self.saved_tensors = None
+        return dX, dY
+
+
+def gsddmm(gidx, op, lhs_data, rhs_data, lhs_target="u", rhs_target="v"):
+    func = GSDDMM(gidx, op, lhs_target, rhs_target)
+    ctx = to_backend_ctx(gidx.ctx)
+    if lhs_data is None:
+        lhs_data = nd.zeros((1,), ctx=ctx)
+    if rhs_data is None:
+        rhs_data = nd.zeros((1,), ctx=ctx)
+    return func(lhs_data, rhs_data)
+
+
+class EdgeSoftmax(mx.autograd.Function):
+    def __init__(self, gidx, eids, norm_by):
+        super(EdgeSoftmax, self).__init__()
+        if not is_all(eids):
+            gidx = gidx.edge_subgraph([eids], True).graph
+        if norm_by == "src":
+            gidx = gidx.reverse()
+        self.gidx = gidx
+
+    def forward(self, score):
+        """Forward function.
+
+        Pseudo-code:
+
+        .. code:: python
+
+            score = dgl.EData(g, score)
+            score_max = score.dst_max()  # of type dgl.NData
+            score = score - score_max  # edge_sub_dst, ret dgl.EData
+            score_sum = score.dst_sum()  # of type dgl.NData
+            out = score / score_sum    # edge_div_dst, ret dgl.EData
+            return out.data
+        """
+        gidx = self.gidx
+        score_max = _gspmm(gidx, "copy_rhs", "max", None, score)[0]
+        score = mx.nd.exp(_gsddmm(gidx, "sub", score, score_max, "e", "v"))
+        score_sum = _gspmm(gidx, "copy_rhs", "sum", None, score)[0]
+        out = _gsddmm(gidx, "div", score, score_sum, "e", "v")
+        self.save_for_backward(out)
+        return out
+
+    def backward(self, grad_out):
+        """Backward function.
+
+        Pseudo-code:
+
+        .. code:: python
+
+            g, out = ctx.backward_cache
+            grad_out = dgl.EData(g, grad_out)
+            out = dgl.EData(g, out)
+            sds = out * grad_out  # type dgl.EData
+            sds_sum = sds.dst_sum()  # type dgl.NData
+            grad_score = sds - sds * sds_sum  # multiple expressions
+        """
+        (out,) = self.saved_tensors
+        gidx = self.gidx
+        sds = out * grad_out
+        accum = gspmm(gidx, "copy_rhs", "sum", None, sds)
+        grad_score = sds - gsddmm(gidx, "mul", out, accum, "e", "v")
+        self.save_tensors = None
+        return grad_score
+
+
+def edge_softmax(gidx, logits, eids=ALL, norm_by="dst"):
+    softmax_op = EdgeSoftmax(gidx, eids, norm_by)
+    return softmax_op(logits)
+
+
+class SegmentReduce(mx.autograd.Function):
+    def __init__(self, op, offsets):
+        super(SegmentReduce, self).__init__()
+        self.op = op
+        self.offsets = offsets
+
+    def forward(self, x):
+        y, arg = _segment_reduce(self.op, x, self.offsets)
+        self.save_for_backward(arg)
+        return y
+
+    def backward(self, dy):
+        (arg,) = self.saved_tensors
+        offsets = self.offsets
+        m = offsets[-1].asscalar()
+        if self.op == "sum":
+            offsets_np = asnumpy(offsets[1:])
+            indices_np = np.zeros((m + 1,), dtype=offsets_np.dtype)
+            np.add.at(indices_np, offsets_np, np.ones_like(offsets_np))
+            indices_np = np.cumsum(indices_np, -1)[:-1]
+            indices = zerocopy_from_numpy(indices_np)
+            dx = dy[indices]
+        else:
+            dx = _bwd_segment_cmp(dy, arg, m)
+        return dx
+
+
+def segment_reduce(op, x, offsets):
+    segment_reduce_op = SegmentReduce(op, offsets)
+    return segment_reduce_op(x)
+
+
+class ScatterAdd(mx.autograd.Function):
+    def __init__(self, idx, m):
+        super(ScatterAdd, self).__init__()
+        self.idx = idx
+        self.m = m
+
+    def forward(self, x):
+        y = _scatter_add(x, self.idx, self.m)
+        return y
+
+    def backward(self, dy):
+        return dy[self.idx]
+
+
+def scatter_add(x, idx, m):
+    scatter_add_op = ScatterAdd(idx, m)
+    return scatter_add_op(x)
+
+
+class CSRMM(mx.autograd.Function):
+    def __init__(self, gidxA, gidxB, num_vtypes):
+        super().__init__()
+        self.gidxA = gidxA
+        self.gidxB = gidxB
+        self.num_vtypes = num_vtypes
+
+    def forward(self, A_weights, B_weights):
+        gidxC, C_weights = _csrmm(
+            self.gidxA, A_weights, self.gidxB, B_weights, self.num_vtypes
+        )
+        (
+            nrows,
+            ncols,
+            C_indptr,
+            C_indices,
+            C_eids,
+        ) = gidxC.adjacency_matrix_tensors(0, False, "csr")
+        # Note: the returned C_indptr, C_indices and C_eids tensors MUST be the same
+        # as the underlying tensors of the created graph gidxC.
+        self.backward_cache = gidxC
+        self.save_for_backward(A_weights, B_weights)
+        nrows = nd.array([nrows], dtype="int64")
+        ncols = nd.array([ncols], dtype="int64")
+        return nrows, ncols, C_indptr, C_indices, C_eids, C_weights
+
+    def backward(
+        self, dnrows, dncols, dC_indptr, dC_indices, dC_eids, dC_weights
+    ):
+        # Only the last argument is meaningful.
+        gidxC = self.backward_cache
+        A_weights, B_weights = self.saved_tensors
+        dgidxA, dA_weights = _csrmm(
+            gidxC,
+            dC_weights,
+            self.gidxB.reverse(),
+            B_weights,
+            self.gidxA.number_of_ntypes(),
+        )
+        dgidxB, dB_weights = _csrmm(
+            self.gidxA.reverse(),
+            A_weights,
+            gidxC,
+            dC_weights,
+            self.gidxB.number_of_ntypes(),
+        )
+        dA_weights = _csrmask(dgidxA, dA_weights, self.gidxA)
+        dB_weights = _csrmask(dgidxB, dB_weights, self.gidxB)
+        return dA_weights, dB_weights
+
+
+def csrmm(gidxA, A_weights, gidxB, B_weights, num_vtypes):
+    op = CSRMM(gidxA, gidxB, num_vtypes)
+    nrows, ncols, C_indptr, C_indices, C_eids, C_weights = op(
+        A_weights, B_weights
+    )
+    gidxC = create_unitgraph_from_csr(
+        num_vtypes,
+        nrows.asscalar(),
+        ncols.asscalar(),
+        C_indptr,
+        C_indices,
+        C_eids,
+        ["coo", "csr", "csc"],
+    )
+    return gidxC, C_weights
+
+
+class CSRSum(mx.autograd.Function):
+    def __init__(self, gidxs):
+        super().__init__()
+        self.gidxs = gidxs
+
+    def forward(self, *weights):
+        gidxC, C_weights = _csrsum(self.gidxs, weights)
+        (
+            nrows,
+            ncols,
+            C_indptr,
+            C_indices,
+            C_eids,
+        ) = gidxC.adjacency_matrix_tensors(0, False, "csr")
+        # Note: the returned C_indptr, C_indices and C_eids tensors MUST be the same
+        # as the underlying tensors of the created graph gidxC.
+        self.backward_cache = gidxC
+        nrows = nd.array([nrows], dtype="int64")
+        ncols = nd.array([ncols], dtype="int64")
+        return nrows, ncols, C_indptr, C_indices, C_eids, C_weights
+
+    def backward(
+        self, dnrows, dncols, dC_indptr, dC_indices, dC_eids, dC_weights
+    ):
+        # Only the last argument is meaningful.
+        gidxC = self.backward_cache
+        return tuple(csrmask(gidxC, dC_weights, gidx) for gidx in self.gidxs)
+
+
+def csrsum(gidxs, weights):
+    op = CSRSum(gidxs)
+    nrows, ncols, C_indptr, C_indices, C_eids, C_weights = op(*weights)
+    num_vtypes = gidxs[0].number_of_ntypes()
+    gidxC = create_unitgraph_from_csr(
+        num_vtypes,
+        nrows.asscalar(),
+        ncols.asscalar(),
+        C_indptr,
+        C_indices,
+        C_eids,
+        ["coo", "csr", "csc"],
+    )
+    return gidxC, C_weights
+
+
+class CSRMask(mx.autograd.Function):
+    def __init__(self, gidxA, gidxB):
+        super().__init__()
+        self.gidxA = gidxA
+        self.gidxB = gidxB
+
+    def forward(self, A_weights):
+        return _csrmask(self.gidxA, A_weights, self.gidxB)
+
+    def backward(self, dB_weights):
+        return _csrmask(self.gidxB, dB_weights, self.gidxA)
+
+
+def csrmask(gidxA, A_weights, gidxB):
+    op = CSRMask(gidxA, gidxB)
+    return op(A_weights)

python/dgl/backend/mxnet/sparse_optim.py

@@ -0,0 +1 @@
+"""Sparse optimizer is not supported for mxnet"""

python/dgl/backend/mxnet/tensor.py

@@ -0,0 +1,573 @@
+from __future__ import absolute_import
+
+import builtins
+import numbers
+import os
+
+import mxnet as mx
+import mxnet.ndarray as nd
+import numpy as np
+
+from ... import ndarray as dglnd
+from ...function.base import TargetCode
+from ...utils import version
+
+if version.parse(mx.__version__) < version.parse("1.6.0"):
+    raise RuntimeError("DGL requires MXNet >= 1.6")
+
+# After MXNet 1.5, empty tensors aren't supprted by default.
+# After we turn on the numpy compatible flag, MXNet supports empty NDArray.
+mx.set_np_shape(bool(os.environ.get("DGL_MXNET_SET_NP_SHAPE", True)))
+
+
+def data_type_dict():
+    return {
+        "float16": np.float16,
+        "float32": np.float32,
+        "float64": np.float64,
+        "uint8": np.uint8,
+        "int8": np.int8,
+        "int16": np.int16,
+        "int32": np.int32,
+        "int64": np.int64,
+        "bool": np.bool_,
+    }  # mxnet does not support bool
+
+
+def cpu():
+    return mx.cpu()
+
+
+def tensor(data, dtype=None):
+    if dtype == np.bool_:
+        # mxnet doesn't support bool
+        dtype = np.int32
+    if isinstance(data, nd.NDArray):
+        if dtype is None or data.dtype == dtype:
+            return data
+        else:
+            return data.astype(dtype)
+    else:
+        if isinstance(data, numbers.Number):
+            data = [data]
+        if dtype is None:
+            if isinstance(data, np.ndarray):
+                dtype = np.int32 if data.dtype == np.bool_ else data.dtype
+            elif len(data) == 0:
+                dtype = np.int64
+            else:
+                dtype = (
+                    np.int64
+                    if isinstance(data[0], numbers.Integral)
+                    else np.float32
+                )
+        return nd.array(data, dtype=dtype)
+
+
+def as_scalar(data):
+    if data.size != 1:
+        raise ValueError("The current array is not a scalar")
+    if data.shape != (1,):
+        data = data.expand_dims(axis=0)
+    return data.asscalar()
+
+
+def get_preferred_sparse_format():
+    """Get the preferred sparse matrix format supported by the backend.
+
+    Different backends have their preferred backend. This info is useful when
+    constructing a sparse matrix.
+    """
+    return "csr"
+
+
+def sparse_matrix(data, index, shape, force_format=False):
+    fmt = index[0]
+    if fmt == "coo":
+        if force_format:
+            raise TypeError(
+                "MXNet backend only supports CSR format,"
+                " but COO format is forced."
+            )
+        coord = index[1]
+        # generate convert idx
+        # FIXME: cannot use int64
+        tmp_data = nd.arange(
+            len(coord[0]), dtype=data.dtype, ctx=coord[0].context
+        )
+        tmp_spmat = nd.sparse.csr_matrix(
+            (tmp_data, (coord[0], coord[1])), tuple(shape), ctx=data.context
+        )
+        convert_idx = nd.cast(tmp_spmat.data, dtype="int64")
+        # shuffle the data
+        data = data[convert_idx]
+        spmat = nd.sparse.csr_matrix(
+            (data, tmp_spmat.indices, tmp_spmat.indptr),
+            tuple(shape),
+            ctx=data.context,
+        )
+        return spmat, convert_idx
+    elif fmt == "csr":
+        indices = index[1]
+        indptr = index[2]
+        spmat = nd.sparse.csr_matrix(
+            (data, indices, indptr), tuple(shape), ctx=data.context
+        )
+        # No conversion is required.
+        return spmat, None
+    else:
+        raise TypeError("Invalid format: %s." % fmt)
+
+
+def sparse_matrix_indices(spmat):
+    return ("csr", spmat.indices, spmat.indptr)
+
+
+def is_tensor(obj):
+    return isinstance(obj, nd.NDArray)
+
+
+def shape(input):
+    # NOTE: the input cannot be a symbol
+    return input.shape
+
+
+def dtype(input):
+    # NOTE: the input cannot be a symbol
+    return input.dtype
+
+
+def ndim(input):
+    return input.ndim
+
+
+def context(input):
+    return input.context
+
+
+def device_type(ctx):
+    return ctx.device_type
+
+
+def device_id(ctx):
+    return ctx.device_id
+
+
+def to_backend_ctx(dglctx):
+    dev_type = dglctx.device_type
+    if dev_type == 1:
+        return mx.cpu()
+    elif dev_type == 2:
+        return mx.gpu(dglctx.device_id)
+    else:
+        raise ValueError("Unsupported DGL device context:", dglctx)
+
+
+def astype(input, ty):
+    if ty == np.bool_:
+        ty = np.int32
+    return input.astype(ty)
+
+
+def asnumpy(input):
+    return input.asnumpy()
+
+
+def copy_to(input, ctx, **kwargs):
+    return input.as_in_context(ctx)
+
+
+def is_pinned(input):
+    return input.context == mx.cpu_pinned()
+
+
+def sum(input, dim, keepdims=False):
+    if len(input) == 0:
+        return nd.array([0.0], dtype=input.dtype, ctx=input.context)
+    return nd.sum(input, axis=dim, keepdims=keepdims)
+
+
+def floor_div(in1, in2):
+    return in1 / in2
+
+
+def reduce_sum(input):
+    return input.sum()
+
+
+def cumsum(input, dim):
+    return nd.cumsum(input, axis=dim)
+
+
+def mean(input, dim):
+    return nd.mean(input, axis=dim)
+
+
+def reduce_mean(input):
+    return input.mean()
+
+
+def max(input, dim):
+    return nd.max(input, axis=dim)
+
+
+def reduce_max(input):
+    return input.max()
+
+
+def min(input, dim):
+    return nd.min(input, axis=dim)
+
+
+def reduce_min(input):
+    return input.min()
+
+
+def topk(input, k, dim, descending=True):
+    return nd.topk(
+        input, axis=dim, k=k, ret_typ="value", is_ascend=not descending
+    )
+
+
+def argtopk(input, k, dim, descending=True):
+    idx = nd.argsort(input, dim, is_ascend=not descending)
+    return nd.slice_axis(input, dim, 0, k)
+
+
+def argsort(input, dim, descending):
+    idx = nd.argsort(input, dim, is_ascend=not descending)
+    idx = nd.cast(idx, dtype="int64")
+    return idx
+
+
+def exp(input):
+    return nd.exp(input)
+
+
+def inverse(input):
+    return nd.linalg_inverse(input)
+
+
+def sqrt(input):
+    return nd.sqrt(input)
+
+
+def softmax(input, dim=-1):
+    return nd.softmax(input, axis=dim)
+
+
+def cat(seq, dim):
+    return nd.concat(*seq, dim=dim)
+
+
+def stack(seq, dim):
+    return nd.stack(*seq, axis=dim)
+
+
+def split(x, sizes_or_sections, dim):
+    if isinstance(sizes_or_sections, list) and len(sizes_or_sections) == 1:
+        assert len(x) == sizes_or_sections[0]
+        return [x]
+
+    if isinstance(sizes_or_sections, (np.ndarray, list)):
+        sizes_or_sections1 = tuple(np.cumsum(sizes_or_sections)[:-1])
+    return nd.split_v2(x, sizes_or_sections1, axis=dim)
+
+
+def repeat(input, repeats, dim):
+    if isinstance(repeats, nd.NDArray):
+        return nd.array(
+            np.repeat(input.asnumpy(), repeats.asnumpy(), axis=dim),
+            ctx=input.context,
+            dtype=input.dtype,
+        )
+    else:
+        return nd.repeat(input, repeats, axis=dim)
+
+
+def gather_row(data, row_index):
+    # MXNet workaround for empty row index
+    if len(row_index) == 0:
+        if data.shape[0] == 0:
+            return data
+        else:
+            return data[0:0]
+
+    if isinstance(row_index, nd.NDArray):
+        return nd.take(data, row_index)
+    else:
+        return data[
+            row_index,
+        ]
+
+
+def slice_axis(data, axis, begin, end):
+    dim = data.shape[axis]
+    if begin < 0:
+        begin += dim
+    if end <= 0:
+        end += dim
+    return nd.slice_axis(data, axis, begin, end)
+
+
+def take(data, indices, dim):
+    return nd.take(data, indices, dim)
+
+
+def narrow_row(data, start, stop):
+    return data[start:stop]
+
+
+def index_add_inplace(data, row_idx, value):
+    raise NotImplementedError("MXNet doesn't support inplace index_add")
+
+
+def scatter_row(data, row_index, value):
+    return mx.nd.contrib.index_copy(data, row_index, value)
+
+
+def scatter_row_inplace(data, row_index, value):
+    data[row_index] = value
+
+
+def squeeze(input, dim):
+    return nd.squeeze(input, axis=dim)
+
+
+def unsqueeze(input, dim):
+    return nd.expand_dims(input, axis=dim)
+
+
+def reshape(input, shape):
+    # NOTE: the input cannot be a symbol
+    return nd.reshape(input, shape)
+
+
+def swapaxes(input, axis1, axis2):
+    return nd.swapaxes(input, axis1, axis2)
+
+
+def empty(shape, dtype, ctx):
+    return nd.empty(shape, dtype=dtype, ctx=ctx)
+
+
+def zeros(shape, dtype, ctx):
+    return nd.zeros(shape, dtype=dtype, ctx=ctx)
+
+
+def zeros_like(input):
+    return nd.zeros_like(input)
+
+
+def ones(shape, dtype, ctx):
+    return nd.ones(shape, dtype=dtype, ctx=ctx)
+
+
+def uniform(shape, dtype, ctx, low, high):
+    return nd.random.uniform(low, high, ctx=ctx, dtype=dtype, shape=shape)
+
+
+def randint(shape, dtype, ctx, low, high):
+    return nd.random.randint(low, high, ctx=ctx, dtype=dtype, shape=shape)
+
+
+def pad_packed_tensor(input, lengths, value, l_min=None):
+    old_shape = input.shape
+    if isinstance(lengths, nd.NDArray):
+        lengths = list(lengths.asnumpy())
+    max_len = builtins.max(lengths)
+
+    if l_min is not None:
+        max_len = builtins.max(max_len, l_min)
+
+    batch_size = len(lengths)
+    ctx = input.context
+    dtype = input.dtype
+    x = nd.full(
+        (batch_size * max_len, *old_shape[1:]), value, ctx=ctx, dtype=dtype
+    )
+    index = []
+    for i, l in enumerate(lengths):
+        index.extend(range(i * max_len, i * max_len + l))
+    index = nd.array(index, ctx=ctx)
+    return scatter_row(x, index, input).reshape(
+        batch_size, max_len, *old_shape[1:]
+    )
+
+
+def pack_padded_tensor(input, lengths):
+    batch_size, max_len = input.shape[:2]
+    ctx = input.context
+    index = []
+    for i, l in enumerate(lengths):
+        index.extend(range(i * max_len, i * max_len + l))
+    index = nd.array(index, ctx=ctx)
+    return gather_row(input.reshape(batch_size * max_len, -1), index)
+
+
+def boolean_mask(input, mask):
+    return mx.contrib.nd.boolean_mask(input, mask)
+
+
+def equal(x, y):
+    return x == y
+
+
+def allclose(x, y, rtol=1e-4, atol=1e-4):
+    return np.allclose(x.asnumpy(), y.asnumpy(), rtol=rtol, atol=atol)
+
+
+def logical_not(input):
+    return nd.logical_not(input)
+
+
+def logical_and(input1, input2):
+    return nd.logical_and(input1, input2)
+
+
+def clone(input):
+    return input.copy()
+
+
+def clamp(data, min_val, max_val):
+    return nd.clip(data, min_val, max_val)
+
+
+def replace_inf_with_zero(x):
+    return nd.where(nd.abs(x) == np.inf, nd.zeros_like(x), x)
+
+
+def count_nonzero(input):
+    # TODO: fallback to numpy is unfortunate
+    tmp = input.asnumpy()
+    return np.count_nonzero(tmp)
+
+
+def unique(input, return_inverse=False, return_counts=False):
+    # TODO: fallback to numpy is unfortunate
+    tmp = input.asnumpy()
+    if return_inverse and return_counts:
+        tmp, inv, count = np.unique(
+            tmp, return_inverse=True, return_counts=True
+        )
+        tmp = nd.array(tmp, ctx=input.context, dtype=input.dtype)
+        inv = nd.array(inv, ctx=input.context)
+        count = nd.array(count, ctx=input.context)
+        return tmp, inv, count
+    elif return_inverse or return_counts:
+        tmp, tmp2 = np.unique(
+            tmp, return_inverse=return_inverse, return_counts=return_counts
+        )
+        tmp = nd.array(tmp, ctx=input.context, dtype=input.dtype)
+        tmp2 = nd.array(tmp2, ctx=input.context)
+        return tmp, tmp2
+    else:
+        tmp = np.unique(tmp)
+        return nd.array(tmp, ctx=input.context, dtype=input.dtype)
+
+
+def full_1d(length, fill_value, dtype, ctx):
+    return nd.full((length,), fill_value, dtype=dtype, ctx=ctx)
+
+
+def nonzero_1d(input):
+    # TODO: fallback to numpy is unfortunate
+    tmp = input.asnumpy()
+    tmp = np.nonzero(tmp)[0]
+    r = nd.array(tmp, ctx=input.context, dtype=tmp.dtype)
+    return r
+
+
+def sort_1d(input):
+    # TODO: this isn't an ideal implementation.
+    val = nd.sort(input, axis=None, is_ascend=True)
+    idx = nd.argsort(input, is_ascend=True)
+    idx = nd.cast(idx, dtype="int64")
+    return val, idx
+
+
+def arange(start, stop, dtype=np.int64, ctx=None):
+    if start >= stop:
+        return nd.array([], dtype=dtype, ctx=ctx)
+    else:
+        return nd.arange(start, stop, dtype=dtype, ctx=ctx)
+
+
+def rand_shuffle(arr):
+    return mx.nd.random.shuffle(arr)
+
+
+def zerocopy_to_dlpack(arr):
+    return arr.to_dlpack_for_read()
+
+
+def zerocopy_from_dlpack(dlpack_arr):
+    return nd.from_dlpack(dlpack_arr)
+
+
+def zerocopy_to_numpy(arr):
+    # NOTE: not zerocopy
+    return arr.asnumpy()
+
+
+def zerocopy_from_numpy(np_data):
+    np_data = np.asarray(np_data, order="C")
+    return mx.nd.from_numpy(np_data, zero_copy=True)
+
+
+def zerocopy_to_dgl_ndarray(arr):
+    arr.to_dlpack_for_read()
+    return dglnd.from_dlpack(arr.to_dlpack_for_read())
+
+
+def zerocopy_to_dgl_ndarray_for_write(arr):
+    return dglnd.from_dlpack(arr.to_dlpack_for_write())
+
+
+def zerocopy_from_dgl_ndarray(arr):
+    return nd.from_dlpack(arr.to_dlpack())
+
+
+def sync():
+    """Synchronize computation.
+
+    In DL frameworks such as MXNet and TensorFlow, the computation in operators
+    are done asynchronously. This is to synchronize computation and makes sure
+    that all computation is complete after this function call.
+    """
+    mx.nd.waitall()
+
+
+def attach_grad(tensor):
+    tensor.attach_grad()
+    return tensor
+
+
+def backward(x, head_gradient=None):
+    x.backward(head_gradient)
+
+
+def grad(x):
+    return x.grad
+
+
+def is_no_grad(x):
+    return (x != 0).sum() == 0
+
+
+def is_recording():
+    return mx.autograd.is_recording()
+
+
+record_grad = mx.autograd.record
+
+
+class no_grad(object):
+    def __init__(self):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        pass

python/dgl/backend/pytorch/__init__.py

@@ -0,0 +1,2 @@
+from .sparse import *
+from .tensor import *

python/dgl/backend/pytorch/tensor.py

@@ -0,0 +1,535 @@
+from __future__ import absolute_import
+
+import builtins
+import numbers
+
+import numpy as np
+import scipy  # Weird bug in new pytorch when import scipy after import torch
+import torch as th
+from torch.utils import dlpack
+
+from ... import ndarray as nd
+from ...function.base import TargetCode
+from ...utils import version
+
+if version.parse(th.__version__) < version.parse("2.1.0"):
+    raise RuntimeError("DGL requires PyTorch >= 2.1.0")
+
+
+def data_type_dict():
+    return {
+        "bfloat16": th.bfloat16,
+        "float16": th.float16,
+        "float32": th.float32,
+        "float64": th.float64,
+        "uint8": th.uint8,
+        "int8": th.int8,
+        "int16": th.int16,
+        "int32": th.int32,
+        "int64": th.int64,
+        "bool": th.bool,
+    }
+
+
+def cpu():
+    return th.device("cpu")
+
+
+def tensor(data, dtype=None):
+    if isinstance(data, numbers.Number):
+        data = [data]
+    if (
+        isinstance(data, list)
+        and len(data) > 0
+        and isinstance(data[0], th.Tensor)
+    ):
+        # prevent GPU->CPU->GPU copies
+        if data[0].ndim == 0:
+            # zero dimenion scalar tensors
+            return th.stack(data)
+    if isinstance(data, th.Tensor):
+        return th.as_tensor(data, dtype=dtype, device=data.device)
+    else:
+        return th.as_tensor(data, dtype=dtype)
+
+
+def as_scalar(data):
+    return data.item()
+
+
+def get_preferred_sparse_format():
+    """Get the preferred sparse matrix format supported by the backend.
+
+    Different backends have their preferred backend. This info is useful when
+    constructing a sparse matrix.
+    """
+    return "coo"
+
+
+def sparse_matrix(data, index, shape, force_format=False):
+    fmt = index[0]
+    if fmt != "coo":
+        raise TypeError(
+            "Pytorch backend only supports COO format. But got %s." % fmt
+        )
+    spmat = th.sparse_coo_tensor(index[1], data, shape)
+    return spmat, None
+
+
+def sparse_matrix_indices(spmat):
+    return ("coo", spmat._indices())
+
+
+def is_tensor(obj):
+    return isinstance(obj, th.Tensor)
+
+
+def shape(input):
+    return input.shape
+
+
+def dtype(input):
+    return input.dtype
+
+
+def ndim(input):
+    return input.dim()
+
+
+def context(input):
+    return input.device
+
+
+def device_type(ctx):
+    return th.device(ctx).type
+
+
+def device_id(ctx):
+    ctx = th.device(ctx)
+    if ctx.index is None:
+        return 0 if ctx.type == "cpu" else th.cuda.current_device()
+    else:
+        return ctx.index
+
+
+def to_backend_ctx(dglctx):
+    dev_type = dglctx.device_type
+    if dev_type == 1:
+        return th.device("cpu")
+    elif dev_type == 2:
+        return th.device("cuda", dglctx.device_id)
+    else:
+        raise ValueError("Unsupported DGL device context:", dglctx)
+
+
+def astype(input, ty):
+    return input.type(ty)
+
+
+def asnumpy(input):
+    if isinstance(input, th.sparse.FloatTensor):
+        return input.to_dense().cpu().detach().numpy()
+    else:
+        return input.cpu().detach().numpy()
+
+
+def copy_to(input, ctx, **kwargs):
+    ctx = th.device(ctx)
+    if ctx.type == "cpu":
+        return input.cpu()
+    elif ctx.type == "cuda":
+        if ctx.index is not None:
+            th.cuda.set_device(ctx.index)
+        return input.cuda(**kwargs)
+    else:
+        raise RuntimeError("Invalid context", ctx)
+
+
+def is_pinned(input):
+    return input.is_pinned()
+
+
+def sum(input, dim, keepdims=False):
+    return th.sum(input, dim=dim, keepdim=keepdims)
+
+
+def floor_div(in1, in2):
+    return in1 // in2
+
+
+def reduce_sum(input):
+    return input.sum()
+
+
+def cumsum(input, dim):
+    return th.cumsum(input, dim=dim)
+
+
+def mean(input, dim):
+    return th.mean(input, dim=dim)
+
+
+def reduce_mean(input):
+    return input.mean()
+
+
+def max(input, dim):
+    # NOTE: the second argmax array is not returned
+    return th.max(input, dim=dim)[0]
+
+
+def reduce_max(input):
+    return input.max()
+
+
+def min(input, dim):
+    # NOTE: the second argmin array is not returned
+    return th.min(input, dim=dim)[0]
+
+
+def reduce_min(input):
+    return input.min()
+
+
+def argsort(input, dim, descending):
+    return th.argsort(input, dim=dim, descending=descending)
+
+
+def topk(input, k, dim, descending=True):
+    return th.topk(input, k, dim, largest=descending)[0]
+
+
+def argtopk(input, k, dim, descending=True):
+    return th.topk(input, k, dim, largest=descending)[1]
+
+
+def exp(input):
+    return th.exp(input)
+
+
+def inverse(input):
+    return th.inverse(input)
+
+
+def sqrt(input):
+    return th.sqrt(input)
+
+
+def softmax(input, dim=-1):
+    return th.softmax(input, dim=dim)
+
+
+def cat(seq, dim):
+    return th.cat(seq, dim=dim)
+
+
+def stack(seq, dim):
+    return th.stack(seq, dim=dim)
+
+
+def split(input, sizes_or_sections, dim):
+    return th.split(input, sizes_or_sections, dim)
+
+
+def repeat(input, repeats, dim):
+    return th.repeat_interleave(input, repeats, dim)  # PyTorch 1.1
+
+
+def gather_row(data, row_index):
+    return th.index_select(data, 0, row_index.long())
+
+
+def slice_axis(data, axis, begin, end):
+    return th.narrow(data, axis, begin, end - begin)
+
+
+def take(data, indices, dim):
+    new_shape = data.shape[:dim] + indices.shape + data.shape[dim + 1 :]
+    return th.index_select(data, dim, indices.view(-1)).view(new_shape)
+
+
+def narrow_row(x, start, stop):
+    return x[start:stop]
+
+
+def index_add_inplace(data, row_idx, value):
+    data.index_add_(0, row_idx, value)
+
+
+def scatter_row(data, row_index, value):
+    return data.index_copy(0, row_index.long(), value)
+
+
+def scatter_row_inplace(data, row_index, value):
+    data[row_index.long()] = value
+
+
+def squeeze(input, dim):
+    return th.squeeze(input, dim)
+
+
+def unsqueeze(input, dim):
+    return th.unsqueeze(input, dim)
+
+
+def reshape(input, shape):
+    return th.reshape(input, shape)
+
+
+def swapaxes(input, axis1, axis2):
+    return th.transpose(input, axis1, axis2)
+
+
+def empty(shape, dtype, ctx):
+    return th.empty(shape, dtype=dtype, device=ctx)
+
+
+def zeros(shape, dtype, ctx):
+    return th.zeros(shape, dtype=dtype, device=ctx)
+
+
+def zeros_like(input):
+    return th.zeros_like(input)
+
+
+def ones(shape, dtype, ctx):
+    return th.ones(shape, dtype=dtype, device=ctx)
+
+
+def uniform(shape, dtype, ctx, low, high):
+    return th.empty(shape, dtype=dtype, device=ctx).uniform_(low, high)
+
+
+def randint(shape, dtype, ctx, low, high):
+    return th.randint(low, high, shape, dtype=dtype, device=ctx)
+
+
+def pad_packed_tensor(input, lengths, value, l_min=None):
+    old_shape = input.shape
+    device = input.device
+    if not is_tensor(lengths):
+        lengths = th.tensor(lengths, dtype=th.int64, device=device)
+    else:
+        lengths = lengths.to(device)
+    max_len = as_scalar(lengths.max())
+
+    if l_min is not None:
+        max_len = builtins.max(max_len, l_min)
+
+    batch_size = len(lengths)
+    x = input.new(batch_size * max_len, *old_shape[1:])
+    x.fill_(value)
+    index = th.ones(len(input), dtype=th.int64, device=device)
+    cum_lengths = th.cumsum(lengths, 0)
+    index[cum_lengths[:-1]] += max_len - lengths[:-1]
+    index = th.cumsum(index, 0) - 1
+    x[index] = input
+    return x.view(batch_size, max_len, *old_shape[1:])
+
+
+def pack_padded_tensor(input, lengths):
+    max_len = input.shape[1]
+    device = input.device
+    if not is_tensor(lengths):
+        lengths = th.tensor(lengths, dtype=th.int64, device=device)
+    else:
+        lengths = lengths.to(device)
+    input = input.view(-1, *input.shape[2:])
+    out_len = lengths.sum().item()
+    index = th.ones(out_len, dtype=th.int64, device=device)
+    cum_lengths = th.cumsum(lengths, 0)
+    index[cum_lengths[:-1]] += max_len - lengths[:-1]
+    index = th.cumsum(index, 0) - 1
+    return input[index]
+
+
+def boolean_mask(input, mask):
+    if "bool" not in str(mask.dtype):
+        mask = th.as_tensor(mask, dtype=th.bool)
+    return input[mask]
+
+
+def equal(x, y):
+    return x == y
+
+
+def allclose(x, y, rtol=1e-4, atol=1e-4):
+    return th.allclose(x, y, rtol=rtol, atol=atol)
+
+
+def logical_not(input):
+    return ~input
+
+
+def logical_and(input1, input2):
+    return input1 & input2
+
+
+def clone(input):
+    return input.clone()
+
+
+def clamp(data, min_val, max_val):
+    return th.clamp(data, min_val, max_val)
+
+
+def replace_inf_with_zero(x):
+    return th.masked_fill(x, th.isinf(x), 0)
+
+
+def count_nonzero(input):
+    # TODO: fallback to numpy for backward compatibility
+    return np.count_nonzero(input)
+
+
+def unique(input, return_inverse=False, return_counts=False):
+    if input.dtype == th.bool:
+        input = input.type(th.int8)
+    return th.unique(
+        input, return_inverse=return_inverse, return_counts=return_counts
+    )
+
+
+def full_1d(length, fill_value, dtype, ctx):
+    return th.full((length,), fill_value, dtype=dtype, device=ctx)
+
+
+def nonzero_1d(input):
+    x = th.nonzero(input, as_tuple=False).squeeze()
+    return x if x.dim() == 1 else x.view(-1)
+
+
+def sort_1d(input):
+    return th.sort(input)
+
+
+def arange(start, stop, dtype=th.int64, ctx=None):
+    return th.arange(start, stop, dtype=dtype, device=ctx)
+
+
+def rand_shuffle(arr):
+    idx = th.randperm(len(arr))
+    return arr[idx]
+
+
+def zerocopy_to_dlpack(input):
+    return dlpack.to_dlpack(input.contiguous())
+
+
+def zerocopy_from_dlpack(dlpack_tensor):
+    return dlpack.from_dlpack(dlpack_tensor)
+
+
+def zerocopy_to_numpy(input):
+    # NOTE: not zerocopy
+    return asnumpy(input)
+
+
+def zerocopy_from_numpy(np_array):
+    return th.as_tensor(np_array)
+
+
+def zerocopy_to_dgl_ndarray(data):
+    if data.dtype == th.bool:
+        data = data.byte()
+    return nd.from_dlpack(dlpack.to_dlpack(data.contiguous()))
+
+
+# NGC PyTorch containers are shipping alpha version PyTorch.
+if version.parse(th.__version__) >= version.parse("2.0.0a0"):
+
+    def check_is_view(input):
+        assert (
+            input.data_ptr() == input.untyped_storage().data_ptr()
+        ), "Cannot convert view tensors to dgl ndarray for write."
+
+else:
+
+    def check_is_view(input):
+        assert (
+            input.data_ptr() == input._storage().data_ptr()
+        ), "Cannot convert view tensors to dgl ndarray for write."
+
+
+def zerocopy_to_dgl_ndarray_for_write(input):
+    if input.numel() > 0:
+        # only check non-empty tensors
+        assert input.is_contiguous(), (
+            "Cannot convert non-contiguous tensors "
+            "to dgl ndarray for write. Call .to_contiguous() first."
+        )
+        check_is_view(input)
+    return zerocopy_to_dgl_ndarray(input)
+
+
+def zerocopy_from_dgl_ndarray(data):
+    if data.shape == (0,):
+        # NOTE: PyTorch v1.5 does not accept DLPack object representing empty CUDA tensor.
+        #  Related issue: https://github.com/pytorch/pytorch/issues/41182
+        #  The issue will be fixed in v1.6 and later.
+        return th.tensor(
+            [], dtype=getattr(th, data.dtype), device=to_backend_ctx(data.ctx)
+        )
+    elif len(data.shape) == 0 or builtins.min(data.shape) == 0:
+        # Workaround the same issue as above, but preserve the shape of the
+        # empty tensor. This is needed by the sparse optimizer when one of
+        # processors may receive no gradients to update, but we want to keep
+        # the dimension of the embedding.
+        return th.empty(
+            data.shape,
+            dtype=getattr(th, data.dtype),
+            device=to_backend_ctx(data.ctx),
+        )
+    else:
+        return dlpack.from_dlpack(data.to_dlpack())
+
+
+def sync():
+    # Pytorch performs computation synchronously, so no need for synchronization.
+    pass
+
+
+def attach_grad(x):
+    if x.grad is not None:
+        x.grad.zero_()
+        return x
+    else:
+        return x.requires_grad_()
+
+
+def backward(x, head_gradient=None):
+    if (
+        head_gradient is not None
+        and head_gradient.shape[0] == 1
+        and len(head_gradient.shape) == 1
+    ):
+        # Fix for torch 1.3.1
+        head_gradient = th.tensor(head_gradient.item()).to(head_gradient.device)
+    x.backward(head_gradient)
+
+
+def grad(x):
+    x.retain_grad()
+    return x.grad
+
+
+def is_no_grad(x):
+    return x.grad is None or (x.grad == 0).all()
+
+
+def is_recording():
+    return th.is_grad_enabled()
+
+
+class record_grad(object):
+    def __init__(self):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        pass
+
+
+no_grad = th.no_grad

python/dgl/backend/set_default_backend.py

@@ -0,0 +1,35 @@
+import argparse
+import json
+import os
+
+
+def set_default_backend(default_dir, backend_name):
+    os.makedirs(default_dir, exist_ok=True)
+    config_path = os.path.join(default_dir, "config.json")
+    with open(config_path, "w") as config_file:
+        json.dump({"backend": backend_name.lower()}, config_file)
+    print(
+        'Setting the default backend to "{}". You can change it in the '
+        "~/.dgl/config.json file or export the DGLBACKEND environment variable.  "
+        "Valid options are: pytorch, mxnet, tensorflow (all lowercase)".format(
+            backend_name
+        )
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "default_dir",
+        type=str,
+        default=os.path.join(os.path.expanduser("~"), ".dgl"),
+    )
+    parser.add_argument(
+        "backend",
+        nargs=1,
+        type=str,
+        choices=["pytorch", "tensorflow", "mxnet"],
+        help="Set default backend",
+    )
+    args = parser.parse_args()
+    set_default_backend(args.default_dir, args.backend[0])

python/dgl/backend/tensorflow/__init__.py

@@ -0,0 +1,6 @@
+import os
+
+os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
+
+from .sparse import *
+from .tensor import *

python/dgl/backend/tensorflow/sparse.py

@@ -0,0 +1,461 @@
+import numpy as np
+import tensorflow as tf
+
+from ..._sparse_ops import (
+    _bwd_segment_cmp,
+    _csrmask,
+    _csrmm,
+    _csrsum,
+    _gsddmm,
+    _gspmm,
+    _scatter_add,
+    _segment_reduce,
+)
+
+from ...base import ALL, is_all
+from ...heterograph_index import create_unitgraph_from_csr
+from .tensor import asnumpy, context, copy_to, tensor, zerocopy_from_numpy
+
+__all__ = [
+    "gspmm",
+    "gsddmm",
+    "edge_softmax",
+    "segment_reduce",
+    "scatter_add",
+    "csrmm",
+    "csrsum",
+    "csrmask",
+]
+
+
+def _scatter_nd(index, src, n_rows):
+    assert index.shape == src.shape
+    shp = index.shape
+    ctx = context(src)
+    ndim = index.ndim
+    offsets = []
+    stride = 1
+    for i in reversed(range(1, ndim)):
+        di = shp[i]
+        offset_i = tf.range(di, dtype=index.dtype)
+        offsets.append(
+            tf.reshape(
+                (stride * offset_i), (1,) * i + (di,) + (1,) * (ndim - 1 - i)
+            )
+        )
+        stride *= di
+    if ndim > 1:
+        new_idx = index * stride + copy_to(sum(offsets), ctx)
+    else:
+        new_idx = index
+    src = tf.reshape(src, (-1,))
+    new_idx = tf.reshape(new_idx, (-1, 1))
+    rst = tf.reshape(
+        tf.scatter_nd(new_idx, src, (stride * n_rows,)), (n_rows, *shp[1:])
+    )
+    return rst
+
+
+def _gather_nd(index, src):
+    shp = index.shape
+    ctx = context(src)
+    ndim = index.ndim
+    offsets = []
+    stride = 1
+    for i in reversed(range(1, ndim)):
+        di = shp[i]
+        offset_i = tf.range(di, dtype=index.dtype)
+        offsets.append(
+            tf.reshape(
+                (stride * offset_i), (1,) * i + (di,) + (1,) * (ndim - 1 - i)
+            )
+        )
+        stride *= di
+    if ndim > 1:
+        new_idx = index * stride + copy_to(sum(offsets), ctx)
+    else:
+        new_idx = index
+    src = tf.reshape(src, (-1,))
+    new_idx = tf.reshape(new_idx, (-1))
+    rst = tf.reshape(tf.gather(src, new_idx), shp)
+    return rst
+
+
+def _reduce_grad(grad, shape):
+    """Reduce gradient on the broadcast dimension
+    If there is broadcast in forward pass, gradients need to be reduced on
+    broadcast dimension. This function checks the input tensor shape and
+    gradient shape and perform the reduction.
+    Parameters
+    ----------
+    grad: Tensor
+        Gradient tensor
+    shape: tuple
+        Shape of input tensor
+    Returns
+    -------
+    Tensor
+    """
+    grad_shape = grad.shape[1:]
+    in_shape = shape[1:]
+    if in_shape == grad_shape:
+        # no need to reduce
+        return grad
+    num_to_squeeze = len(grad_shape) - len(in_shape)
+    # pad inshape
+    in_shape = (1,) * num_to_squeeze + in_shape
+    reduce_idx = np.asarray(
+        np.nonzero(np.asarray(grad_shape) - np.asarray(in_shape))
+    )
+    reduce_idx += 1  # skip batch dim
+    reduce_idx_tensor = tf.constant(
+        tuple(reduce_idx.flatten().tolist()), dtype=tf.int32
+    )
+    grad = tf.reduce_sum(grad, axis=reduce_idx_tensor, keepdims=True)
+    return tf.reshape(grad, shape)
+
+
+def _need_reduce_last_dim(ufeat, efeat):
+    """Indicates whether to reduce the last dimension on edges
+    in the backward pass of spmm,
+    if so, use dot instead of mul."""
+    ushp = ufeat.shape
+    eshp = efeat.shape
+    return ushp[1:-1] == eshp[1:-1] and eshp[-1] == 1 and ushp[-1] > 1
+
+
+def _muldiv(op, x):
+    return 1.0 / x if op == "div" else x
+
+
+def _addsub(op, x):
+    return -x if op == "sub" else x
+
+
+def _expand(x, shape):
+    return tf.broadcast_to(x, (x.shape[0], *shape))
+
+
+def gspmm_real(gidx, op, reduce_op, X, Y):
+    out, (argX, argY) = _gspmm(gidx, op, reduce_op, X, Y)
+
+    def grad(dZ):
+        dZ = tensor(dZ)
+        if op != "copy_rhs":
+            g_rev = gidx.reverse()
+            if reduce_op == "sum":
+                if op in ["mul", "div"]:
+                    dX = _gspmm(g_rev, "mul", "sum", dZ, _muldiv(op, Y))[0]
+                elif op in ["add", "sub"]:
+                    dX = _gspmm(g_rev, "copy_lhs", "sum", dZ, Y)[0]
+                elif op == "copy_lhs":
+                    dX = _gspmm(g_rev, "copy_lhs", "sum", dZ, None)[0]
+            else:
+                if op in ["mul", "div"]:
+                    dX = _scatter_nd(
+                        argX,
+                        _muldiv(op, _gather_nd(argY, _expand(Y, dZ.shape[1:])))
+                        * dZ,
+                        X.shape[0],
+                    )
+                elif op in ["add", "sub", "copy_lhs"]:
+                    dX = _scatter_nd(argX, dZ, X.shape[0])
+            dX = _reduce_grad(dX, X.shape)
+        else:
+            dX = tf.zeros_like(X)
+        if op != "copy_lhs":
+            if reduce_op == "sum":
+                if op == "mul" and _need_reduce_last_dim(X, Y):
+                    dY = _gsddmm(gidx, "dot", X, dZ)
+                elif op in ["mul", "div"]:
+                    dY = _gsddmm(gidx, "mul", X, dZ)
+                    if op == "div":
+                        dY = -dY / (Y**2)
+                elif op in ["add", "sub", "copy_rhs"]:
+                    dY = _gsddmm(gidx, "copy_rhs", X, _addsub(op, dZ))
+            else:
+                out_shp = (Y.shape[0],) + dZ.shape[1:]
+                if op in ["mul", "div"]:
+                    dY = _scatter_nd(
+                        argY,
+                        _gather_nd(argX, _expand(X, dZ.shape[1:])) * dZ,
+                        Y.shape[0],
+                    )
+                    if op == "div":
+                        dY = -dY / (Y**2)
+                elif op in ["add", "sub", "copy_rhs"]:
+                    dY = _scatter_nd(argY, _addsub(op, dZ), Y.shape[0])
+            dY = _reduce_grad(dY, Y.shape)
+        else:
+            dY = tf.zeros_like(Y)
+        return dX, dY
+
+    return out, grad
+
+
+def gspmm(gidx, op, reduce_op, X, Y):
+    @tf.custom_gradient
+    def _lambda(X, Y):
+        return gspmm_real(gidx, op, reduce_op, X, Y)
+
+    if X is None:
+        X = tf.zeros(())
+    if Y is None:
+        Y = tf.zeros(())
+    return _lambda(X, Y)
+
+
+def gsddmm_real(gidx, op, X, Y, lhs_target, rhs_target):
+    out = _gsddmm(gidx, op, X, Y, lhs_target, rhs_target)
+
+    def grad(dZ):
+        if op != "copy_rhs":
+            if lhs_target in ["u", "v"]:
+                _gidx = gidx if lhs_target == "v" else gidx.reverse()
+                if op in ["add", "sub", "copy_lhs"]:
+                    dX = _gspmm(_gidx, "copy_rhs", "sum", None, dZ)[0]
+                else:  # mul, div, dot
+                    if rhs_target == lhs_target:
+                        dX = _gspmm(_gidx, "copy_rhs", "sum", None, dZ)[
+                            0
+                        ] * _muldiv(op, Y)
+                    elif rhs_target == "e":
+                        dX = _gspmm(
+                            _gidx, "copy_rhs", "sum", None, dZ * _muldiv(op, Y)
+                        )[0]
+                    else:  # rhs_target = !lhs_target
+                        dX = _gspmm(_gidx, "mul", "sum", _muldiv(op, Y), dZ)[0]
+            else:  # lhs_target == 'e'
+                if op in ["add", "sub", "copy_lhs"]:
+                    dX = dZ
+                else:  # mul, div, dot
+                    dX = _gsddmm(
+                        gidx, "mul", dZ, _muldiv(op, Y), "e", rhs_target
+                    )
+            dX = _reduce_grad(dX, X.shape)
+        else:
+            dX = tf.zeros_like(X)
+        if op != "copy_lhs":
+            if rhs_target in ["u", "v"]:
+                _gidx = gidx if rhs_target == "v" else gidx.reverse()
+                if op in ["add", "sub", "copy_rhs"]:
+                    dY = _gspmm(
+                        _gidx, "copy_rhs", "sum", None, _addsub(op, dZ)
+                    )[0]
+                else:  # mul, div, dot
+                    if lhs_target == rhs_target:
+                        dY = _gspmm(_gidx, "copy_rhs", "sum", None, dZ)[0] * X
+                    elif lhs_target == "e":
+                        dY = _gspmm(_gidx, "copy_rhs", "sum", None, dZ * X)[0]
+                    else:  # rhs_target = !lhs_target
+                        dY = _gspmm(_gidx, "mul", "sum", X, dZ)[0]
+                    if op == "div":
+                        dY = -dY / (Y**2)
+            else:
+                if op in ["add", "sub", "copy_rhs"]:
+                    dY = _addsub(op, dZ)
+                else:  # mul, div, dot
+                    dY = _gsddmm(gidx, "mul", dZ, X, "e", lhs_target)
+                    if op == "div":
+                        dY = -dY / (Y**2)
+            dY = _reduce_grad(dY, Y.shape)
+        else:
+            dY = tf.zeros_like(Y)
+        return dX, dY
+
+    return out, grad
+
+
+def gsddmm(gidx, op, X, Y, lhs_target="u", rhs_target="v"):
+    @tf.custom_gradient
+    def _lambda(X, Y):
+        return gsddmm_real(gidx, op, X, Y, lhs_target, rhs_target)
+
+    if X is None:
+        X = tf.zeros(())
+    if Y is None:
+        Y = tf.zeros(())
+    return _lambda(X, Y)
+
+
+def edge_softmax_real(gidx, score, eids=ALL, norm_by="dst"):
+    if not is_all(eids):
+        gidx = gidx.edge_subgraph([eids], True).graph
+    if norm_by == "src":
+        gidx = gidx.reverse()
+    score_max = _gspmm(gidx, "copy_rhs", "max", None, score)[0]
+    score = tf.math.exp(_gsddmm(gidx, "sub", score, score_max, "e", "v"))
+    score_sum = _gspmm(gidx, "copy_rhs", "sum", None, score)[0]
+    out = _gsddmm(gidx, "div", score, score_sum, "e", "v")
+
+    def edge_softmax_backward(grad_out):
+        sds = out * grad_out
+        accum = gspmm(gidx, "copy_rhs", "sum", None, sds)
+        grad_score = sds - gsddmm(gidx, "mul", out, accum, "e", "v")
+        return grad_score
+
+    return out, edge_softmax_backward
+
+
+def edge_softmax(gidx, logits, eids=ALL, norm_by="dst"):
+    @tf.custom_gradient
+    def _lambda(logits):
+        return edge_softmax_real(gidx, logits, eids, norm_by)
+
+    return _lambda(logits)
+
+
+def segment_reduce_real(op, x, offsets):
+    y, arg = _segment_reduce(op, x, offsets)
+
+    def segment_reduce_backward(dy):
+        m = x.shape[0]
+        if op == "sum":
+            offsets_np = asnumpy(offsets[1:])
+            indices_np = np.zeros((m + 1,), dtype=offsets_np.dtype)
+            np.add.at(indices_np, offsets_np, np.ones_like(offsets_np))
+            indices_np = np.cumsum(indices_np, -1)[:-1]
+            indices = zerocopy_from_numpy(indices_np)
+            dx = tf.gather(dy, indices)
+        else:
+            dx = _bwd_segment_cmp(dy, arg, m)
+        return dx
+
+    return y, segment_reduce_backward
+
+
+def segment_reduce(op, x, offsets):
+    @tf.custom_gradient
+    def _lambda(x):
+        return segment_reduce_real(op, x, offsets)
+
+    return _lambda(x)
+
+
+def scatter_add_real(x, idx, m):
+    y = _scatter_add(x, idx, m)
+
+    def scatter_add_backward(dy):
+        return tf.gather(dy, idx)
+
+    return y, scatter_add_backward
+
+
+def scatter_add(x, idx, m):
+    @tf.custom_gradient
+    def _lambda(x):
+        return scatter_add_real(x, idx, m)
+
+    return _lambda(x)
+
+
+def csrmm_real(gidxA, A_weights, gidxB, B_weights, num_vtypes):
+    gidxC, C_weights = _csrmm(gidxA, A_weights, gidxB, B_weights, num_vtypes)
+    nrows, ncols, C_indptr, C_indices, C_eids = gidxC.adjacency_matrix_tensors(
+        0, False, "csr"
+    )
+
+    def grad(dnrows, dncols, dC_indptr, dC_indices, dC_eids, dC_weights):
+        # Only the last argument is meaningful.
+        dgidxA, dA_weights = _csrmm(
+            gidxC,
+            dC_weights,
+            gidxB.reverse(),
+            B_weights,
+            gidxA.number_of_ntypes(),
+        )
+        dgidxB, dB_weights = _csrmm(
+            gidxA.reverse(),
+            A_weights,
+            gidxC,
+            dC_weights,
+            gidxB.number_of_ntypes(),
+        )
+        dA_weights = _csrmask(dgidxA, dA_weights, gidxA)
+        dB_weights = _csrmask(dgidxB, dB_weights, gidxB)
+        return dA_weights, dB_weights
+
+    return (
+        tf.constant(nrows),
+        tf.constant(ncols),
+        C_indptr,
+        C_indices,
+        C_eids,
+        C_weights,
+    ), grad
+
+
+def csrmm(gidxA, A_weights, gidxB, B_weights, num_vtypes):
+    @tf.custom_gradient
+    def _lambda(A_weights, B_weights):
+        return csrmm_real(gidxA, A_weights, gidxB, B_weights, num_vtypes)
+
+    nrows, ncols, C_indptr, C_indices, C_eids, C_weights = _lambda(
+        A_weights, B_weights
+    )
+    gidxC = create_unitgraph_from_csr(
+        num_vtypes,
+        nrows.numpy(),
+        ncols.numpy(),
+        C_indptr,
+        C_indices,
+        C_eids,
+        ["coo", "csr", "csc"],
+    )
+    return gidxC, C_weights
+
+
+def csrsum_real(gidxs, weights):
+    gidxC, C_weights = _csrsum(gidxs, weights)
+    nrows, ncols, C_indptr, C_indices, C_eids = gidxC.adjacency_matrix_tensors(
+        0, False, "csr"
+    )
+
+    def grad(dnrows, dncols, dC_indptr, dC_indices, dC_eids, dC_weights):
+        # Only the last argument is meaningful.
+        return tuple(_csrmask(gidxC, dC_weights, gidx) for gidx in gidxs)
+
+    return (
+        tf.constant(nrows),
+        tf.constant(ncols),
+        C_indptr,
+        C_indices,
+        C_eids,
+        C_weights,
+    ), grad
+
+
+def csrsum(gidxs, weights):
+    @tf.custom_gradient
+    def _lambda(*weights):
+        return csrsum_real(gidxs, weights)
+
+    nrows, ncols, C_indptr, C_indices, C_eids, C_weights = _lambda(*weights)
+    num_vtypes = gidxs[0].number_of_ntypes()
+    gidxC = create_unitgraph_from_csr(
+        num_vtypes,
+        nrows.numpy(),
+        ncols.numpy(),
+        C_indptr,
+        C_indices,
+        C_eids,
+        ["coo", "csr", "csc"],
+    )
+    return gidxC, C_weights
+
+
+def csrmask_real(gidxA, A_weights, gidxB):
+    B_weights = _csrmask(gidxA, A_weights, gidxB)
+
+    def grad(dB_weights):
+        return _csrmask(gidxB, dB_weights, gidxA)
+
+    return B_weights, grad
+
+
+def csrmask(gidxA, A_weights, gidxB):
+    @tf.custom_gradient
+    def _lambda(A_weights):
+        return csrmask_real(gidxA, A_weights, gidxB)
+
+    return _lambda(A_weights)

python/dgl/backend/tensorflow/sparse_optim.py

@@ -0,0 +1 @@
+"""Sparse optimizer is not supported for tensorflow"""

python/dgl/backend/tensorflow/tensor.py

@@ -0,0 +1,619 @@
+"""Tensorflow backend implementation"""
+from __future__ import absolute_import
+
+import builtins
+import numbers
+
+import numpy as np
+import tensorflow as tf
+
+from ... import ndarray as nd
+from ...function.base import TargetCode
+from ...utils import version
+
+if version.parse(tf.__version__) < version.parse("2.3.0"):
+    raise RuntimeError(
+        "DGL requires TensorFlow>=2.3.0 for the official DLPack support."
+    )
+
+
+def zerocopy_to_dlpack(data):
+    return tf.experimental.dlpack.to_dlpack(data)
+
+
+def zerocopy_from_dlpack(dlpack_tensor):
+    # TODO(Jinjing): Tensorflow requires memory to be 64-bytes aligned. We check the
+    #   alignment and make a copy if needed. The functionality is better in TF's main repo.
+    aligned = nd.from_dlpack(dlpack_tensor).to_dlpack(64)
+    return tf.experimental.dlpack.from_dlpack(aligned)
+
+
+def data_type_dict():
+    return {
+        "bfloat16": tf.bfloat16,
+        "float16": tf.float16,
+        "float32": tf.float32,
+        "float64": tf.float64,
+        "uint8": tf.uint8,
+        "int8": tf.int8,
+        "int16": tf.int16,
+        "int32": tf.int32,
+        "int64": tf.int64,
+        "bool": tf.bool,
+    }
+
+
+def cpu():
+    return "/cpu:0"
+
+
+def tensor(data, dtype=None):
+    if isinstance(data, tf.Tensor):
+        if dtype is None or data.dtype == dtype:
+            return data
+        else:
+            return tf.cast(data, dtype=dtype)
+    else:
+        if isinstance(data, numbers.Number):
+            data = [data]
+        return tf.convert_to_tensor(data, dtype=dtype)
+
+
+def initialize_context():
+    tf.zeros(1)
+
+
+def as_scalar(data):
+    data = data.numpy()
+    return data if np.isscalar(data) else data.item()
+
+
+def get_preferred_sparse_format():
+    """Get the preferred sparse matrix format supported by the backend.
+
+    Different backends have their preferred backend. This info is useful when
+    constructing a sparse matrix.
+    """
+    return "coo"
+
+
+def sparse_matrix(data, index, shape, force_format=False):
+    fmt = index[0]
+    if fmt != "coo":
+        raise TypeError(
+            "Tensorflow backend only supports COO format. But got %s." % fmt
+        )
+    # tf.SparseTensor only supports int64 indexing,
+    # therefore manually casting to int64 when input in int32
+    spmat = tf.SparseTensor(
+        indices=tf.cast(tf.transpose(index[1], (1, 0)), tf.int64),
+        values=data,
+        dense_shape=shape,
+    )
+    return spmat, None
+
+
+def sparse_matrix_indices(spmat):
+    return ("coo", spmat.indices)
+
+
+def is_tensor(obj):
+    return isinstance(obj, tf.Tensor)
+
+
+def shape(input):
+    return input.shape
+
+
+def dtype(input):
+    return input.dtype
+
+
+def ndim(input):
+    return input.ndim
+
+
+def context(input):
+    spec = tf.DeviceSpec.from_string(input.device)
+    return "/{}:{}".format(spec.device_type.lower(), spec.device_index)
+
+
+def device_type(ctx):
+    return tf.DeviceSpec.from_string(ctx).device_type.lower()
+
+
+def device_id(ctx):
+    return tf.DeviceSpec.from_string(ctx).device_index
+
+
+def to_backend_ctx(dglctx):
+    dev_type = dglctx.device_type
+    if dev_type == 1:
+        return "/cpu:0"
+    elif dev_type == 2:
+        return "/gpu:%d" % (dglctx.device_id)
+    else:
+        raise ValueError("Unsupported DGL device context:", dglctx)
+
+
+def astype(input, ty):
+    with tf.device(input.device):
+        return tf.cast(input, dtype=ty)
+
+
+def asnumpy(input):
+    if isinstance(input, tf.SparseTensor):
+        # tf.sparse.to_dense assume sorted indices, need to turn off validate_indices in our cases
+        return tf.sparse.to_dense(input, validate_indices=False).numpy()
+    else:
+        return input.numpy()
+
+
+def copy_to(input, ctx, **kwargs):
+    with tf.device(ctx):
+        new_tensor = tf.identity(input)
+    return new_tensor
+
+
+def is_pinned(input):
+    return False  # not sure how to do this
+
+
+def sum(input, dim, keepdims=False):
+    if input.dtype == tf.bool:
+        input = tf.cast(input, tf.int32)
+    return tf.reduce_sum(input, axis=dim, keepdims=keepdims)
+
+
+def floor_div(in1, in2):
+    return astype(in1 / in2, dtype(in1))
+
+
+def reduce_sum(input):
+    if input.dtype == tf.bool:
+        input = tf.cast(input, tf.int32)
+    return tf.reduce_sum(input)
+
+
+def cumsum(input, dim):
+    if input.dtype == tf.bool:
+        input = tf.cast(input, tf.int32)
+    return tf.cumsum(input, axis=dim)
+
+
+def mean(input, dim):
+    return tf.reduce_mean(input, axis=dim)
+
+
+def reduce_mean(input):
+    return tf.reduce_mean(input)
+
+
+def max(input, dim):
+    return tf.reduce_max(input, axis=dim)
+
+
+def reduce_max(input):
+    return tf.reduce_max(input)
+
+
+def min(input, dim):
+    return tf.reduce_min(input, axis=dim)
+
+
+def reduce_min(input):
+    return tf.reduce_min(input)
+
+
+def argsort(input, dim, descending):
+    if descending:
+        return tf.cast(
+            tf.argsort(input, axis=dim, direction="DESCENDING"), dtype=tf.int64
+        )
+    else:
+        return tf.cast(
+            tf.argsort(input, axis=dim, direction="ASCENDING"), dtype=tf.int64
+        )
+
+
+def topk(input, k, dim, descending=True):
+    if not descending:
+        input = -input
+    shape = np.arange(input.ndim)
+    shape[dim], shape[-1] = shape[-1], shape[dim]
+    out1 = tf.transpose(input, perm=shape)
+    out2 = tf.math.top_k(out1, k=k, sorted=True)
+    out = tf.transpose(out2[0], shape)
+    if not descending:
+        out = -out
+    return out
+
+
+def argtopk(input, k, dim, descending=True):
+    if not descending:
+        input = -input
+    shape = np.arange(input.ndim)
+    shape[dim], shape[-1] = shape[-1], shape[dim]
+    out1 = tf.transpose(input, perm=shape)
+    out2 = tf.math.top_k(out1, k=k, sorted=True)
+    out = tf.transpose(out2[1], shape)
+    if not descending:
+        out = -out
+    return out
+
+
+def exp(input):
+    return tf.exp(input)
+
+
+def inverse(input):
+    return tf.linalg.inv(input)
+
+
+def sqrt(input):
+    return tf.sqrt(input)
+
+
+def softmax(input, dim=-1):
+    return tf.math.softmax(input, axis=dim)
+
+
+def cat(seq, dim):
+    return tf.concat(seq, axis=dim)
+
+
+def stack(seq, dim):
+    return tf.stack(seq, axis=dim)
+
+
+def split(input, sizes_or_sections, dim):
+    return [
+        copy_to(_, input.device)
+        for _ in tf.split(input, sizes_or_sections, axis=dim)
+    ]
+
+
+def repeat(input, repeats, dim):
+    return tf.repeat(input, repeats, dim)
+
+
+def gather_row(data, row_index):
+    return tf.gather(data, row_index)
+
+
+def slice_axis(data, axis, begin, end):
+    # assert axis == 0
+    # tf doesn't behave well with negative
+    s = [slice(None) for i in range(data.ndim)]
+    if end == 0:
+        end = data.shape[axis]
+    s[axis] = slice(begin, end, None)
+    return data[tuple(s)]
+
+
+def take(data, indices, dim):
+    return tf.gather_nd(data, indices, dim)
+
+
+def narrow_row(x, start, stop):
+    return x[start:stop]
+
+
+def scatter_row(data, row_index, value):
+    row_index = tf.expand_dims(row_index, 1)
+    # XXX(minjie): Normally, the copy_to here is unnecessary. However, TF has this
+    #   notorious legacy issue that int32 type data is always on CPU, which will
+    #   crash the program since DGL requires feature data to be on the same device
+    #   as graph structure.
+    return copy_to(
+        tf.tensor_scatter_nd_update(data, row_index, value), data.device
+    )
+
+
+def index_add_inplace(data, row_idx, value):
+    raise NotImplementedError("Tensorflow doesn't support inplace index_add")
+
+
+def scatter_row_inplace(data, row_index, value):
+    raise NotImplementedError("Tensorflow doesn't support inplace update")
+
+
+def squeeze(input, dim):
+    return tf.squeeze(input, axis=dim)
+
+
+def unsqueeze(input, dim):
+    return tf.expand_dims(input, axis=dim)
+
+
+def reshape(input, shape):
+    return tf.reshape(input, shape)
+
+
+def swapaxes(input, axis1, axis2):
+    ndim = input.ndim
+    t = list(range(ndim))
+    t[axis1], t[axis2] = axis2 % ndim, axis1 % ndim
+    return tf.transpose(input, perm=t)
+
+
+def empty(shape, dtype, ctx):
+    # tf doesn't have tf.empty(), use zeros() as a workaround
+    return zeros(shape, dtype, ctx)
+
+
+def zeros(shape, dtype, ctx):
+    with tf.device(ctx):
+        t = tf.zeros(shape, dtype=dtype)
+    return t
+
+
+def zeros_like(input):
+    return tf.zeros_like(input)
+
+
+def ones(shape, dtype, ctx):
+    with tf.device(ctx):
+        t = tf.ones(shape, dtype=dtype)
+    return t
+
+
+def uniform(shape, dtype, ctx, low, high):
+    with tf.device(ctx):
+        t = tf.random.uniform(shape, dtype=dtype, minval=low, maxval=high)
+    return t
+
+
+def randint(shape, dtype, ctx, low, high):
+    with tf.device(ctx):
+        t = tf.random.uniform(shape, dtype=dtype, minval=low, maxval=high)
+    return t
+
+
+def pad_packed_tensor(input, lengths, value, l_min=None):
+    old_shape = input.shape
+    if isinstance(lengths, tf.Tensor):
+        max_len = as_scalar(tf.reduce_max(lengths))
+    else:
+        max_len = builtins.max(lengths)
+
+    if l_min is not None:
+        max_len = builtins.max(max_len, l_min)
+
+    batch_size = len(lengths)
+    ndim = input.ndim
+    tensor_list = []
+    cum_row = 0
+    pad_nparray = np.zeros((ndim, 2), dtype=np.int32)
+    for l in lengths:
+        t = input[cum_row : cum_row + l]
+        pad_nparray[0, 1] = max_len - l
+        t = tf.pad(
+            t, tf.constant(pad_nparray), mode="CONSTANT", constant_values=value
+        )
+        tensor_list.append(t)
+        cum_row += l
+    return tf.stack(tensor_list, axis=0)
+
+
+def pack_padded_tensor(input, lengths):
+    out_list = []
+    for i, l in enumerate(lengths):
+        t = input[i]
+        out = t[:l]
+        out_list.append(out)
+    return tf.concat(out_list, axis=0)
+
+
+def boolean_mask(input, mask):
+    return tf.boolean_mask(input, mask)
+
+
+def equal(x, y):
+    return x == y
+
+
+def allclose(x, y, rtol=1e-4, atol=1e-4):
+    return np.allclose(
+        tf.convert_to_tensor(x).numpy(),
+        tf.convert_to_tensor(y).numpy(),
+        rtol=rtol,
+        atol=atol,
+    )
+
+
+def logical_not(input):
+    return ~input
+
+
+def logical_and(input1, input2):
+    return tf.math.logical_and(input1, input2)
+
+
+def clone(input):
+    # TF tensor is always immutable so returning the input is safe.
+    return input
+
+
+def clamp(data, min_val, max_val):
+    return tf.clip_by_value(data, min_val, max_val)
+
+
+def replace_inf_with_zero(x):
+    return tf.where(tf.abs(x) == np.inf, 0, x)
+
+
+def count_nonzero(input):
+    return int(tf.math.count_nonzero(input))
+
+
+def unique(input, return_inverse=False, return_counts=False):
+    if return_inverse and return_counts:
+        return tf.unique_with_counts(input)
+    elif return_counts:
+        result = tf.unique_with_counts(input)
+        return result.y, result.count
+    elif return_inverse:
+        return tf.unique(input)
+    else:
+        return tf.unique(input).y
+
+
+def full_1d(length, fill_value, dtype, ctx):
+    with tf.device(ctx):
+        t = tf.fill([length], value=fill_value)
+        t = tf.cast(t, dtype=dtype)
+    return t
+
+
+def nonzero_1d(input):
+    nonzero_bool = tf.cast(input, tf.bool)
+    return tf.reshape(tf.where(nonzero_bool), (-1,))
+
+
+def sort_1d(input):
+    return tf.sort(input), tf.cast(tf.argsort(input), dtype=tf.int64)
+
+
+def arange(start, stop, dtype=tf.int64, ctx=None):
+    if not ctx:
+        ctx = "/cpu:0"
+    with tf.device(ctx):
+        t = tf.range(start, stop, dtype=dtype)
+    return t
+
+
+def rand_shuffle(arr):
+    return tf.random.shuffle(arr)
+
+
+def zerocopy_to_numpy(input):
+    return np.asarray(memoryview(input))
+
+
+def zerocopy_from_numpy(np_array):
+    # NOTE: not zerocopy
+    # This assumes tensor should be on cpu
+    with tf.device("/cpu:0"):
+        t = tf.convert_to_tensor(np_array)
+    return t
+
+
+def zerocopy_to_dgl_ndarray(data):
+    if device_type(data.device) == "gpu" and data.dtype in (tf.int32, tf.int64):
+        # NOTE: TF doesn't keep signed tensors on GPU due to legacy issues with
+        #   shape inference. Convert it to unsigned and cast it back afterwards.
+        if data.dtype == tf.int32:
+            data = tf.cast(data, tf.uint32)
+        elif data.dtype == tf.int64:
+            data = tf.cast(data, tf.uint64)
+        return nd.cast_to_signed(nd.from_dlpack(zerocopy_to_dlpack(data)))
+    else:
+        return nd.from_dlpack(zerocopy_to_dlpack(data))
+
+
+def zerocopy_to_dgl_ndarray_for_write(input):
+    return zerocopy_to_dgl_ndarray(input)
+
+
+def zerocopy_from_dgl_ndarray(input):
+    return zerocopy_from_dlpack(input.to_dlpack())
+
+
+def sync():
+    context = context().context()
+    context.async_wait()
+
+
+class GradContext:
+    def __init__(self):
+        self.tensor_for_grad = []
+        self.grad_list = []
+        self.tape = None
+
+    def set_tape(self, tape):
+        self.tape = tape
+
+    def add_tensor(self, x):
+        idx_pop = []
+        for idx, ele in enumerate(self.tensor_for_grad):
+            if ele._id == x._id:
+                idx_pop.append(idx)
+        if len(idx_pop) > 0:
+            self.tensor_for_grad.pop(idx_pop[0])
+        if self.tape is not None:
+            self.tape.watch(x)
+        self.tensor_for_grad.append(x)
+
+    def backward(self, x, head_gradient=None):
+        if head_gradient is not None:
+            x = x * head_gradient
+        self.grad_list = self.tape.gradient(x, self.tensor_for_grad)
+
+    def is_no_grad(self, x):
+        idx_pop = []
+        for idx, ele in enumerate(self.tensor_for_grad):
+            if ele._id == x._id:
+                idx_pop.append(idx)
+        if len(idx_pop) == 0:
+            return True
+        else:
+            return self.grad_list[idx_pop[0]] is None
+
+    def grad(self, x):
+        idx_pop = []
+        for idx, ele in enumerate(self.tensor_for_grad):
+            if ele._id == x._id:
+                idx_pop.append(idx)
+        assert len(idx_pop) == 1
+        t = self.grad_list[idx_pop[0]]
+        return tf.convert_to_tensor(t)
+
+
+cgrad = GradContext()
+
+
+def get_cgrad():
+    return cgrad
+
+
+class record_grad:
+    def __init__(self):
+        self.tape = tf.GradientTape()
+
+    def __enter__(self):
+        cgrad.set_tape(self.tape)
+        self.tape.__enter__()
+        for x in cgrad.tensor_for_grad:
+            self.tape.watch(x)
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        # pass
+        self.tape.__exit__(exc_type, exc_value, exc_traceback)
+        cgrad.tape = None
+
+
+def attach_grad(x):
+    cgrad.add_tensor(x)
+    return x
+
+
+def backward(x, head_gradient=None):
+    cgrad.backward(x, head_gradient)
+
+
+def grad(x):
+    return cgrad.grad(x)
+
+
+def is_no_grad(x):
+    return cgrad.is_no_grad(x)
+
+
+def is_recording():
+    raise NotImplementedError("Tensorflow doesn't support is_recording")
+
+
+no_grad = None
+
+initialize_context()

python/dgl/base.py

@@ -0,0 +1,58 @@
+"""Module for base types and utilities."""
+from __future__ import absolute_import
+
+import warnings
+
+from ._ffi.base import DGLError  # pylint: disable=unused-import
+from ._ffi.function import _init_internal_api
+
+# A special symbol for selecting all nodes or edges.
+ALL = "__ALL__"
+# An alias for [:]
+SLICE_FULL = slice(None, None, None)
+# Reserved column names for storing parent node/edge types and IDs in flattened heterographs
+NTYPE = "_TYPE"
+NID = "_ID"
+ETYPE = "_TYPE"
+EID = "_ID"
+
+_INTERNAL_COLUMNS = {NTYPE, NID, ETYPE, EID}
+
+
+def is_internal_column(name):
+    """Return true if the column name is reversed by DGL."""
+    return name in _INTERNAL_COLUMNS
+
+
+def is_all(arg):
+    """Return true if the argument is a special symbol for all nodes or edges."""
+    return isinstance(arg, str) and arg == ALL
+
+
+# pylint: disable=invalid-name
+_default_formatwarning = warnings.formatwarning
+
+
+class DGLWarning(UserWarning):
+    """DGL Warning class."""
+
+
+# pylint: disable=unused-argument
+def dgl_warning_format(message, category, filename, lineno, line=None):
+    """Format DGL warnings."""
+    if isinstance(category, DGLWarning):
+        return "DGL Warning: {}\n".format(message)
+    else:
+        return _default_formatwarning(
+            message, category, filename, lineno, line=None
+        )
+
+
+def dgl_warning(message, category=DGLWarning, stacklevel=2):
+    """DGL warning wrapper that defaults to ``DGLWarning`` instead of ``UserWarning`` category."""
+    return warnings.warn(message, category=category, stacklevel=stacklevel)
+
+
+warnings.formatwarning = dgl_warning_format
+
+_init_internal_api()

python/dgl/batch.py

@@ -0,0 +1,544 @@
+"""Utilities for batching/unbatching graphs."""
+from collections.abc import Mapping
+
+from . import backend as F, convert, utils
+from .base import ALL, DGLError, EID, is_all, NID
+from .heterograph import DGLGraph
+from .heterograph_index import disjoint_union, slice_gidx
+
+
+__all__ = ["batch", "unbatch", "slice_batch"]
+
+
+def batch(graphs, ndata=ALL, edata=ALL):
+    r"""Batch a collection of :class:`DGLGraph` s into one graph for more efficient
+    graph computation.
+
+    Each input graph becomes one disjoint component of the batched graph. The nodes
+    and edges are relabeled to be disjoint segments:
+
+    =================  =========  =================  ===  =========
+                       graphs[0]  graphs[1]          ...  graphs[k]
+    =================  =========  =================  ===  =========
+    Original node ID   0 ~ N_0    0 ~ N_1            ...  0 ~ N_k
+    New node ID        0 ~ N_0    N_0 ~ N_0+N_1      ...  \sum_{i=0}^{k-1} N_i ~
+                                                          \sum_{i=0}^k N_i
+    =================  =========  =================  ===  =========
+
+    Because of this, many of the computations on a batched graph are the same as if
+    performed on each graph individually, but become much more efficient
+    since they can be parallelized easily. This makes ``dgl.batch`` very useful
+    for tasks dealing with many graph samples such as graph classification tasks.
+
+    For heterograph inputs, they must share the same set of relations (i.e., node types
+    and edge types) and the function will perform batching on each relation one by one.
+    Thus, the result is also a heterograph and has the same set of relations as the inputs.
+
+    The numbers of nodes and edges of the input graphs are accessible via the
+    :func:`DGLGraph.batch_num_nodes` and :func:`DGLGraph.batch_num_edges` attributes
+    of the resulting graph. For homogeneous graphs, they are 1D integer tensors,
+    with each element being the number of nodes/edges of the corresponding input graph. For
+    heterographs, they are dictionaries of 1D integer tensors, with node
+    type or edge type as the keys.
+
+    The function supports batching batched graphs. The batch size of the result
+    graph is the sum of the batch sizes of all the input graphs.
+
+    By default, node/edge features are batched by concatenating the feature tensors
+    of all input graphs. This thus requires features of the same name to have
+    the same data type and feature size. One can pass ``None`` to the ``ndata``
+    or ``edata`` argument to prevent feature batching, or pass a list of strings
+    to specify which features to batch.
+
+    To unbatch the graph back to a list, use the :func:`dgl.unbatch` function.
+
+    Parameters
+    ----------
+    graphs : list[DGLGraph]
+        Input graphs.
+    ndata : list[str], None, optional
+        Node features to batch.
+    edata : list[str], None, optional
+        Edge features to batch.
+
+    Returns
+    -------
+    DGLGraph
+        Batched graph.
+
+    Examples
+    --------
+
+    Batch homogeneous graphs
+
+    >>> import dgl
+    >>> import torch as th
+    >>> # 4 nodes, 3 edges
+    >>> g1 = dgl.graph((th.tensor([0, 1, 2]), th.tensor([1, 2, 3])))
+    >>> # 3 nodes, 4 edges
+    >>> g2 = dgl.graph((th.tensor([0, 0, 0, 1]), th.tensor([0, 1, 2, 0])))
+    >>> bg = dgl.batch([g1, g2])
+    >>> bg
+    Graph(num_nodes=7, num_edges=7,
+          ndata_schemes={}
+          edata_schemes={})
+    >>> bg.batch_size
+    2
+    >>> bg.batch_num_nodes()
+    tensor([4, 3])
+    >>> bg.batch_num_edges()
+    tensor([3, 4])
+    >>> bg.edges()
+    (tensor([0, 1, 2, 4, 4, 4, 5], tensor([1, 2, 3, 4, 5, 6, 4]))
+
+    Batch batched graphs
+
+    >>> bbg = dgl.batch([bg, bg])
+    >>> bbg.batch_size
+    4
+    >>> bbg.batch_num_nodes()
+    tensor([4, 3, 4, 3])
+    >>> bbg.batch_num_edges()
+    tensor([3, 4, 3, 4])
+
+    Batch graphs with feature data
+
+    >>> g1.ndata['x'] = th.zeros(g1.num_nodes(), 3)
+    >>> g1.edata['w'] = th.ones(g1.num_edges(), 2)
+    >>> g2.ndata['x'] = th.ones(g2.num_nodes(), 3)
+    >>> g2.edata['w'] = th.zeros(g2.num_edges(), 2)
+    >>> bg = dgl.batch([g1, g2])
+    >>> bg.ndata['x']
+    tensor([[0, 0, 0],
+            [0, 0, 0],
+            [0, 0, 0],
+            [0, 0, 0],
+            [1, 1, 1],
+            [1, 1, 1],
+            [1, 1, 1]])
+    >>> bg.edata['w']
+    tensor([[1, 1],
+            [1, 1],
+            [1, 1],
+            [0, 0],
+            [0, 0],
+            [0, 0],
+            [0, 0]])
+
+    Batch heterographs
+
+    >>> hg1 = dgl.heterograph({
+    ...     ('user', 'plays', 'game') : (th.tensor([0, 1]), th.tensor([0, 0]))})
+    >>> hg2 = dgl.heterograph({
+    ...     ('user', 'plays', 'game') : (th.tensor([0, 0, 0]), th.tensor([1, 0, 2]))})
+    >>> bhg = dgl.batch([hg1, hg2])
+    >>> bhg
+    Graph(num_nodes={'user': 3, 'game': 4},
+          num_edges={('user', 'plays', 'game'): 5},
+          metagraph=[('drug', 'game')])
+    >>> bhg.batch_size
+    2
+    >>> bhg.batch_num_nodes()
+    {'user' : tensor([2, 1]), 'game' : tensor([1, 3])}
+    >>> bhg.batch_num_edges()
+    {('user', 'plays', 'game') : tensor([2, 3])}
+
+    See Also
+    --------
+    unbatch
+    """
+    if len(graphs) == 0:
+        raise DGLError("The input list of graphs cannot be empty.")
+    if not (is_all(ndata) or isinstance(ndata, list) or ndata is None):
+        raise DGLError(
+            "Invalid argument ndata: must be a string list but got {}.".format(
+                type(ndata)
+            )
+        )
+    if not (is_all(edata) or isinstance(edata, list) or edata is None):
+        raise DGLError(
+            "Invalid argument edata: must be a string list but got {}.".format(
+                type(edata)
+            )
+        )
+    if any(g.is_block for g in graphs):
+        raise DGLError("Batching a MFG is not supported.")
+
+    relations = list(graphs[0].canonical_etypes)
+    relation_ids = [graphs[0].get_etype_id(r) for r in relations]
+    ntypes = list(graphs[0].ntypes)
+    ntype_ids = [graphs[0].get_ntype_id(n) for n in ntypes]
+    etypes = [etype for _, etype, _ in relations]
+
+    gidx = disjoint_union(
+        graphs[0]._graph.metagraph, [g._graph for g in graphs]
+    )
+    retg = DGLGraph(gidx, ntypes, etypes)
+
+    # Compute batch num nodes
+    bnn = {}
+    for ntype in ntypes:
+        bnn[ntype] = F.cat([g.batch_num_nodes(ntype) for g in graphs], 0)
+    retg.set_batch_num_nodes(bnn)
+
+    # Compute batch num edges
+    bne = {}
+    for etype in relations:
+        bne[etype] = F.cat([g.batch_num_edges(etype) for g in graphs], 0)
+    retg.set_batch_num_edges(bne)
+
+    # Batch node feature
+    if ndata is not None:
+        for ntype_id, ntype in zip(ntype_ids, ntypes):
+            all_empty = all(g._graph.num_nodes(ntype_id) == 0 for g in graphs)
+            frames = [
+                g._node_frames[ntype_id]
+                for g in graphs
+                if g._graph.num_nodes(ntype_id) > 0 or all_empty
+            ]
+            # TODO: do we require graphs with no nodes/edges to have the same schema?  Currently
+            # we allow empty graphs to have no features during batching.
+            ret_feat = _batch_feat_dicts(
+                frames, ndata, 'nodes["{}"].data'.format(ntype)
+            )
+            retg.nodes[ntype].data.update(ret_feat)
+
+    # Batch edge feature
+    if edata is not None:
+        for etype_id, etype in zip(relation_ids, relations):
+            all_empty = all(g._graph.num_edges(etype_id) == 0 for g in graphs)
+            frames = [
+                g._edge_frames[etype_id]
+                for g in graphs
+                if g._graph.num_edges(etype_id) > 0 or all_empty
+            ]
+            # TODO: do we require graphs with no nodes/edges to have the same schema?  Currently
+            # we allow empty graphs to have no features during batching.
+            ret_feat = _batch_feat_dicts(
+                frames, edata, "edges[{}].data".format(etype)
+            )
+            retg.edges[etype].data.update(ret_feat)
+
+    return retg
+
+
+def _batch_feat_dicts(frames, keys, feat_dict_name):
+    """Internal function to batch feature dictionaries.
+
+    Parameters
+    ----------
+    frames : list[Frame]
+        List of frames
+    keys : list[str]
+        Feature keys. Can be '__ALL__', meaning batching all features.
+    feat_dict_name : str
+        Name of the feature dictionary for reporting errors.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        New feature dict.
+    """
+    if len(frames) == 0:
+        return {}
+    schemas = [frame.schemes for frame in frames]
+    # sanity checks
+    if is_all(keys):
+        utils.check_all_same_schema(schemas, feat_dict_name)
+        keys = schemas[0].keys()
+    else:
+        utils.check_all_same_schema_for_keys(schemas, keys, feat_dict_name)
+    # concat features
+    ret_feat = {k: F.cat([fd[k] for fd in frames], 0) for k in keys}
+    return ret_feat
+
+
+def unbatch(g, node_split=None, edge_split=None):
+    """Revert the batch operation by split the given graph into a list of small ones.
+
+    This is the reverse operation of :func:``dgl.batch``. If the ``node_split``
+    or the ``edge_split`` is not given, it calls :func:`DGLGraph.batch_num_nodes`
+    and :func:`DGLGraph.batch_num_edges` of the input graph to get the information.
+
+    If the ``node_split`` or the ``edge_split`` arguments are given,
+    it will partition the graph according to the given segments. One must assure
+    that the partition is valid -- edges of the i^th graph only connect nodes
+    belong to the i^th graph. Otherwise, DGL will throw an error.
+
+    The function supports heterograph input, in which case the two split
+    section arguments shall be of dictionary type -- similar to the
+    :func:`DGLGraph.batch_num_nodes`
+    and :func:`DGLGraph.batch_num_edges` attributes of a heterograph.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        Input graph to unbatch.
+    node_split : Tensor, dict[str, Tensor], optional
+        Number of nodes of each result graph.
+    edge_split : Tensor, dict[str, Tensor], optional
+        Number of edges of each result graph.
+
+    Returns
+    -------
+    list[DGLGraph]
+        Unbatched list of graphs.
+
+    Examples
+    --------
+
+    Unbatch a batched graph
+
+    >>> import dgl
+    >>> import torch as th
+    >>> # 4 nodes, 3 edges
+    >>> g1 = dgl.graph((th.tensor([0, 1, 2]), th.tensor([1, 2, 3])))
+    >>> # 3 nodes, 4 edges
+    >>> g2 = dgl.graph((th.tensor([0, 0, 0, 1]), th.tensor([0, 1, 2, 0])))
+    >>> # add features
+    >>> g1.ndata['x'] = th.zeros(g1.num_nodes(), 3)
+    >>> g1.edata['w'] = th.ones(g1.num_edges(), 2)
+    >>> g2.ndata['x'] = th.ones(g2.num_nodes(), 3)
+    >>> g2.edata['w'] = th.zeros(g2.num_edges(), 2)
+    >>> bg = dgl.batch([g1, g2])
+    >>> f1, f2 = dgl.unbatch(bg)
+    >>> f1
+    Graph(num_nodes=4, num_edges=3,
+          ndata_schemes={‘x’ : Scheme(shape=(3,), dtype=torch.float32)}
+          edata_schemes={‘w’ : Scheme(shape=(2,), dtype=torch.float32)})
+    >>> f2
+    Graph(num_nodes=3, num_edges=4,
+          ndata_schemes={‘x’ : Scheme(shape=(3,), dtype=torch.float32)}
+          edata_schemes={‘w’ : Scheme(shape=(2,), dtype=torch.float32)})
+
+    With provided split arguments:
+
+    >>> g1 = dgl.graph((th.tensor([0, 1, 2]), th.tensor([1, 2, 3])))
+    >>> g2 = dgl.graph((th.tensor([0, 0, 0, 1]), th.tensor([0, 1, 2, 0])))
+    >>> g3 = dgl.graph((th.tensor([0]), th.tensor([1])))
+    >>> bg = dgl.batch([g1, g2, g3])
+    >>> bg.batch_num_nodes()
+    tensor([4, 3, 2])
+    >>> bg.batch_num_edges()
+    tensor([3, 4, 1])
+    >>> # unbatch but merge g2 and g3
+    >>> f1, f2 = dgl.unbatch(bg, th.tensor([4, 5]), th.tensor([3, 5]))
+    >>> f1
+    Graph(num_nodes=4, num_edges=3,
+          ndata_schemes={}
+          edata_schemes={})
+    >>> f2
+    Graph(num_nodes=5, num_edges=5,
+          ndata_schemes={}
+          edata_schemes={})
+
+    Heterograph input
+
+    >>> hg1 = dgl.heterograph({
+    ...     ('user', 'plays', 'game') : (th.tensor([0, 1]), th.tensor([0, 0]))})
+    >>> hg2 = dgl.heterograph({
+    ...     ('user', 'plays', 'game') : (th.tensor([0, 0, 0]), th.tensor([1, 0, 2]))})
+    >>> bhg = dgl.batch([hg1, hg2])
+    >>> f1, f2 = dgl.unbatch(bhg)
+    >>> f1
+    Graph(num_nodes={'user': 2, 'game': 1},
+          num_edges={('user', 'plays', 'game'): 2},
+          metagraph=[('drug', 'game')])
+    >>> f2
+    Graph(num_nodes={'user': 1, 'game': 3},
+          num_edges={('user', 'plays', 'game'): 3},
+          metagraph=[('drug', 'game')])
+
+    See Also
+    --------
+    batch
+    """
+    num_split = None
+    # Parse node_split
+    if node_split is None:
+        node_split = {ntype: g.batch_num_nodes(ntype) for ntype in g.ntypes}
+    elif not isinstance(node_split, Mapping):
+        if len(g.ntypes) != 1:
+            raise DGLError(
+                "Must provide a dictionary for argument node_split when"
+                " there are multiple node types."
+            )
+        node_split = {g.ntypes[0]: node_split}
+    if node_split.keys() != set(g.ntypes):
+        raise DGLError("Must specify node_split for each node type.")
+    for split in node_split.values():
+        if num_split is not None and num_split != len(split):
+            raise DGLError(
+                "All node_split and edge_split must specify the same number"
+                " of split sizes."
+            )
+        num_split = len(split)
+
+    # Parse edge_split
+    if edge_split is None:
+        edge_split = {
+            etype: g.batch_num_edges(etype) for etype in g.canonical_etypes
+        }
+    elif not isinstance(edge_split, Mapping):
+        if len(g.etypes) != 1:
+            raise DGLError(
+                "Must provide a dictionary for argument edge_split when"
+                " there are multiple edge types."
+            )
+        edge_split = {g.canonical_etypes[0]: edge_split}
+    if edge_split.keys() != set(g.canonical_etypes):
+        raise DGLError("Must specify edge_split for each canonical edge type.")
+    for split in edge_split.values():
+        if num_split is not None and num_split != len(split):
+            raise DGLError(
+                "All edge_split and edge_split must specify the same number"
+                " of split sizes."
+            )
+        num_split = len(split)
+
+    node_split = {
+        k: F.asnumpy(split).tolist() for k, split in node_split.items()
+    }
+    edge_split = {
+        k: F.asnumpy(split).tolist() for k, split in edge_split.items()
+    }
+
+    # Split edges for each relation
+    edge_dict_per = [{} for i in range(num_split)]
+    for rel in g.canonical_etypes:
+        srctype, etype, dsttype = rel
+        srcnid_off = dstnid_off = 0
+        u, v = g.edges(order="eid", etype=rel)
+        us = F.split(u, edge_split[rel], 0)
+        vs = F.split(v, edge_split[rel], 0)
+        for i, (subu, subv) in enumerate(zip(us, vs)):
+            edge_dict_per[i][rel] = (subu - srcnid_off, subv - dstnid_off)
+            srcnid_off += node_split[srctype][i]
+            dstnid_off += node_split[dsttype][i]
+    num_nodes_dict_per = [
+        {k: split[i] for k, split in node_split.items()}
+        for i in range(num_split)
+    ]
+
+    # Create graphs
+    gs = [
+        convert.heterograph(edge_dict, num_nodes_dict, idtype=g.idtype)
+        for edge_dict, num_nodes_dict in zip(edge_dict_per, num_nodes_dict_per)
+    ]
+
+    # Unbatch node features
+    for ntype in g.ntypes:
+        for key, feat in g.nodes[ntype].data.items():
+            subfeats = F.split(feat, node_split[ntype], 0)
+            for subg, subf in zip(gs, subfeats):
+                subg.nodes[ntype].data[key] = subf
+
+    # Unbatch edge features
+    for etype in g.canonical_etypes:
+        for key, feat in g.edges[etype].data.items():
+            subfeats = F.split(feat, edge_split[etype], 0)
+            for subg, subf in zip(gs, subfeats):
+                subg.edges[etype].data[key] = subf
+
+    return gs
+
+
+def slice_batch(g, gid, store_ids=False):
+    """Get a particular graph from a batch of graphs.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        Input batched graph.
+    gid : int
+        The ID of the graph to retrieve.
+    store_ids : bool
+        If True, it will store the raw IDs of the extracted nodes and edges in the ``ndata`` and
+        ``edata`` of the resulting graph under name ``dgl.NID`` and ``dgl.EID``, respectively.
+
+    Returns
+    -------
+    DGLGraph
+        Retrieved graph.
+
+    Examples
+    --------
+
+    The following example uses PyTorch backend.
+
+    >>> import dgl
+    >>> import torch
+
+    Create a batched graph.
+
+    >>> g1 = dgl.graph(([0, 1], [2, 3]))
+    >>> g2 = dgl.graph(([1], [2]))
+    >>> bg = dgl.batch([g1, g2])
+
+    Get the second component graph.
+
+    >>> g = dgl.slice_batch(bg, 1)
+    >>> print(g)
+    Graph(num_nodes=3, num_edges=1,
+          ndata_schemes={}
+          edata_schemes={})
+    """
+    start_nid = []
+    num_nodes = []
+    for ntype in g.ntypes:
+        batch_num_nodes = g.batch_num_nodes(ntype)
+        num_nodes.append(F.as_scalar(batch_num_nodes[gid]))
+        if gid == 0:
+            start_nid.append(0)
+        else:
+            start_nid.append(
+                F.as_scalar(F.sum(F.slice_axis(batch_num_nodes, 0, 0, gid), 0))
+            )
+
+    start_eid = []
+    num_edges = []
+    for etype in g.canonical_etypes:
+        batch_num_edges = g.batch_num_edges(etype)
+        num_edges.append(F.as_scalar(batch_num_edges[gid]))
+        if gid == 0:
+            start_eid.append(0)
+        else:
+            start_eid.append(
+                F.as_scalar(F.sum(F.slice_axis(batch_num_edges, 0, 0, gid), 0))
+            )
+
+    # Slice graph structure
+    gidx = slice_gidx(
+        g._graph,
+        utils.toindex(num_nodes),
+        utils.toindex(start_nid),
+        utils.toindex(num_edges),
+        utils.toindex(start_eid),
+    )
+    retg = DGLGraph(gidx, g.ntypes, g.etypes)
+
+    # Slice node features
+    for ntid, ntype in enumerate(g.ntypes):
+        stnid = start_nid[ntid]
+        for key, feat in g.nodes[ntype].data.items():
+            subfeats = F.slice_axis(feat, 0, stnid, stnid + num_nodes[ntid])
+            retg.nodes[ntype].data[key] = subfeats
+
+        if store_ids:
+            retg.nodes[ntype].data[NID] = F.arange(
+                stnid, stnid + num_nodes[ntid], retg.idtype, retg.device
+            )
+
+    # Slice edge features
+    for etid, etype in enumerate(g.canonical_etypes):
+        steid = start_eid[etid]
+        for key, feat in g.edges[etype].data.items():
+            subfeats = F.slice_axis(feat, 0, steid, steid + num_edges[etid])
+            retg.edges[etype].data[key] = subfeats
+
+        if store_ids:
+            retg.edges[etype].data[EID] = F.arange(
+                steid, steid + num_edges[etid], retg.idtype, retg.device
+            )
+
+    return retg

python/dgl/container.py

@@ -0,0 +1,102 @@
+"""Container data structures used in DGL runtime.
+reference: tvm/python/tvm/collections.py
+"""
+from __future__ import absolute_import as _abs
+
+from . import _api_internal
+from ._ffi.object import ObjectBase, register_object
+from ._ffi.object_generic import convert_to_object
+
+
+@register_object
+class List(ObjectBase):
+    """List container of DGL.
+
+    You do not need to create List explicitly.
+    Normally python list and tuple will be converted automatically
+    to List during dgl function call.
+    You may get List in return values of DGL function call.
+    """
+
+    def __getitem__(self, i):
+        if isinstance(i, slice):
+            start = i.start if i.start is not None else 0
+            stop = i.stop if i.stop is not None else len(self)
+            step = i.step if i.step is not None else 1
+            if start < 0:
+                start += len(self)
+            if stop < 0:
+                stop += len(self)
+            return [self[idx] for idx in range(start, stop, step)]
+
+        if i < -len(self) or i >= len(self):
+            raise IndexError(
+                "List index out of range. List size: {}, got index {}".format(
+                    len(self), i
+                )
+            )
+        if i < 0:
+            i += len(self)
+        ret = _api_internal._ListGetItem(self, i)
+        if isinstance(ret, Value):
+            ret = ret.data
+        return ret
+
+    def __len__(self):
+        return _api_internal._ListSize(self)
+
+
+@register_object
+class Map(ObjectBase):
+    """Map container of DGL.
+
+    You do not need to create Map explicitly.
+    Normally python dict will be converted automaticall to Map during dgl function call.
+    You can use convert to create a dict[ObjectBase-> ObjectBase] into a Map
+    """
+
+    def __getitem__(self, k):
+        return _api_internal._MapGetItem(self, k)
+
+    def __contains__(self, k):
+        return _api_internal._MapCount(self, k) != 0
+
+    def items(self):
+        """Get the items from the map"""
+        akvs = _api_internal._MapItems(self)
+        return [(akvs[i], akvs[i + 1]) for i in range(0, len(akvs), 2)]
+
+    def __len__(self):
+        return _api_internal._MapSize(self)
+
+
+@register_object
+class StrMap(Map):
+    """A special map container that has str as key.
+
+    You can use convert to create a dict[str->ObjectBase] into a Map.
+    """
+
+    def items(self):
+        """Get the items from the map"""
+        akvs = _api_internal._MapItems(self)
+        return [(akvs[i], akvs[i + 1]) for i in range(0, len(akvs), 2)]
+
+
+@register_object
+class Value(ObjectBase):
+    """Object wrapper for various values."""
+
+    @property
+    def data(self):
+        """Return the value data."""
+        return _api_internal._ValueGet(self)
+
+
+def convert_to_strmap(value):
+    """Convert a python dictionary to a dgl.contrainer.StrMap"""
+    assert isinstance(value, dict), "Only support dict"
+    if len(value) == 0:
+        return _api_internal._EmptyStrMap()
+    else:
+        return convert_to_object(value)

python/dgl/core.py

@@ -0,0 +1,425 @@
+"""Implementation for core graph computation."""
+# pylint: disable=not-callable
+import numpy as np
+
+from . import backend as F, function as fn, ops
+from .base import ALL, dgl_warning, DGLError, EID, is_all, NID
+from .frame import Frame
+from .udf import EdgeBatch, NodeBatch
+
+
+def is_builtin(func):
+    """Return true if the function is a DGL builtin function."""
+    return isinstance(func, fn.BuiltinFunction)
+
+
+def invoke_node_udf(graph, nid, ntype, func, *, ndata=None, orig_nid=None):
+    """Invoke user-defined node function on the given nodes.
+
+    Parameters
+    ----------
+    graph : DGLGraph
+        The input graph.
+    nid : Tensor
+        The IDs of the nodes to invoke UDF on.
+    ntype : str
+        Node type.
+    func : callable
+        The user-defined function.
+    ndata : dict[str, Tensor], optional
+        If provided, apply the UDF on this ndata instead of the ndata of the graph.
+    orig_nid : Tensor, optional
+        Original node IDs. Useful if the input graph is an extracted subgraph.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        Results from running the UDF.
+    """
+    ntid = graph.get_ntype_id(ntype)
+    if ndata is None:
+        if is_all(nid):
+            ndata = graph._node_frames[ntid]
+            nid = graph.nodes(ntype=ntype)
+        else:
+            ndata = graph._node_frames[ntid].subframe(nid)
+    nbatch = NodeBatch(
+        graph, nid if orig_nid is None else orig_nid, ntype, ndata
+    )
+    return func(nbatch)
+
+
+def invoke_edge_udf(graph, eid, etype, func, *, orig_eid=None):
+    """Invoke user-defined edge function on the given edges.
+
+    Parameters
+    ----------
+    graph : DGLGraph
+        The input graph.
+    eid : Tensor
+        The IDs of the edges to invoke UDF on.
+    etype : (str, str, str)
+        Edge type.
+    func : callable
+        The user-defined function.
+    orig_eid : Tensor, optional
+        Original edge IDs. Useful if the input graph is an extracted subgraph.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        Results from running the UDF.
+    """
+    etid = graph.get_etype_id(etype)
+    stid, dtid = graph._graph.metagraph.find_edge(etid)
+    if is_all(eid):
+        u, v, eid = graph.edges(form="all")
+        edata = graph._edge_frames[etid]
+    else:
+        u, v = graph.find_edges(eid)
+        edata = graph._edge_frames[etid].subframe(eid)
+    if len(u) == 0:
+        dgl_warning(
+            "The input graph for the user-defined edge function "
+            "does not contain valid edges"
+        )
+    srcdata = graph._node_frames[stid].subframe(u)
+    dstdata = graph._node_frames[dtid].subframe(v)
+    ebatch = EdgeBatch(
+        graph,
+        eid if orig_eid is None else orig_eid,
+        etype,
+        srcdata,
+        edata,
+        dstdata,
+    )
+    return func(ebatch)
+
+
+def invoke_udf_reduce(graph, func, msgdata, *, orig_nid=None):
+    """Invoke user-defined reduce function on all the nodes in the graph.
+
+    It analyzes the graph, groups nodes by their degrees and applies the UDF on each
+    group -- a strategy called *degree-bucketing*.
+
+    Parameters
+    ----------
+    graph : DGLGraph
+        The input graph.
+    func : callable
+        The user-defined function.
+    msgdata : dict[str, Tensor]
+        Message data.
+    orig_nid : Tensor, optional
+        Original node IDs. Useful if the input graph is an extracted subgraph.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        Results from running the UDF.
+    """
+    degs = graph.in_degrees()
+    nodes = graph.dstnodes()
+    if orig_nid is None:
+        orig_nid = nodes
+    ntype = graph.dsttypes[0]
+    ntid = graph.get_ntype_id_from_dst(ntype)
+    dstdata = graph._node_frames[ntid]
+    msgdata = Frame(msgdata)
+
+    # degree bucketing
+    unique_degs, bucketor = _bucketing(degs)
+    bkt_rsts = []
+    bkt_nodes = []
+    for deg, node_bkt, orig_nid_bkt in zip(
+        unique_degs, bucketor(nodes), bucketor(orig_nid)
+    ):
+        if deg == 0:
+            # skip reduce function for zero-degree nodes
+            continue
+        bkt_nodes.append(node_bkt)
+        ndata_bkt = dstdata.subframe(node_bkt)
+
+        # order the incoming edges per node by edge ID
+        eid_bkt = F.zerocopy_to_numpy(graph.in_edges(node_bkt, form="eid"))
+        assert len(eid_bkt) == deg * len(node_bkt)
+        eid_bkt = np.sort(eid_bkt.reshape((len(node_bkt), deg)), 1)
+        eid_bkt = F.zerocopy_from_numpy(eid_bkt.flatten())
+
+        msgdata_bkt = msgdata.subframe(eid_bkt)
+        # reshape all msg tensors to (num_nodes_bkt, degree, feat_size)
+        maildata = {}
+        for k, msg in msgdata_bkt.items():
+            newshape = (len(node_bkt), deg) + F.shape(msg)[1:]
+            maildata[k] = F.reshape(msg, newshape)
+        # invoke udf
+        nbatch = NodeBatch(graph, orig_nid_bkt, ntype, ndata_bkt, msgs=maildata)
+        bkt_rsts.append(func(nbatch))
+
+    # prepare a result frame
+    retf = Frame(num_rows=len(nodes))
+    retf._initializers = dstdata._initializers
+    retf._default_initializer = dstdata._default_initializer
+
+    # merge bucket results and write to the result frame
+    if (
+        len(bkt_rsts) != 0
+    ):  # if all the nodes have zero degree, no need to merge results.
+        merged_rst = {}
+        for k in bkt_rsts[0].keys():
+            merged_rst[k] = F.cat([rst[k] for rst in bkt_rsts], dim=0)
+        merged_nodes = F.cat(bkt_nodes, dim=0)
+        retf.update_row(merged_nodes, merged_rst)
+
+    return retf
+
+
+def _bucketing(val):
+    """Internal function to create groups on the values.
+
+    Parameters
+    ----------
+    val : Tensor
+        Value tensor.
+
+    Returns
+    -------
+    unique_val : Tensor
+        Unique values.
+    bucketor : callable[Tensor -> list[Tensor]]
+        A bucketing function that splits the given tensor data as the same
+        way of how the :attr:`val` tensor is grouped.
+    """
+    sorted_val, idx = F.sort_1d(val)
+    unique_val = F.asnumpy(F.unique(sorted_val))
+    bkt_idx = []
+    for v in unique_val:
+        eqidx = F.nonzero_1d(F.equal(sorted_val, v))
+        bkt_idx.append(F.gather_row(idx, eqidx))
+
+    def bucketor(data):
+        bkts = [F.gather_row(data, idx) for idx in bkt_idx]
+        return bkts
+
+    return unique_val, bucketor
+
+
+def data_dict_to_list(graph, data_dict, func, target):
+    """Get node or edge feature data of the given name for all the types.
+
+    Parameters
+    -------------
+    graph :  DGLGraph
+        The input graph.
+    data_dict : dict[str, Tensor] or dict[(str, str, str), Tensor]] or Tensor
+        Node or edge data stored in DGLGraph. The key of the dictionary
+        is the node type name or edge type name. If there is only single source
+        node type, data_dict is the value of feature(a Tensor) not a dict.
+    func : dgl.function.BaseMessageFunction
+        Built-in message function.
+    target : 'u', 'v' or 'e'
+        The target of the lhs or rhs data
+
+    Returns
+    --------
+    data_list : list(Tensor)
+        Feature data stored in a list of tensors. The i^th tensor stores the feature
+        data of type ``types[i]``.
+    """
+    if isinstance(func, fn.BinaryMessageFunction):
+        if target in ["u", "v"]:
+            output_list = [None] * graph._graph.number_of_ntypes()
+            # If there is only single source node type, data_dict should be the value of
+            # feature, namely, a tensor.
+            if not isinstance(data_dict, dict):
+                src_id, dst_id = graph._graph.metagraph.find_edge(0)
+                if target == "u":
+                    output_list[src_id] = data_dict
+                else:
+                    output_list[dst_id] = data_dict
+            else:
+                for srctype, _, dsttype in graph.canonical_etypes:
+                    if target == "u":
+                        src_id = graph.get_ntype_id(srctype)
+                        output_list[src_id] = data_dict[srctype]
+                    else:
+                        dst_id = graph.get_ntype_id(dsttype)
+                        output_list[dst_id] = data_dict[dsttype]
+        else:  # target == 'e'
+            output_list = [None] * graph._graph.number_of_etypes()
+            for rel in graph.canonical_etypes:
+                etid = graph.get_etype_id(rel)
+                output_list[etid] = data_dict[rel]
+        return output_list
+    else:
+        if target == "u":
+            lhs_list = [None] * graph._graph.number_of_ntypes()
+            if not isinstance(data_dict, dict):
+                src_id, _ = graph._graph.metagraph.find_edge(0)
+                lhs_list[src_id] = data_dict
+            else:
+                for srctype, _, _ in graph.canonical_etypes:
+                    src_id = graph.get_ntype_id(srctype)
+                    lhs_list[src_id] = data_dict[srctype]
+            return lhs_list
+        else:  # target == 'e':
+            rhs_list = [None] * graph._graph.number_of_etypes()
+            for rel in graph.canonical_etypes:
+                etid = graph.get_etype_id(rel)
+                rhs_list[etid] = data_dict[rel]
+            return rhs_list
+
+
+def invoke_gsddmm(graph, func):
+    """Invoke g-SDDMM computation on the graph.
+
+    Parameters
+    ----------
+    graph :  DGLGraph
+        The input graph.
+    func : dgl.function.BaseMessageFunction
+        Built-in message function.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        Results from the g-SDDMM computation.
+    """
+    alldata = [graph.srcdata, graph.dstdata, graph.edata]
+    if isinstance(func, fn.BinaryMessageFunction):
+        x = alldata[func.lhs][func.lhs_field]
+        y = alldata[func.rhs][func.rhs_field]
+        op = getattr(ops, func.name)
+        if graph._graph.number_of_etypes() > 1:
+            lhs_target, _, rhs_target = func.name.split("_", 2)
+            x = data_dict_to_list(graph, x, func, lhs_target)
+            y = data_dict_to_list(graph, y, func, rhs_target)
+        z = op(graph, x, y)
+    else:
+        x = alldata[func.target][func.in_field]
+        op = getattr(ops, func.name)
+        if graph._graph.number_of_etypes() > 1:
+            # Convert to list as dict is unordered.
+            if func.name == "copy_u":
+                x = data_dict_to_list(graph, x, func, "u")
+            else:  # "copy_e"
+                x = data_dict_to_list(graph, x, func, "e")
+        z = op(graph, x)
+    return {func.out_field: z}
+
+
+def invoke_gspmm(
+    graph, mfunc, rfunc, *, srcdata=None, dstdata=None, edata=None
+):
+    """Invoke g-SPMM computation on the graph.
+
+    Parameters
+    ----------
+    graph :  DGLGraph
+        The input graph.
+    mfunc : dgl.function.BaseMessageFunction
+        Built-in message function.
+    rfunc : dgl.function.BaseReduceFunction
+        Built-in reduce function.
+    srcdata : dict[str, Tensor], optional
+        Source node feature data. If not provided, it use ``graph.srcdata``.
+    dstdata : dict[str, Tensor], optional
+        Destination node feature data. If not provided, it use ``graph.dstdata``.
+    edata : dict[str, Tensor], optional
+        Edge feature data. If not provided, it use ``graph.edata``.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        Results from the g-SPMM computation.
+    """
+    # sanity check
+    if mfunc.out_field != rfunc.msg_field:
+        raise DGLError(
+            "Invalid message ({}) and reduce ({}) function pairs."
+            " The output field of the message function must be equal to the"
+            " message field of the reduce function.".format(mfunc, rfunc)
+        )
+    if edata is None:
+        edata = graph.edata
+    if srcdata is None:
+        srcdata = graph.srcdata
+    if dstdata is None:
+        dstdata = graph.dstdata
+    alldata = [srcdata, dstdata, edata]
+
+    if isinstance(mfunc, fn.BinaryMessageFunction):
+        x = alldata[mfunc.lhs][mfunc.lhs_field]
+        y = alldata[mfunc.rhs][mfunc.rhs_field]
+        op = getattr(ops, "{}_{}".format(mfunc.name, rfunc.name))
+        if graph._graph.number_of_etypes() > 1:
+            lhs_target, _, rhs_target = mfunc.name.split("_", 2)
+            x = data_dict_to_list(graph, x, mfunc, lhs_target)
+            y = data_dict_to_list(graph, y, mfunc, rhs_target)
+        z = op(graph, x, y)
+    else:
+        x = alldata[mfunc.target][mfunc.in_field]
+        op = getattr(ops, "{}_{}".format(mfunc.name, rfunc.name))
+        if graph._graph.number_of_etypes() > 1 and not isinstance(x, tuple):
+            if mfunc.name == "copy_u":
+                x = data_dict_to_list(graph, x, mfunc, "u")
+            else:  # "copy_e"
+                x = data_dict_to_list(graph, x, mfunc, "e")
+        z = op(graph, x)
+    return {rfunc.out_field: z}
+
+
+def message_passing(g, mfunc, rfunc, afunc):
+    """Invoke message passing computation on the whole graph.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The input graph.
+    mfunc : callable or dgl.function.BuiltinFunction
+        Message function.
+    rfunc : callable or dgl.function.BuiltinFunction
+        Reduce function.
+    afunc : callable or dgl.function.BuiltinFunction
+        Apply function.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        Results from the message passing computation.
+    """
+    if (
+        is_builtin(mfunc)
+        and is_builtin(rfunc)
+        and getattr(ops, "{}_{}".format(mfunc.name, rfunc.name), None)
+        is not None
+    ):
+        # invoke fused message passing
+        ndata = invoke_gspmm(g, mfunc, rfunc)
+    else:
+        # invoke message passing in two separate steps
+        # message phase
+        if is_builtin(mfunc):
+            msgdata = invoke_gsddmm(g, mfunc)
+        else:
+            orig_eid = g.edata.get(EID, None)
+            msgdata = invoke_edge_udf(
+                g, ALL, g.canonical_etypes[0], mfunc, orig_eid=orig_eid
+            )
+        # reduce phase
+        if is_builtin(rfunc):
+            msg = rfunc.msg_field
+            ndata = invoke_gspmm(g, fn.copy_e(msg, msg), rfunc, edata=msgdata)
+        else:
+            orig_nid = g.dstdata.get(NID, None)
+            ndata = invoke_udf_reduce(g, rfunc, msgdata, orig_nid=orig_nid)
+    # apply phase
+    if afunc is not None:
+        for k, v in g.dstdata.items():  # include original node features
+            if k not in ndata:
+                ndata[k] = v
+        orig_nid = g.dstdata.get(NID, None)
+        ndata = invoke_node_udf(
+            g, ALL, g.dsttypes[0], afunc, ndata=ndata, orig_nid=orig_nid
+        )
+    return ndata

python/dgl/cuda/__init__.py

@@ -0,0 +1,7 @@
+""" CUDA wrappers """
+from .. import backend as F
+
+from .gpu_cache import GPUCache
+
+if F.get_preferred_backend() == "pytorch":
+    from . import nccl

python/dgl/cuda/gpu_cache.py

@@ -0,0 +1,86 @@
+"""API wrapping HugeCTR gpu_cache."""
+#    Copyright (c) 2022, NVIDIA Corporation
+#    All rights reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+#  @file gpu_cache.py
+#  @brief API for managing a GPU Cache
+
+from .. import backend as F
+from .._ffi.function import _init_api
+
+
+class GPUCache(object):
+    """High-level wrapper for GPU embedding cache"""
+
+    def __init__(self, num_items, num_feats, idtype=F.int64):
+        assert idtype in [F.int32, F.int64]
+        self._cache = _CAPI_DGLGpuCacheCreate(
+            num_items, num_feats, 32 if idtype == F.int32 else 64
+        )
+        self.idtype = idtype
+        self.total_miss = 0
+        self.total_queries = 0
+
+    def query(self, keys):
+        """Queries the GPU cache.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the GPU cache with.
+
+        Returns
+        -------
+        tuple(Tensor, Tensor, Tensor)
+            A tuple containing (values, missing_indices, missing_keys) where
+            values[missing_indices] corresponds to cache misses that should be
+            filled by quering another source with missing_keys.
+        """
+        self.total_queries += keys.shape[0]
+        keys = F.astype(keys, self.idtype)
+        values, missing_index, missing_keys = _CAPI_DGLGpuCacheQuery(
+            self._cache, F.to_dgl_nd(keys)
+        )
+        self.total_miss += missing_keys.shape[0]
+        return (
+            F.from_dgl_nd(values),
+            F.from_dgl_nd(missing_index),
+            F.from_dgl_nd(missing_keys),
+        )
+
+    def replace(self, keys, values):
+        """Inserts key-value pairs into the GPU cache using the Least-Recently
+        Used (LRU) algorithm to remove old key-value pairs if it is full.
+
+        Parameters
+        ----------
+        keys: Tensor
+            The keys to insert to the GPU cache.
+        values: Tensor
+            The values to insert to the GPU cache.
+        """
+        keys = F.astype(keys, self.idtype)
+        values = F.astype(values, F.float32)
+        _CAPI_DGLGpuCacheReplace(
+            self._cache, F.to_dgl_nd(keys), F.to_dgl_nd(values)
+        )
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self.total_miss / self.total_queries
+
+
+_init_api("dgl.cuda", __name__)

python/dgl/cuda/nccl.py

@@ -0,0 +1,189 @@
+"""API wrapping NCCL primitives."""
+
+import torch
+import torch.distributed as dist
+
+
+def sparse_all_to_all_push(idx, value, partition):
+    """Perform an all-to-all-v operation, where by all processors send out
+    a set of indices and corresponding values. Indices and values,
+    corresponding to the current process, will copied into the output
+    arrays.
+
+    Note: This method requires 'torch.distributed.get_backend() == "nccl"'.
+
+    Parameters
+    ----------
+    idx : torch.Tensor
+        The 1D set of indices to send to other processors.
+    value : torch.Tensor
+        The multi-dimension set of values to send to other processors.
+        The first dimension must match that of `idx`.
+    partition : NDArrayPartition
+        The object containing information for assigning indices to
+        processors.
+
+    Returns
+    -------
+    torch.Tensor
+        The 1D tensor of the recieved indices.
+    torch.Tensor
+        The set of recieved values.
+
+    Examples
+    --------
+
+    To perform a sparse_all_to_all_push(), a partition object must be
+    provided. A partition of a homgeonous graph, where the vertices are
+    striped across processes can be generated via:
+
+    >>> from dgl.partition import NDArrayPartition
+    >>> part = NDArrayPartition(g.num_nodes(), world_size, mode='remainder')
+
+    With this partition, each processor can send values to be associatd
+    with vertices in the graph. So if we have an array `global_idxs` of all of
+    the neighbors updated during mini-batch processing, and an array
+    `global_values` containing the new values associated with the neighbors,
+    we communicate them to the own processes via:
+
+    >>> my_idxs, my_values = nccl.sparse_all_to_all_push(global_idxs, global_values, part)
+
+    This communication pattern is common when communicating gradient
+    updates for node embeddings.
+
+    Indices the current process owns, do not need to treated specially,
+    as internally they will be copied to the output array. If we have a
+    set of indices in process 0 '[0, 3, 8, 9, 10]` and for process 1
+    '[0, 2, 4, 5, 8, 8, 9]'. Using a remainder partition will result
+    indices for processe 0 of '[0, 8, 10, 0, 2, 4, 8, 8]', and for
+    process 1 of '[3, 9, 5, 9]'.
+    """
+    if not dist.is_initialized() or dist.get_world_size() == 1:
+        return idx, value
+    assert (
+        dist.get_backend() == "nccl"
+    ), "requires NCCL backend to communicate CUDA tensors."
+
+    perm, send_splits = partition.generate_permutation(idx)
+    perm = perm.long()
+
+    # Get receive splits.
+    recv_splits = torch.empty_like(send_splits)
+    dist.all_to_all_single(recv_splits, send_splits)
+
+    # Use pinned memory to speedup D2H copy.
+    recv_splits = recv_splits.to("cpu", non_blocking=True)
+    send_splits = send_splits.to("cpu", non_blocking=True)
+    send_idx = idx[perm]
+    send_value = value[perm]
+    # Wait D2H copy finish.
+    torch.cuda.current_stream().synchronize()
+    recv_sum = recv_splits.sum()
+    recv_splits = recv_splits.tolist()
+    send_splits = send_splits.tolist()
+
+    # Send idx.
+    recv_idx = torch.empty((recv_sum,), dtype=idx.dtype, device=idx.device)
+    dist.all_to_all_single(recv_idx, send_idx, recv_splits, send_splits)
+
+    # Send value.
+    recv_value = torch.empty(
+        (recv_sum, *value.shape[1:]), dtype=value.dtype, device=value.device
+    )
+    dist.all_to_all_single(recv_value, send_value, recv_splits, send_splits)
+
+    return recv_idx, recv_value
+
+
+def sparse_all_to_all_pull(req_idx, value, partition):
+    """Perform an all-to-all-v operation, where by all processors request
+    the values corresponding to their set of indices.
+
+    Note: This method requires 'torch.distributed.get_backend() == "nccl"'.
+
+    Parameters
+    ----------
+    req_idx : torch.Tensor
+        The set of indices this processor is requesting.
+    value : torch.Tensor
+        The multi-dimension set of values that can be requested from
+        this processor.
+    partition : NDArrayPartition
+        The object containing information for assigning indices to
+        processors.
+
+    Returns
+    -------
+    torch.Tensor
+        The set of recieved values, corresponding to `req_idx`.
+
+    Examples
+    --------
+
+    To perform a sparse_all_to_all_pull(), a partition object must be
+    provided. A partition of a homgeonous graph, where the vertices are
+    striped across processes can be generated via:
+
+    >>> from dgl.partition import NDArrayPartition
+    >>> part = NDArrayPartition(g.num_nodes(), world_size, mode='remainder')
+
+    With this partition, each processor can request values/features
+    associated with vertices in the graph. So in the case where we have
+    a set of neighbors 'nbr_idxs' we need features for, and each process
+    has a tensor 'node_feat' storing the features of nodes it owns in
+    the partition, the features can be requested via:
+
+    >>> nbr_values = nccl.sparse_all_to_all_pull(nbr_idxs, node_feat, part)
+
+    Then two the arrays 'nbr_idxs' and 'nbr_values' forms the sparse
+    set of features, where 'nbr_idxs[i]' is the global node id, and
+    'nbr_values[i]' is the feature vector for that node. This
+    communication pattern is useful for node features or node
+    embeddings.
+    """
+    if not dist.is_initialized() or dist.get_world_size() == 1:
+        return value[req_idx.long()]
+    assert (
+        dist.get_backend() == "nccl"
+    ), "requires NCCL backend to communicate CUDA tensors."
+
+    perm, req_splits = partition.generate_permutation(req_idx)
+    perm = perm.long()
+
+    # Get response splits.
+    resp_splits = torch.empty_like(req_splits)
+    dist.all_to_all_single(resp_splits, req_splits)
+
+    # Use pinned memory to speedup D2H copy.
+    resp_splits = resp_splits.to("cpu", non_blocking=True)
+    req_splits = req_splits.to("cpu", non_blocking=True)
+    req_idx = req_idx[perm]
+    # Wait D2H copy finish.
+    torch.cuda.current_stream().synchronize()
+    resp_sum = resp_splits.sum()
+    resp_splits = resp_splits.tolist()
+    req_splits = req_splits.tolist()
+
+    # Gather requested indices.
+    resp_idx = torch.empty(
+        (resp_sum,), dtype=req_idx.dtype, device=req_idx.device
+    )
+    dist.all_to_all_single(resp_idx, req_idx, resp_splits, req_splits)
+
+    # Convert requested indices to local indices depending on partition.
+    if resp_sum > 0:
+        resp_idx = partition.map_to_local(resp_idx)
+
+    # Collect the request value.
+    req_value = torch.empty(
+        (req_idx.size(0), *value.shape[1:]),
+        dtype=value.dtype,
+        device=value.device,
+    )
+    dist.all_to_all_single(req_value, value[resp_idx], req_splits, resp_splits)
+
+    # Permute the value back into the requested order.
+    return_value = torch.empty_like(req_value)
+    return_value[perm] = req_value
+
+    return return_value

python/dgl/data/__init__.py

@@ -0,0 +1,112 @@
+"""The ``dgl.data`` package contains datasets hosted by DGL and also utilities
+for downloading, processing, saving and loading data from external resources.
+"""
+
+from __future__ import absolute_import
+
+from . import citation_graph as citegrh
+from .actor import ActorDataset
+from .movielens import MovieLensDataset
+from .adapter import *
+from .bitcoinotc import BitcoinOTC, BitcoinOTCDataset
+from .citation_graph import (
+    CitationGraphDataset,
+    CiteseerGraphDataset,
+    CoraBinary,
+    CoraGraphDataset,
+    PubmedGraphDataset,
+)
+from .csv_dataset import CSVDataset
+from .dgl_dataset import DGLBuiltinDataset, DGLDataset
+from .fakenews import FakeNewsDataset
+from .flickr import FlickrDataset
+from .fraud import FraudAmazonDataset, FraudDataset, FraudYelpDataset
+from .gdelt import GDELT, GDELTDataset
+from .gindt import GINDataset
+from .gnn_benchmark import (
+    AmazonCoBuy,
+    AmazonCoBuyComputerDataset,
+    AmazonCoBuyPhotoDataset,
+    Coauthor,
+    CoauthorCSDataset,
+    CoauthorPhysicsDataset,
+    CoraFull,
+    CoraFullDataset,
+)
+from .icews18 import ICEWS18, ICEWS18Dataset
+from .karate import KarateClub, KarateClubDataset
+from .knowledge_graph import FB15k237Dataset, FB15kDataset, WN18Dataset
+from .minigc import *
+from .ppi import LegacyPPIDataset, PPIDataset
+from .qm7b import QM7b, QM7bDataset
+from .qm9 import QM9, QM9Dataset
+from .qm9_edge import QM9Edge, QM9EdgeDataset
+from .rdf import AIFBDataset, AMDataset, BGSDataset, MUTAGDataset
+from .reddit import RedditDataset
+from .sbm import SBMMixture, SBMMixtureDataset
+from .synthetic import (
+    BA2MotifDataset,
+    BACommunityDataset,
+    BAShapeDataset,
+    TreeCycleDataset,
+    TreeGridDataset,
+)
+from .tree import SST, SSTDataset
+from .tu import LegacyTUDataset, TUDataset
+from .utils import *
+from .cluster import CLUSTERDataset
+from .geom_gcn import (
+    ChameleonDataset,
+    CornellDataset,
+    SquirrelDataset,
+    TexasDataset,
+    WisconsinDataset,
+)
+
+from .heterophilous_graphs import (
+    AmazonRatingsDataset,
+    MinesweeperDataset,
+    QuestionsDataset,
+    RomanEmpireDataset,
+    TolokersDataset,
+)
+
+# RDKit is required for Peptides-Structural, Peptides-Functional dataset.
+# Exception handling was added to prevent crashes for users who are using other
+# datasets.
+try:
+    from .lrgb import (
+        COCOSuperpixelsDataset,
+        PeptidesFunctionalDataset,
+        PeptidesStructuralDataset,
+        VOCSuperpixelsDataset,
+    )
+except ImportError:
+    pass
+from .pattern import PATTERNDataset
+from .superpixel import CIFAR10SuperPixelDataset, MNISTSuperPixelDataset
+from .wikics import WikiCSDataset
+from .yelp import YelpDataset
+from .zinc import ZINCDataset
+
+
+def register_data_args(parser):
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        required=False,
+        help="The input dataset. Can be cora, citeseer, pubmed, syn(synthetic dataset) or reddit",
+    )
+
+
+def load_data(args):
+    if args.dataset == "cora":
+        return citegrh.load_cora()
+    elif args.dataset == "citeseer":
+        return citegrh.load_citeseer()
+    elif args.dataset == "pubmed":
+        return citegrh.load_pubmed()
+    elif args.dataset is not None and args.dataset.startswith("reddit"):
+        return RedditDataset(self_loop=("self-loop" in args.dataset))
+    else:
+        raise ValueError("Unknown dataset: {}".format(args.dataset))

python/dgl/data/actor.py

@@ -0,0 +1,138 @@
+"""
+Actor-only induced subgraph of the film-directoractor-writer network.
+"""
+import os
+
+import numpy as np
+
+from ..convert import graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url
+
+
+class ActorDataset(DGLBuiltinDataset):
+    r"""Actor-only induced subgraph of the film-directoractor-writer network
+    from `Social Influence Analysis in Large-scale Networks
+    <https://dl.acm.org/doi/10.1145/1557019.1557108>`, introduced by
+    `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`
+
+    Nodes represent actors, and edges represent co-occurrence on the same
+    Wikipedia page. Node features correspond to some keywords in the Wikipedia
+    pages.
+
+    Statistics:
+
+    - Nodes: 7600
+    - Edges: 33391
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 3648
+        - Val: 2432
+        - Test: 1520
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(ActorDataset, self).__init__(
+            name="actor",
+            url=_get_dgl_url("dataset/actor.zip"),
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """Load and process the data."""
+        try:
+            import torch
+        except ImportError:
+            raise ModuleNotFoundError(
+                "This dataset requires PyTorch to be the backend."
+            )
+
+        # Process node features and labels.
+        with open(f"{self.raw_path}/out1_node_feature_label.txt", "r") as f:
+            data = [x.split("\t") for x in f.read().split("\n")[1:-1]]
+
+            rows, cols = [], []
+            labels = torch.empty(len(data), dtype=torch.long)
+            for n_id, col, label in data:
+                col = [int(x) for x in col.split(",")]
+                rows += [int(n_id)] * len(col)
+                cols += col
+
+                labels[int(n_id)] = int(label)
+
+            row, col = torch.tensor(rows), torch.tensor(cols)
+            features = torch.zeros(len(data), int(col.max()) + 1)
+            features[row, col] = 1.0
+
+            self._num_classes = int(labels.max().item()) + 1
+
+        # Process graph structure.
+        with open(f"{self.raw_path}/out1_graph_edges.txt", "r") as f:
+            data = f.read().split("\n")[1:-1]
+            data = [[int(v) for v in r.split("\t")] for r in data]
+        dst, src = torch.tensor(data, dtype=torch.long).t().contiguous()
+
+        self._g = graph((src, dst), num_nodes=features.size(0))
+        self._g.ndata["feat"] = features
+        self._g.ndata["label"] = labels
+
+        # Process 10 train/val/test node splits.
+        train_masks, val_masks, test_masks = [], [], []
+        for i in range(10):
+            filepath = f"{self.raw_path}/{self.name}_split_0.6_0.2_{i}.npz"
+            f = np.load(filepath)
+            train_masks += [torch.from_numpy(f["train_mask"])]
+            val_masks += [torch.from_numpy(f["val_mask"])]
+            test_masks += [torch.from_numpy(f["test_mask"])]
+        self._g.ndata["train_mask"] = torch.stack(train_masks, dim=1).bool()
+        self._g.ndata["val_mask"] = torch.stack(val_masks, dim=1).bool()
+        self._g.ndata["test_mask"] = torch.stack(test_masks, dim=1).bool()
+
+    def has_cache(self):
+        return os.path.exists(self.raw_path)
+
+    def load(self):
+        self.process()
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return self._num_classes

python/dgl/data/adapter.py

@@ -0,0 +1,633 @@
+"""Dataset adapters for re-purposing a dataset for a different kind of training task."""
+
+import json
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..base import DGLError
+from ..convert import graph as create_dgl_graph
+from ..sampling.negative import _calc_redundancy
+from . import utils
+from .dgl_dataset import DGLDataset
+
+__all__ = ["AsNodePredDataset", "AsLinkPredDataset", "AsGraphPredDataset"]
+
+
+class AsNodePredDataset(DGLDataset):
+    """Repurpose a dataset for a standard semi-supervised transductive
+    node prediction task.
+
+    The class converts a given dataset into a new dataset object such that:
+
+      - Contains only one graph, accessible from ``dataset[0]``.
+      - The graph stores:
+
+        - Node labels in ``g.ndata['label']``.
+        - Train/val/test masks in ``g.ndata['train_mask']``, ``g.ndata['val_mask']``,
+          and ``g.ndata['test_mask']`` respectively.
+      - In addition, the dataset contains the following attributes:
+
+        - ``num_classes``, the number of classes to predict.
+        - ``train_idx``, ``val_idx``, ``test_idx``, train/val/test indexes.
+
+    If the input dataset contains heterogeneous graphs, users need to specify the
+    ``target_ntype`` argument to indicate which node type to make predictions for.
+    In this case:
+
+      - Node labels are stored in ``g.nodes[target_ntype].data['label']``.
+      - Training masks are stored in ``g.nodes[target_ntype].data['train_mask']``.
+        So do validation and test masks.
+
+    The class will keep only the first graph in the provided dataset and
+    generate train/val/test masks according to the given split ratio. The generated
+    masks will be cached to disk for fast re-loading. If the provided split ratio
+    differs from the cached one, it will re-process the dataset properly.
+
+    Parameters
+    ----------
+    dataset : DGLDataset
+        The dataset to be converted.
+    split_ratio : (float, float, float), optional
+        Split ratios for training, validation and test sets. They must sum to one.
+    target_ntype : str, optional
+        The node type to add split mask for.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes to predict.
+    train_idx : Tensor
+        An 1-D integer tensor of training node IDs.
+    val_idx : Tensor
+        An 1-D integer tensor of validation node IDs.
+    test_idx : Tensor
+        An 1-D integer tensor of test node IDs.
+
+    Examples
+    --------
+    >>> ds = dgl.data.AmazonCoBuyComputerDataset()
+    >>> print(ds)
+    Dataset("amazon_co_buy_computer", num_graphs=1, save_path=...)
+    >>> new_ds = dgl.data.AsNodePredDataset(ds, [0.8, 0.1, 0.1])
+    >>> print(new_ds)
+    Dataset("amazon_co_buy_computer-as-nodepred", num_graphs=1, save_path=...)
+    >>> print('train_mask' in new_ds[0].ndata)
+    True
+    """
+
+    def __init__(self, dataset, split_ratio=None, target_ntype=None, **kwargs):
+        self.dataset = dataset
+        self.split_ratio = split_ratio
+        self.target_ntype = target_ntype
+        super().__init__(
+            self.dataset.name + "-as-nodepred",
+            hash_key=(split_ratio, target_ntype, dataset.name, "nodepred"),
+            **kwargs
+        )
+
+    def process(self):
+        is_ogb = hasattr(self.dataset, "get_idx_split")
+        if is_ogb:
+            g, label = self.dataset[0]
+            self.g = g.clone()
+            self.g.ndata["label"] = F.reshape(label, (g.num_nodes(),))
+        else:
+            self.g = self.dataset[0].clone()
+
+        if "label" not in self.g.nodes[self.target_ntype].data:
+            raise ValueError(
+                "Missing node labels. Make sure labels are stored "
+                "under name 'label'."
+            )
+
+        if self.split_ratio is None:
+            if is_ogb:
+                split = self.dataset.get_idx_split()
+                train_idx, val_idx, test_idx = (
+                    split["train"],
+                    split["valid"],
+                    split["test"],
+                )
+                n = self.g.num_nodes()
+                train_mask = utils.generate_mask_tensor(
+                    utils.idx2mask(train_idx, n)
+                )
+                val_mask = utils.generate_mask_tensor(
+                    utils.idx2mask(val_idx, n)
+                )
+                test_mask = utils.generate_mask_tensor(
+                    utils.idx2mask(test_idx, n)
+                )
+                self.g.ndata["train_mask"] = train_mask
+                self.g.ndata["val_mask"] = val_mask
+                self.g.ndata["test_mask"] = test_mask
+            else:
+                assert (
+                    "train_mask" in self.g.nodes[self.target_ntype].data
+                ), "train_mask is not provided, please specify split_ratio to generate the masks"
+                assert (
+                    "val_mask" in self.g.nodes[self.target_ntype].data
+                ), "val_mask is not provided, please specify split_ratio to generate the masks"
+                assert (
+                    "test_mask" in self.g.nodes[self.target_ntype].data
+                ), "test_mask is not provided, please specify split_ratio to generate the masks"
+        else:
+            if self.verbose:
+                print("Generating train/val/test masks...")
+            utils.add_nodepred_split(self, self.split_ratio, self.target_ntype)
+
+        self._set_split_index()
+
+        self.num_classes = getattr(self.dataset, "num_classes", None)
+        if self.num_classes is None:
+            self.num_classes = len(
+                F.unique(self.g.nodes[self.target_ntype].data["label"])
+            )
+
+    def has_cache(self):
+        return os.path.isfile(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash))
+        )
+
+    def load(self):
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "r"
+        ) as f:
+            info = json.load(f)
+            if (
+                info["split_ratio"] != self.split_ratio
+                or info["target_ntype"] != self.target_ntype
+            ):
+                raise ValueError(
+                    "Provided split ratio is different from the cached file. "
+                    "Re-process the dataset."
+                )
+            self.split_ratio = info["split_ratio"]
+            self.target_ntype = info["target_ntype"]
+            self.num_classes = info["num_classes"]
+        gs, _ = utils.load_graphs(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash))
+        )
+        self.g = gs[0]
+        self._set_split_index()
+
+    def save(self):
+        utils.save_graphs(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash)),
+            [self.g],
+        )
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "w"
+        ) as f:
+            json.dump(
+                {
+                    "split_ratio": self.split_ratio,
+                    "target_ntype": self.target_ntype,
+                    "num_classes": self.num_classes,
+                },
+                f,
+            )
+
+    def __getitem__(self, idx):
+        return self.g
+
+    def __len__(self):
+        return 1
+
+    def _set_split_index(self):
+        """Add train_idx/val_idx/test_idx as dataset attributes according to corresponding mask."""
+        ndata = self.g.nodes[self.target_ntype].data
+        self.train_idx = F.nonzero_1d(ndata["train_mask"])
+        self.val_idx = F.nonzero_1d(ndata["val_mask"])
+        self.test_idx = F.nonzero_1d(ndata["test_mask"])
+
+
+def negative_sample(g, num_samples):
+    """Random sample negative edges from graph, excluding self-loops,
+    the result samples might be less than num_samples
+    """
+    num_nodes = g.num_nodes()
+    redundancy = _calc_redundancy(num_samples, g.num_edges(), num_nodes**2)
+    sample_size = int(num_samples * (1 + redundancy))
+    edges = np.random.randint(0, num_nodes, size=(2, sample_size))
+    edges = np.unique(edges, axis=1)
+    # remove self loop
+    mask_self_loop = edges[0] == edges[1]
+    # remove existing edges
+    has_edges = F.asnumpy(g.has_edges_between(edges[0], edges[1]))
+    mask = ~(np.logical_or(mask_self_loop, has_edges))
+    edges = edges[:, mask]
+    if edges.shape[1] >= num_samples:
+        edges = edges[:, :num_samples]
+    return edges
+
+
+class AsLinkPredDataset(DGLDataset):
+    """Repurpose a dataset for link prediction task.
+
+    The created dataset will include data needed for link prediction.
+    Currently it only supports homogeneous graphs.
+    It will keep only the first graph in the provided dataset and
+    generate train/val/test edges according to the given split ratio,
+    and the correspondent negative edges based on the neg_ratio. The generated
+    edges will be cached to disk for fast re-loading. If the provided split ratio
+    differs from the cached one, it will re-process the dataset properly.
+
+    Parameters
+    ----------
+    dataset : DGLDataset
+        The dataset to be converted.
+    split_ratio : (float, float, float), optional
+        Split ratios for training, validation and test sets. Must sum to one.
+    neg_ratio : int, optional
+        Indicate how much negative samples to be sampled
+        The number of the negative samples will be equal or less than neg_ratio * num_positive_edges.
+
+    Attributes
+    -------
+    feat_size: int
+        The size of the feature dimension in the graph
+    train_graph: DGLGraph
+        The DGLGraph for training
+    val_edges: Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]
+        The validation set edges, encoded as
+        ((positive_edge_src, positive_edge_dst), (negative_edge_src, negative_edge_dst))
+    test_edges: Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]
+        The test set edges, encoded as
+        ((positive_edge_src, positive_edge_dst), (negative_edge_src, negative_edge_dst))
+
+    Examples
+    --------
+    >>> ds = dgl.data.CoraGraphDataset()
+    >>> print(ds)
+    Dataset("cora_v2", num_graphs=1, save_path=...)
+    >>> new_ds = dgl.data.AsLinkPredDataset(ds, [0.8, 0.1, 0.1])
+    >>> print(new_ds)
+    Dataset("cora_v2-as-linkpred", num_graphs=1, save_path=/home/ubuntu/.dgl/cora_v2-as-linkpred)
+    >>> print(hasattr(new_ds, "test_edges"))
+    True
+    """
+
+    def __init__(self, dataset, split_ratio=None, neg_ratio=3, **kwargs):
+        self.g = dataset[0]
+        self.num_nodes = self.g.num_nodes()
+        self.dataset = dataset
+        self.split_ratio = split_ratio
+        self.neg_ratio = neg_ratio
+        super().__init__(
+            dataset.name + "-as-linkpred",
+            hash_key=(neg_ratio, split_ratio, dataset.name, "linkpred"),
+            **kwargs
+        )
+
+    def process(self):
+        if self.split_ratio is None:
+            # Handle logics for OGB link prediction dataset
+            assert hasattr(
+                self.dataset, "get_edge_split"
+            ), "dataset doesn't have get_edge_split method, please specify split_ratio and neg_ratio to generate the split"
+            # This is likely to be an ogb dataset
+            self.edge_split = self.dataset.get_edge_split()
+            self._train_graph = self.g
+            if "source_node" in self.edge_split["test"]:
+                # Probably ogbl-citation2
+                pos_e = (
+                    self.edge_split["valid"]["source_node"],
+                    self.edge_split["valid"]["target_node"],
+                )
+                neg_e_size = self.edge_split["valid"]["target_node_neg"].shape[
+                    -1
+                ]
+                neg_e_src = np.repeat(
+                    self.edge_split["valid"]["source_node"], neg_e_size
+                )
+                neg_e_dst = np.reshape(
+                    self.edge_split["valid"]["target_node_neg"], -1
+                )
+                self._val_edges = pos_e, (neg_e_src, neg_e_dst)
+                pos_e = (
+                    self.edge_split["test"]["source_node"],
+                    self.edge_split["test"]["target_node"],
+                )
+                neg_e_size = self.edge_split["test"]["target_node_neg"].shape[
+                    -1
+                ]
+                neg_e_src = np.repeat(
+                    self.edge_split["test"]["source_node"], neg_e_size
+                )
+                neg_e_dst = np.reshape(
+                    self.edge_split["test"]["target_node_neg"], -1
+                )
+                self._test_edges = pos_e, (neg_e_src, neg_e_dst)
+            elif "edge" in self.edge_split["test"]:
+                # Probably ogbl-collab
+                pos_e_tensor, neg_e_tensor = (
+                    self.edge_split["valid"]["edge"],
+                    self.edge_split["valid"]["edge_neg"],
+                )
+                pos_e = (pos_e_tensor[:, 0], pos_e_tensor[:, 1])
+                neg_e = (neg_e_tensor[:, 0], neg_e_tensor[:, 1])
+                self._val_edges = pos_e, neg_e
+
+                pos_e_tensor, neg_e_tensor = (
+                    self.edge_split["test"]["edge"],
+                    self.edge_split["test"]["edge_neg"],
+                )
+                pos_e = (pos_e_tensor[:, 0], pos_e_tensor[:, 1])
+                neg_e = (neg_e_tensor[:, 0], neg_e_tensor[:, 1])
+                self._test_edges = pos_e, neg_e
+            # delete edge split to save memory
+            self.edge_split = None
+        else:
+            assert self.split_ratio is not None, "Need to specify split_ratio"
+            assert self.neg_ratio is not None, "Need to specify neg_ratio"
+            ratio = self.split_ratio
+            graph = self.dataset[0]
+            n = graph.num_edges()
+            src, dst = graph.edges()
+            src, dst = F.asnumpy(src), F.asnumpy(dst)
+            n_train, n_val, n_test = (
+                int(n * ratio[0]),
+                int(n * ratio[1]),
+                int(n * ratio[2]),
+            )
+
+            idx = np.random.permutation(n)
+            train_pos_idx = idx[:n_train]
+            val_pos_idx = idx[n_train : n_train + n_val]
+            test_pos_idx = idx[n_train + n_val :]
+            neg_src, neg_dst = negative_sample(
+                graph, self.neg_ratio * (n_val + n_test)
+            )
+            neg_n_val, neg_n_test = (
+                self.neg_ratio * n_val,
+                self.neg_ratio * n_test,
+            )
+            neg_val_src, neg_val_dst = neg_src[:neg_n_val], neg_dst[:neg_n_val]
+            neg_test_src, neg_test_dst = (
+                neg_src[neg_n_val:],
+                neg_dst[neg_n_val:],
+            )
+            self._val_edges = (
+                F.tensor(src[val_pos_idx]),
+                F.tensor(dst[val_pos_idx]),
+            ), (F.tensor(neg_val_src), F.tensor(neg_val_dst))
+            self._test_edges = (
+                F.tensor(src[test_pos_idx]),
+                F.tensor(dst[test_pos_idx]),
+            ), (F.tensor(neg_test_src), F.tensor(neg_test_dst))
+            self._train_graph = create_dgl_graph(
+                (src[train_pos_idx], dst[train_pos_idx]),
+                num_nodes=self.num_nodes,
+            )
+            self._train_graph.ndata["feat"] = graph.ndata["feat"]
+
+    def has_cache(self):
+        return os.path.isfile(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash))
+        )
+
+    def load(self):
+        gs, tensor_dict = utils.load_graphs(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash))
+        )
+        self.g = gs[0]
+        self._train_graph = self.g
+        self._val_edges = (
+            tensor_dict["val_pos_src"],
+            tensor_dict["val_pos_dst"],
+        ), (tensor_dict["val_neg_src"], tensor_dict["val_neg_dst"])
+        self._test_edges = (
+            tensor_dict["test_pos_src"],
+            tensor_dict["test_pos_dst"],
+        ), (tensor_dict["test_neg_src"], tensor_dict["test_neg_dst"])
+
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "r"
+        ) as f:
+            info = json.load(f)
+            self.split_ratio = info["split_ratio"]
+            self.neg_ratio = info["neg_ratio"]
+
+    def save(self):
+        tensor_dict = {
+            "val_pos_src": self._val_edges[0][0],
+            "val_pos_dst": self._val_edges[0][1],
+            "val_neg_src": self._val_edges[1][0],
+            "val_neg_dst": self._val_edges[1][1],
+            "test_pos_src": self._test_edges[0][0],
+            "test_pos_dst": self._test_edges[0][1],
+            "test_neg_src": self._test_edges[1][0],
+            "test_neg_dst": self._test_edges[1][1],
+        }
+        utils.save_graphs(
+            os.path.join(self.save_path, "graph_{}.bin".format(self.hash)),
+            [self._train_graph],
+            tensor_dict,
+        )
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "w"
+        ) as f:
+            json.dump(
+                {"split_ratio": self.split_ratio, "neg_ratio": self.neg_ratio},
+                f,
+            )
+
+    @property
+    def feat_size(self):
+        return self._train_graph.ndata["feat"].shape[-1]
+
+    @property
+    def train_graph(self):
+        return self._train_graph
+
+    @property
+    def val_edges(self):
+        return self._val_edges
+
+    @property
+    def test_edges(self):
+        return self._test_edges
+
+    def __getitem__(self, idx):
+        return self.g
+
+    def __len__(self):
+        return 1
+
+
+class AsGraphPredDataset(DGLDataset):
+    """Repurpose a dataset for standard graph property prediction task.
+
+    The created dataset will include data needed for graph property prediction.
+    Currently it only supports homogeneous graphs.
+
+    The class converts a given dataset into a new dataset object such that:
+
+      - It stores ``len(dataset)`` graphs.
+      - The i-th graph and its label is accessible from ``dataset[i]``.
+
+    The class will generate a train/val/test split if :attr:`split_ratio` is provided.
+    The generated split will be cached to disk for fast re-loading. If the provided split
+    ratio differs from the cached one, it will re-process the dataset properly.
+
+    Parameters
+    ----------
+    dataset : DGLDataset
+        The dataset to be converted.
+    split_ratio : (float, float, float), optional
+        Split ratios for training, validation and test sets. They must sum to one.
+
+    Attributes
+    ----------
+    num_tasks : int
+        Number of tasks to predict.
+    num_classes : int
+        Number of classes to predict per task, None for regression datasets.
+    train_idx : Tensor
+        An 1-D integer tensor of training node IDs.
+    val_idx : Tensor
+        An 1-D integer tensor of validation node IDs.
+    test_idx : Tensor
+        An 1-D integer tensor of test node IDs.
+    node_feat_size : int
+        Input node feature size, None if not applicable.
+    edge_feat_size : int
+        Input edge feature size, None if not applicable.
+
+    Examples
+    --------
+
+    >>> from dgl.data import AsGraphPredDataset
+    >>> from ogb.graphproppred import DglGraphPropPredDataset
+    >>> dataset = DglGraphPropPredDataset(name='ogbg-molhiv')
+    >>> new_dataset = AsGraphPredDataset(dataset)
+    >>> print(new_dataset)
+    Dataset("ogbg-molhiv-as-graphpred", num_graphs=41127, save_path=...)
+    >>> print(len(new_dataset))
+    41127
+    >>> print(new_dataset[0])
+    (Graph(num_nodes=19, num_edges=40,
+           ndata_schemes={'feat': Scheme(shape=(9,), dtype=torch.int64)}
+           edata_schemes={'feat': Scheme(shape=(3,), dtype=torch.int64)}), tensor([0]))
+    """
+
+    def __init__(self, dataset, split_ratio=None, **kwargs):
+        self.dataset = dataset
+        self.split_ratio = split_ratio
+        super().__init__(
+            dataset.name + "-as-graphpred",
+            hash_key=(split_ratio, dataset.name, "graphpred"),
+            **kwargs
+        )
+
+    def process(self):
+        is_ogb = hasattr(self.dataset, "get_idx_split")
+        if self.split_ratio is None:
+            if is_ogb:
+                split = self.dataset.get_idx_split()
+                self.train_idx = split["train"]
+                self.val_idx = split["valid"]
+                self.test_idx = split["test"]
+            else:
+                # Handle FakeNewsDataset
+                try:
+                    self.train_idx = F.nonzero_1d(self.dataset.train_mask)
+                    self.val_idx = F.nonzero_1d(self.dataset.val_mask)
+                    self.test_idx = F.nonzero_1d(self.dataset.test_mask)
+                except:
+                    raise DGLError(
+                        "The input dataset does not have default train/val/test\
+                        split. Please specify split_ratio to generate the split."
+                    )
+        else:
+            if self.verbose:
+                print("Generating train/val/test split...")
+            train_ratio, val_ratio, _ = self.split_ratio
+            num_graphs = len(self.dataset)
+            num_train = int(num_graphs * train_ratio)
+            num_val = int(num_graphs * val_ratio)
+
+            idx = np.random.permutation(num_graphs)
+            self.train_idx = F.tensor(idx[:num_train])
+            self.val_idx = F.tensor(idx[num_train : num_train + num_val])
+            self.test_idx = F.tensor(idx[num_train + num_val :])
+
+        if hasattr(self.dataset, "num_classes"):
+            # GINDataset, MiniGCDataset, FakeNewsDataset, TUDataset,
+            # LegacyTUDataset, BA2MotifDataset
+            self.num_classes = self.dataset.num_classes
+        else:
+            # None for multi-label classification and regression
+            self.num_classes = None
+
+        if hasattr(self.dataset, "num_tasks"):
+            # OGB datasets
+            self.num_tasks = self.dataset.num_tasks
+        else:
+            self.num_tasks = 1
+
+    def has_cache(self):
+        return os.path.isfile(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash))
+        )
+
+    def load(self):
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "r"
+        ) as f:
+            info = json.load(f)
+            if info["split_ratio"] != self.split_ratio:
+                raise ValueError(
+                    "Provided split ratio is different from the cached file. "
+                    "Re-process the dataset."
+                )
+            self.split_ratio = info["split_ratio"]
+            self.num_tasks = info["num_tasks"]
+            self.num_classes = info["num_classes"]
+
+        split = np.load(
+            os.path.join(self.save_path, "split_{}.npz".format(self.hash))
+        )
+        self.train_idx = F.zerocopy_from_numpy(split["train_idx"])
+        self.val_idx = F.zerocopy_from_numpy(split["val_idx"])
+        self.test_idx = F.zerocopy_from_numpy(split["test_idx"])
+
+    def save(self):
+        if not os.path.exists(self.save_path):
+            os.makedirs(self.save_path)
+        with open(
+            os.path.join(self.save_path, "info_{}.json".format(self.hash)), "w"
+        ) as f:
+            json.dump(
+                {
+                    "split_ratio": self.split_ratio,
+                    "num_tasks": self.num_tasks,
+                    "num_classes": self.num_classes,
+                },
+                f,
+            )
+        np.savez(
+            os.path.join(self.save_path, "split_{}.npz".format(self.hash)),
+            train_idx=F.zerocopy_to_numpy(self.train_idx),
+            val_idx=F.zerocopy_to_numpy(self.val_idx),
+            test_idx=F.zerocopy_to_numpy(self.test_idx),
+        )
+
+    def __getitem__(self, idx):
+        return self.dataset[idx]
+
+    def __len__(self):
+        return len(self.dataset)
+
+    @property
+    def node_feat_size(self):
+        g = self[0][0]
+        return g.ndata["feat"].shape[-1] if "feat" in g.ndata else None
+
+    @property
+    def edge_feat_size(self):
+        g = self[0][0]
+        return g.edata["feat"].shape[-1] if "feat" in g.edata else None

python/dgl/data/bitcoinotc.py

@@ -0,0 +1,191 @@
+""" BitcoinOTC dataset for fraud detection """
+import datetime
+import gzip
+import os
+import shutil
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import check_sha1, download, load_graphs, makedirs, save_graphs
+
+
+class BitcoinOTCDataset(DGLBuiltinDataset):
+    r"""BitcoinOTC dataset for fraud detection
+
+    This is who-trusts-whom network of people who trade using Bitcoin on
+    a platform called Bitcoin OTC. Since Bitcoin users are anonymous,
+    there is a need to maintain a record of users' reputation to prevent
+    transactions with fraudulent and risky users.
+
+    Offical website: `<https://snap.stanford.edu/data/soc-sign-bitcoin-otc.html>`_
+
+    Bitcoin OTC dataset statistics:
+
+    - Nodes: 5,881
+    - Edges: 35,592
+    - Range of edge weight: -10 to +10
+    - Percentage of positive edges: 89%
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose: bool
+        Whether to print out progress information.
+        Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    graphs : list
+        A list of DGLGraph objects
+    is_temporal : bool
+        Indicate whether the graphs are temporal graphs
+
+    Raises
+    ------
+    UserWarning
+        If the raw data is changed in the remote server by the author.
+
+    Examples
+    --------
+    >>> dataset = BitcoinOTCDataset()
+    >>> len(dataset)
+    136
+    >>> for g in dataset:
+    ....    # get edge feature
+    ....    edge_weights = g.edata['h']
+    ....    # your code here
+    >>>
+    """
+
+    _url = "https://snap.stanford.edu/data/soc-sign-bitcoinotc.csv.gz"
+    _sha1_str = "c14281f9e252de0bd0b5f1c6e2bae03123938641"
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(BitcoinOTCDataset, self).__init__(
+            name="bitcoinotc",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        gz_file_path = os.path.join(self.raw_dir, self.name + ".csv.gz")
+        download(self.url, path=gz_file_path)
+        if not check_sha1(gz_file_path, self._sha1_str):
+            raise UserWarning(
+                "File {} is downloaded but the content hash does not match."
+                "The repo may be outdated or download may be incomplete. "
+                "Otherwise you can create an issue for it.".format(
+                    self.name + ".csv.gz"
+                )
+            )
+        self._extract_gz(gz_file_path, self.raw_path)
+
+    def process(self):
+        filename = os.path.join(self.save_path, self.name + ".csv")
+        data = np.loadtxt(filename, delimiter=",").astype(np.int64)
+        data[:, 0:2] = data[:, 0:2] - data[:, 0:2].min()
+        delta = datetime.timedelta(days=14).total_seconds()
+        # The source code is not released, but the paper indicates there're
+        # totally 137 samples. The cutoff below has exactly 137 samples.
+        time_index = np.around((data[:, 3] - data[:, 3].min()) / delta).astype(
+            np.int64
+        )
+
+        self._graphs = []
+        for i in range(time_index.max()):
+            row_mask = time_index <= i
+            edges = data[row_mask][:, 0:2]
+            rate = data[row_mask][:, 2]
+            g = dgl_graph((edges[:, 0], edges[:, 1]))
+            g.edata["h"] = F.tensor(
+                rate.reshape(-1, 1), dtype=F.data_type_dict["int64"]
+            )
+            self._graphs.append(g)
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, "dgl_graph.bin")
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def save(self):
+        save_graphs(self.graph_path, self.graphs)
+
+    def load(self):
+        self._graphs = load_graphs(self.graph_path)[0]
+
+    @property
+    def graphs(self):
+        return self._graphs
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return len(self.graphs)
+
+    def __getitem__(self, item):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        item : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``edata['h']`` : edge weights
+        """
+        if self._transform is None:
+            return self.graphs[item]
+        else:
+            return self._transform(self.graphs[item])
+
+    @property
+    def is_temporal(self):
+        r"""Are the graphs temporal graphs
+
+        Returns
+        -------
+        bool
+        """
+        return True
+
+    def _extract_gz(self, file, target_dir, overwrite=False):
+        if os.path.exists(target_dir) and not overwrite:
+            return
+        print("Extracting file to {}".format(target_dir))
+        fname = os.path.basename(file)
+        makedirs(target_dir)
+        out_file_path = os.path.join(target_dir, fname[:-3])
+        with gzip.open(file, "rb") as f_in:
+            with open(out_file_path, "wb") as f_out:
+                shutil.copyfileobj(f_in, f_out)
+
+
+BitcoinOTC = BitcoinOTCDataset

python/dgl/data/citation_graph.py

@@ -0,0 +1,953 @@
+"""Cora, citeseer, pubmed dataset.
+
+(lingfan): following dataset loading and preprocessing code from tkipf/gcn
+https://github.com/tkipf/gcn/blob/master/gcn/utils.py
+"""
+
+from __future__ import absolute_import
+
+import os, sys
+import pickle as pkl
+import warnings
+
+import networkx as nx
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F, convert
+from ..batch import batch as batch_graphs
+from ..convert import from_networkx, graph as dgl_graph, to_networkx
+from ..transforms import reorder_graph
+from .dgl_dataset import DGLBuiltinDataset
+
+from .utils import (
+    _get_dgl_url,
+    deprecate_function,
+    deprecate_property,
+    generate_mask_tensor,
+    load_graphs,
+    load_info,
+    makedirs,
+    save_graphs,
+    save_info,
+)
+
+backend = os.environ.get("DGLBACKEND", "pytorch")
+
+
+def _pickle_load(pkl_file):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=DeprecationWarning)
+        if sys.version_info > (3, 0):
+            return pkl.load(pkl_file, encoding="latin1")
+        else:
+            return pkl.load(pkl_file)
+
+
+class CitationGraphDataset(DGLBuiltinDataset):
+    r"""The citation graph dataset, including cora, citeseer and pubmeb.
+    Nodes mean authors and edges mean citation relationships.
+
+    Parameters
+    -----------
+    name: str
+      name can be 'cora', 'citeseer' or 'pubmed'.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`. Default: False.
+    """
+
+    _urls = {
+        "cora_v2": "dataset/cora_v2.zip",
+        "citeseer": "dataset/citeseer.zip",
+        "pubmed": "dataset/pubmed.zip",
+    }
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        reverse_edge=True,
+        transform=None,
+        reorder=False,
+    ):
+        assert name.lower() in ["cora", "citeseer", "pubmed"]
+
+        # Previously we use the pre-processing in pygcn (https://github.com/tkipf/pygcn)
+        # for Cora, which is slightly different from the one used in the GCN paper
+        if name.lower() == "cora":
+            name = "cora_v2"
+
+        url = _get_dgl_url(self._urls[name])
+        self._reverse_edge = reverse_edge
+        self._reorder = reorder
+
+        super(CitationGraphDataset, self).__init__(
+            name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """Loads input data from data directory and reorder graph for better locality
+
+        ind.name.x => the feature vectors of the training instances as scipy.sparse.csr.csr_matrix object;
+        ind.name.tx => the feature vectors of the test instances as scipy.sparse.csr.csr_matrix object;
+        ind.name.allx => the feature vectors of both labeled and unlabeled training instances
+            (a superset of ind.name.x) as scipy.sparse.csr.csr_matrix object;
+        ind.name.y => the one-hot labels of the labeled training instances as numpy.ndarray object;
+        ind.name.ty => the one-hot labels of the test instances as numpy.ndarray object;
+        ind.name.ally => the labels for instances in ind.name.allx as numpy.ndarray object;
+        ind.name.graph => a dict in the format {index: [index_of_neighbor_nodes]} as collections.defaultdict
+            object;
+        ind.name.test.index => the indices of test instances in graph, for the inductive setting as list object.
+        """
+        root = self.raw_path
+        objnames = ["x", "y", "tx", "ty", "allx", "ally", "graph"]
+        objects = []
+        for i in range(len(objnames)):
+            with open(
+                "{}/ind.{}.{}".format(root, self.name, objnames[i]), "rb"
+            ) as f:
+                objects.append(_pickle_load(f))
+
+        x, y, tx, ty, allx, ally, graph = tuple(objects)
+        test_idx_reorder = _parse_index_file(
+            "{}/ind.{}.test.index".format(root, self.name)
+        )
+        test_idx_range = np.sort(test_idx_reorder)
+
+        if self.name == "citeseer":
+            # Fix citeseer dataset (there are some isolated nodes in the graph)
+            # Find isolated nodes, add them as zero-vecs into the right position
+            test_idx_range_full = range(
+                min(test_idx_reorder), max(test_idx_reorder) + 1
+            )
+            tx_extended = sp.lil_matrix((len(test_idx_range_full), x.shape[1]))
+            tx_extended[test_idx_range - min(test_idx_range), :] = tx
+            tx = tx_extended
+            ty_extended = np.zeros((len(test_idx_range_full), y.shape[1]))
+            ty_extended[test_idx_range - min(test_idx_range), :] = ty
+            ty = ty_extended
+
+        features = sp.vstack((allx, tx)).tolil()
+        features[test_idx_reorder, :] = features[test_idx_range, :]
+
+        if self.reverse_edge:
+            graph = nx.DiGraph(nx.from_dict_of_lists(graph))
+            g = from_networkx(graph)
+        else:
+            graph = nx.Graph(nx.from_dict_of_lists(graph))
+            edges = list(graph.edges())
+            u, v = map(list, zip(*edges))
+            g = dgl_graph((u, v))
+
+        onehot_labels = np.vstack((ally, ty))
+        onehot_labels[test_idx_reorder, :] = onehot_labels[test_idx_range, :]
+        labels = np.argmax(onehot_labels, 1)
+
+        idx_test = test_idx_range.tolist()
+        idx_train = range(len(y))
+        idx_val = range(len(y), len(y) + 500)
+
+        train_mask = generate_mask_tensor(
+            _sample_mask(idx_train, labels.shape[0])
+        )
+        val_mask = generate_mask_tensor(_sample_mask(idx_val, labels.shape[0]))
+        test_mask = generate_mask_tensor(
+            _sample_mask(idx_test, labels.shape[0])
+        )
+
+        g.ndata["train_mask"] = train_mask
+        g.ndata["val_mask"] = val_mask
+        g.ndata["test_mask"] = test_mask
+        g.ndata["label"] = F.tensor(labels)
+        g.ndata["feat"] = F.tensor(
+            _preprocess_features(features), dtype=F.data_type_dict["float32"]
+        )
+        self._num_classes = onehot_labels.shape[1]
+        self._labels = labels
+        if self._reorder:
+            self._g = reorder_graph(
+                g,
+                node_permute_algo="rcmk",
+                edge_permute_algo="dst",
+                store_ids=False,
+            )
+        else:
+            self._g = g
+
+        if self.verbose:
+            print("Finished data loading and preprocessing.")
+            print("  NumNodes: {}".format(self._g.num_nodes()))
+            print("  NumEdges: {}".format(self._g.num_edges()))
+            print("  NumFeats: {}".format(self._g.ndata["feat"].shape[1]))
+            print("  NumClasses: {}".format(self.num_classes))
+            print(
+                "  NumTrainingSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["train_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumValidationSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["val_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumTestSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["test_mask"]).shape[0]
+                )
+            )
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.save_name + ".bin")
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, self.save_name + ".pkl")
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+
+        return False
+
+    def save(self):
+        """save the graph list and the labels"""
+        save_graphs(str(self.graph_path), self._g)
+        save_info(str(self.info_path), {"num_classes": self.num_classes})
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+
+        info = load_info(str(self.info_path))
+        graph = graphs[0]
+        self._g = graph
+        # for compatability
+        graph = graph.clone()
+        graph.ndata.pop("train_mask")
+        graph.ndata.pop("val_mask")
+        graph.ndata.pop("test_mask")
+        graph.ndata.pop("feat")
+        graph.ndata.pop("label")
+        graph = to_networkx(graph)
+
+        self._num_classes = info["num_classes"]
+        self._g.ndata["train_mask"] = generate_mask_tensor(
+            F.asnumpy(self._g.ndata["train_mask"])
+        )
+        self._g.ndata["val_mask"] = generate_mask_tensor(
+            F.asnumpy(self._g.ndata["val_mask"])
+        )
+        self._g.ndata["test_mask"] = generate_mask_tensor(
+            F.asnumpy(self._g.ndata["test_mask"])
+        )
+        # hack for mxnet compatability
+
+        if self.verbose:
+            print("  NumNodes: {}".format(self._g.num_nodes()))
+            print("  NumEdges: {}".format(self._g.num_edges()))
+            print("  NumFeats: {}".format(self._g.ndata["feat"].shape[1]))
+            print("  NumClasses: {}".format(self.num_classes))
+            print(
+                "  NumTrainingSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["train_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumValidationSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["val_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumTestSamples: {}".format(
+                    F.nonzero_1d(self._g.ndata["test_mask"]).shape[0]
+                )
+            )
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def save_name(self):
+        return self.name + "_dgl_graph"
+
+    @property
+    def num_labels(self):
+        deprecate_property("dataset.num_labels", "dataset.num_classes")
+        return self.num_classes
+
+    @property
+    def num_classes(self):
+        return self._num_classes
+
+    """ Citation graph is used in many examples
+        We preserve these properties for compatability.
+    """
+
+    @property
+    def reverse_edge(self):
+        return self._reverse_edge
+
+
+def _preprocess_features(features):
+    """Row-normalize feature matrix and convert to tuple representation"""
+    features = _normalize(features)
+    return np.asarray(features.todense())
+
+
+def _parse_index_file(filename):
+    """Parse index file."""
+    index = []
+    for line in open(filename):
+        index.append(int(line.strip()))
+    return index
+
+
+def _sample_mask(idx, l):
+    """Create mask."""
+    mask = np.zeros(l)
+    mask[idx] = 1
+    return mask
+
+
+class CoraGraphDataset(CitationGraphDataset):
+    r"""Cora citation network dataset.
+
+    Nodes mean paper and edges mean citation
+    relationships. Each node has a predefined
+    feature with 1433 dimensions. The dataset is
+    designed for the node classification task.
+    The task is to predict the category of
+    certain paper.
+
+    Statistics:
+
+    - Nodes: 2708
+    - Edges: 10556
+    - Number of Classes: 7
+    - Label split:
+
+        - Train: 140
+        - Valid: 500
+        - Test: 1000
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`. Default: False.
+
+    Attributes
+    ----------
+    num_classes: int
+        Number of label classes
+
+    Notes
+    -----
+    The node feature is row-normalized.
+
+    Examples
+    --------
+    >>> dataset = CoraGraphDataset()
+    >>> g = dataset[0]
+    >>> num_class = dataset.num_classes
+    >>>
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>>
+    >>> # get labels
+    >>> label = g.ndata['label']
+
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        reverse_edge=True,
+        transform=None,
+        reorder=False,
+    ):
+        name = "cora"
+
+        super(CoraGraphDataset, self).__init__(
+            name,
+            raw_dir,
+            force_reload,
+            verbose,
+            reverse_edge,
+            transform,
+            reorder,
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, CoraGraphDataset has only one graph object
+
+        Return
+        ------
+        :class:`dgl.DGLGraph`
+
+            graph structure, node features and labels.
+
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+            - ``ndata['feat']``: node feature
+            - ``ndata['label']``: ground truth labels
+        """
+        return super(CoraGraphDataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(CoraGraphDataset, self).__len__()
+
+
+class CiteseerGraphDataset(CitationGraphDataset):
+    r"""Citeseer citation network dataset.
+
+    Nodes mean scientific publications and edges
+    mean citation relationships. Each node has a
+    predefined feature with 3703 dimensions. The
+    dataset is designed for the node classification
+    task. The task is to predict the category of
+    certain publication.
+
+    Statistics:
+
+    - Nodes: 3327
+    - Edges: 9228
+    - Number of Classes: 6
+    - Label Split:
+
+        - Train: 120
+        - Valid: 500
+        - Test: 1000
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`. Default: False.
+
+    Attributes
+    ----------
+    num_classes: int
+        Number of label classes
+
+    Notes
+    -----
+    The node feature is row-normalized.
+
+    In citeseer dataset, there are some isolated nodes in the graph.
+    These isolated nodes are added as zero-vecs into the right position.
+
+    Examples
+    --------
+    >>> dataset = CiteseerGraphDataset()
+    >>> g = dataset[0]
+    >>> num_class = dataset.num_classes
+    >>>
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>>
+    >>> # get labels
+    >>> label = g.ndata['label']
+
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        reverse_edge=True,
+        transform=None,
+        reorder=False,
+    ):
+        name = "citeseer"
+
+        super(CiteseerGraphDataset, self).__init__(
+            name,
+            raw_dir,
+            force_reload,
+            verbose,
+            reverse_edge,
+            transform,
+            reorder,
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, CiteseerGraphDataset has only one graph object
+
+        Return
+        ------
+        :class:`dgl.DGLGraph`
+
+            graph structure, node features and labels.
+
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+            - ``ndata['feat']``: node feature
+            - ``ndata['label']``: ground truth labels
+        """
+        return super(CiteseerGraphDataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(CiteseerGraphDataset, self).__len__()
+
+
+class PubmedGraphDataset(CitationGraphDataset):
+    r"""Pubmed citation network dataset.
+
+    Nodes mean scientific publications and edges
+    mean citation relationships. Each node has a
+    predefined feature with 500 dimensions. The
+    dataset is designed for the node classification
+    task. The task is to predict the category of
+    certain publication.
+
+    Statistics:
+
+    - Nodes: 19717
+    - Edges: 88651
+    - Number of Classes: 3
+    - Label Split:
+
+        - Train: 60
+        - Valid: 500
+        - Test: 1000
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`. Default: False.
+
+    Attributes
+    ----------
+    num_classes: int
+        Number of label classes
+
+    Notes
+    -----
+    The node feature is row-normalized.
+
+    Examples
+    --------
+    >>> dataset = PubmedGraphDataset()
+    >>> g = dataset[0]
+    >>> num_class = dataset.num_of_class
+    >>>
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>>
+    >>> # get labels
+    >>> label = g.ndata['label']
+
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        reverse_edge=True,
+        transform=None,
+        reorder=False,
+    ):
+        name = "pubmed"
+
+        super(PubmedGraphDataset, self).__init__(
+            name,
+            raw_dir,
+            force_reload,
+            verbose,
+            reverse_edge,
+            transform,
+            reorder,
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, PubmedGraphDataset has only one graph object
+
+        Return
+        ------
+        :class:`dgl.DGLGraph`
+
+            graph structure, node features and labels.
+
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+            - ``ndata['feat']``: node feature
+            - ``ndata['label']``: ground truth labels
+        """
+        return super(PubmedGraphDataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(PubmedGraphDataset, self).__len__()
+
+
+def load_cora(
+    raw_dir=None,
+    force_reload=False,
+    verbose=True,
+    reverse_edge=True,
+    transform=None,
+):
+    """Get CoraGraphDataset
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Return
+    -------
+    CoraGraphDataset
+    """
+    data = CoraGraphDataset(
+        raw_dir, force_reload, verbose, reverse_edge, transform
+    )
+    return data
+
+
+def load_citeseer(
+    raw_dir=None,
+    force_reload=False,
+    verbose=True,
+    reverse_edge=True,
+    transform=None,
+):
+    """Get CiteseerGraphDataset
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Return
+    -------
+    CiteseerGraphDataset
+    """
+    data = CiteseerGraphDataset(
+        raw_dir, force_reload, verbose, reverse_edge, transform
+    )
+    return data
+
+
+def load_pubmed(
+    raw_dir=None,
+    force_reload=False,
+    verbose=True,
+    reverse_edge=True,
+    transform=None,
+):
+    """Get PubmedGraphDataset
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    reverse_edge : bool
+        Whether to add reverse edges in graph. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Return
+    -------
+    PubmedGraphDataset
+    """
+    data = PubmedGraphDataset(
+        raw_dir, force_reload, verbose, reverse_edge, transform
+    )
+    return data
+
+
+class CoraBinary(DGLBuiltinDataset):
+    """A mini-dataset for binary classification task using Cora.
+
+    After loaded, it has following members:
+
+    graphs : list of :class:`~dgl.DGLGraph`
+    pmpds : list of :class:`scipy.sparse.coo_matrix`
+    labels : list of :class:`numpy.ndarray`
+
+    Parameters
+    -----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose: bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        name = "cora_binary"
+        url = _get_dgl_url("dataset/cora_binary.zip")
+        super(CoraBinary, self).__init__(
+            name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        root = self.raw_path
+        # load graphs
+        self.graphs = []
+        with open("{}/graphs.txt".format(root), "r") as f:
+            elist = []
+            for line in f.readlines():
+                if line.startswith("graph"):
+                    if len(elist) != 0:
+                        self.graphs.append(dgl_graph(tuple(zip(*elist))))
+                    elist = []
+                else:
+                    u, v = line.strip().split(" ")
+                    elist.append((int(u), int(v)))
+            if len(elist) != 0:
+                self.graphs.append(dgl_graph(tuple(zip(*elist))))
+        with open("{}/pmpds.pkl".format(root), "rb") as f:
+            self.pmpds = _pickle_load(f)
+        self.labels = []
+        with open("{}/labels.txt".format(root), "r") as f:
+            cur = []
+            for line in f.readlines():
+                if line.startswith("graph"):
+                    if len(cur) != 0:
+                        self.labels.append(np.asarray(cur))
+                    cur = []
+                else:
+                    cur.append(int(line.strip()))
+            if len(cur) != 0:
+                self.labels.append(np.asarray(cur))
+        # sanity check
+        assert len(self.graphs) == len(self.pmpds)
+        assert len(self.graphs) == len(self.labels)
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.save_name + ".bin")
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path):
+            return True
+
+        return False
+
+    def save(self):
+        """save the graph list and the labels"""
+        labels = {}
+        for i, label in enumerate(self.labels):
+            labels["{}".format(i)] = F.tensor(label)
+        save_graphs(str(self.graph_path), self.graphs, labels)
+        if self.verbose:
+            print("Done saving data into cached files.")
+
+    def load(self):
+        self.graphs, labels = load_graphs(str(self.graph_path))
+
+        self.labels = []
+        for i in range(len(labels)):
+            self.labels.append(F.asnumpy(labels["{}".format(i)]))
+        # load pmpds under self.raw_path
+        with open("{}/pmpds.pkl".format(self.raw_path), "rb") as f:
+            self.pmpds = _pickle_load(f)
+        if self.verbose:
+            print("Done loading data into cached files.")
+        # sanity check
+        assert len(self.graphs) == len(self.pmpds)
+        assert len(self.graphs) == len(self.labels)
+
+    def __len__(self):
+        return len(self.graphs)
+
+    def __getitem__(self, i):
+        r"""Gets the idx-th sample.
+
+        Parameters
+        -----------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (dgl.DGLGraph, scipy.sparse.coo_matrix, int)
+            The graph, scipy sparse coo_matrix and its label.
+        """
+        if self._transform is None:
+            g = self.graphs[i]
+        else:
+            g = self._transform(self.graphs[i])
+        return (g, self.pmpds[i], self.labels[i])
+
+    @property
+    def save_name(self):
+        return self.name + "_dgl_graph"
+
+    @staticmethod
+    def collate_fn(cur):
+        graphs, pmpds, labels = zip(*cur)
+        batched_graphs = batch_graphs(graphs)
+        batched_pmpds = sp.block_diag(pmpds)
+        batched_labels = np.concatenate(labels, axis=0)
+        return batched_graphs, batched_pmpds, batched_labels
+
+
+def _normalize(mx):
+    """Row-normalize sparse matrix"""
+    rowsum = np.asarray(mx.sum(1))
+    mask = np.equal(rowsum, 0.0).flatten()
+    rowsum[mask] = np.nan
+    r_inv = np.power(rowsum, -1).flatten()
+    r_inv[mask] = 0.0
+    r_mat_inv = sp.diags(r_inv)
+    return r_mat_inv.dot(mx)
+
+
+def _encode_onehot(labels):
+    classes = list(sorted(set(labels)))
+    classes_dict = {
+        c: np.identity(len(classes))[i, :] for i, c in enumerate(classes)
+    }
+    labels_onehot = np.asarray(
+        list(map(classes_dict.get, labels)), dtype=np.int32
+    )
+    return labels_onehot

python/dgl/data/cluster.py

@@ -0,0 +1,132 @@
+""" CLUSTERDataset for inductive learning. """
+import os
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs
+
+
+class CLUSTERDataset(DGLBuiltinDataset):
+    r"""CLUSTER dataset for semi-supervised clustering task.
+
+    Each graph contains 6 SBM clusters with sizes randomly selected between
+    [5, 35] and probabilities p = 0.55, q = 0.25. The graphs are of sizes 40
+    -190 nodes. Each node can take an input feature value in {0, 1, 2, ..., 6}
+    and values 1~6 correspond to classes 0~5 respectively, while value 0 means
+    that the class of the node is unknown. There is only one labeled node that
+    is randomly assigned to each community and most node features are set to 0.
+
+    Reference `<https://arxiv.org/pdf/2003.00982.pdf>`_
+
+    Statistics:
+
+    - Train examples: 10,000
+    - Valid examples: 1,000
+    - Test examples: 1,000
+    - Number of classes for each node: 6
+
+    Parameters
+    ----------
+    mode : str
+        Must be one of ('train', 'valid', 'test').
+        Default: 'train'
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> from dgl.data import CLUSTERDataset
+    >>>
+    >>> trainset = CLUSTERDataset(mode='train')
+    >>>
+    >>> trainset.num_classes
+    6
+    >>> len(trainset)
+    10000
+    >>> trainset[0]
+    Graph(num_nodes=117, num_edges=4104,
+          ndata_schemes={'label': Scheme(shape=(), dtype=torch.int16),
+                         'feat': Scheme(shape=(), dtype=torch.int64)}
+          edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)})
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self._url = _get_dgl_url("dataset/SBM_CLUSTER.zip")
+        self.mode = mode
+
+        super(CLUSTERDataset, self).__init__(
+            name="cluster",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.load()
+
+    def has_cache(self):
+        graph_path = os.path.join(
+            self.save_path, "CLUSTER_{}.bin".format(self.mode)
+        )
+        return os.path.exists(graph_path)
+
+    def load(self):
+        graph_path = os.path.join(
+            self.save_path, "CLUSTER_{}.bin".format(self.mode)
+        )
+        self._graphs, _ = load_graphs(graph_path)
+
+    @property
+    def num_classes(self):
+        r"""Number of classes for each node."""
+        return 6
+
+    def __len__(self):
+        r"""The number of examples in the dataset."""
+        return len(self._graphs)
+
+    def __getitem__(self, idx):
+        r"""Get the idx^th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node features, node labels and edge features.
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+            - ``edata['feat']``: edge features
+        """
+        if self._transform is None:
+            return self._graphs[idx]
+        else:
+            return self._transform(self._graphs[idx])

python/dgl/data/csv_dataset.py

@@ -0,0 +1,214 @@
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..base import DGLError
+from .dgl_dataset import DGLDataset
+from .utils import load_graphs, save_graphs, Subset
+
+
+class CSVDataset(DGLDataset):
+    """Dataset class that loads and parses graph data from CSV files.
+
+    This class requires the following additional packages:
+
+        - pyyaml >= 5.4.1
+        - pandas >= 1.1.5
+        - pydantic >= 1.9.0
+
+    The parsed graph and feature data will be cached for faster reloading. If
+    the source CSV files are modified, please specify ``force_reload=True``
+    to re-parse from them.
+
+    Parameters
+    ----------
+    data_path : str
+        Directory which contains 'meta.yaml' and CSV files
+    force_reload : bool, optional
+        Whether to reload the dataset. Default: False
+    verbose: bool, optional
+        Whether to print out progress information. Default: True.
+    ndata_parser : dict[str, callable] or callable, optional
+        Callable object which takes in the ``pandas.DataFrame`` object created from
+        CSV file, parses node data and returns a dictionary of parsed data. If given a
+        dictionary, the key is node type and the value is a callable object which is
+        used to parse data of corresponding node type. If given a single callable
+        object, such object is used to parse data of all node type data. Default: None.
+        If None, a default data parser is applied which load data directly and tries to
+        convert list into array.
+    edata_parser : dict[(str, str, str), callable], or callable, optional
+        Callable object which takes in the ``pandas.DataFrame`` object created from
+        CSV file, parses edge data and returns a dictionary of parsed data. If given a
+        dictionary, the key is edge type and the value is a callable object which is
+        used to parse data of corresponding edge type. If given a single callable
+        object, such object is used to parse data of all edge type data. Default: None.
+        If None, a default data parser is applied which load data directly and tries to
+        convert list into array.
+    gdata_parser : callable, optional
+        Callable object which takes in the ``pandas.DataFrame`` object created from
+        CSV file, parses graph data and returns a dictionary of parsed data. Default:
+        None. If None, a default data parser is applied which load data directly and
+        tries to convert list into array.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    graphs : :class:`dgl.DGLGraph`
+        Graphs of the dataset
+    data : dict
+        any available graph-level data such as graph-level feature, labels.
+
+    Examples
+    --------
+    Please refer to :ref:`guide-data-pipeline-loadcsv`.
+
+    """
+
+    META_YAML_NAME = "meta.yaml"
+
+    def __init__(
+        self,
+        data_path,
+        force_reload=False,
+        verbose=True,
+        ndata_parser=None,
+        edata_parser=None,
+        gdata_parser=None,
+        transform=None,
+    ):
+        from .csv_dataset_base import (
+            DefaultDataParser,
+            load_yaml_with_sanity_check,
+        )
+
+        self.graphs = None
+        self.data = None
+        self.ndata_parser = {} if ndata_parser is None else ndata_parser
+        self.edata_parser = {} if edata_parser is None else edata_parser
+        self.gdata_parser = gdata_parser
+        self.default_data_parser = DefaultDataParser()
+        meta_yaml_path = os.path.join(data_path, CSVDataset.META_YAML_NAME)
+        if not os.path.exists(meta_yaml_path):
+            raise DGLError(
+                "'{}' cannot be found under {}.".format(
+                    CSVDataset.META_YAML_NAME, data_path
+                )
+            )
+        self.meta_yaml = load_yaml_with_sanity_check(meta_yaml_path)
+        ds_name = self.meta_yaml.dataset_name
+        super().__init__(
+            ds_name,
+            raw_dir=os.path.dirname(meta_yaml_path),
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """Parse node/edge data from CSV files and construct DGL.Graphs"""
+        from .csv_dataset_base import (
+            DGLGraphConstructor,
+            EdgeData,
+            GraphData,
+            NodeData,
+        )
+
+        meta_yaml = self.meta_yaml
+        base_dir = self.raw_dir
+        node_data = []
+        for meta_node in meta_yaml.node_data:
+            if meta_node is None:
+                continue
+            ntype = meta_node.ntype
+            data_parser = (
+                self.ndata_parser
+                if callable(self.ndata_parser)
+                else self.ndata_parser.get(ntype, self.default_data_parser)
+            )
+            ndata = NodeData.load_from_csv(
+                meta_node,
+                base_dir=base_dir,
+                separator=meta_yaml.separator,
+                data_parser=data_parser,
+            )
+            node_data.append(ndata)
+        edge_data = []
+        for meta_edge in meta_yaml.edge_data:
+            if meta_edge is None:
+                continue
+            etype = tuple(meta_edge.etype)
+            data_parser = (
+                self.edata_parser
+                if callable(self.edata_parser)
+                else self.edata_parser.get(etype, self.default_data_parser)
+            )
+            edata = EdgeData.load_from_csv(
+                meta_edge,
+                base_dir=base_dir,
+                separator=meta_yaml.separator,
+                data_parser=data_parser,
+            )
+            edge_data.append(edata)
+        graph_data = None
+        if meta_yaml.graph_data is not None:
+            meta_graph = meta_yaml.graph_data
+            data_parser = (
+                self.default_data_parser
+                if self.gdata_parser is None
+                else self.gdata_parser
+            )
+            graph_data = GraphData.load_from_csv(
+                meta_graph,
+                base_dir=base_dir,
+                separator=meta_yaml.separator,
+                data_parser=data_parser,
+            )
+        # construct graphs
+        self.graphs, self.data = DGLGraphConstructor.construct_graphs(
+            node_data, edge_data, graph_data
+        )
+        if len(self.data) == 1:
+            self.labels = list(self.data.values())[0]
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, self.name + ".bin")
+        if os.path.exists(graph_path):
+            return True
+
+        return False
+
+    def save(self):
+        if self.graphs is None:
+            raise DGLError("No graphs available in dataset")
+        graph_path = os.path.join(self.save_path, self.name + ".bin")
+        save_graphs(graph_path, self.graphs, labels=self.data)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, self.name + ".bin")
+        self.graphs, self.data = load_graphs(graph_path)
+        if len(self.data) == 1:
+            self.labels = list(self.data.values())[0]
+
+    def __getitem__(self, i):
+        if F.is_tensor(i) and F.ndim(i) == 1:
+            return Subset(self, F.copy_to(i, F.cpu()))
+
+        if self._transform is None:
+            g = self.graphs[i]
+        else:
+            g = self._transform(self.graphs[i])
+
+        if len(self.data) == 1:
+            return g, self.labels[i]
+        elif len(self.data) > 0:
+            data = {k: v[i] for (k, v) in self.data.items()}
+            return g, data
+        else:
+            return g
+
+    def __len__(self):
+        return len(self.graphs)

python/dgl/data/csv_dataset_base.py

@@ -0,0 +1,386 @@
+import ast
+import os
+from typing import Callable, List, Optional
+
+import numpy as np
+import pandas as pd
+import pydantic as dt
+import yaml
+
+from .. import backend as F
+from ..base import dgl_warning, DGLError
+from ..convert import heterograph as dgl_heterograph
+
+
+class MetaNode(dt.BaseModel):
+    """Class of node_data in YAML. Internal use only."""
+
+    file_name: str
+    ntype: Optional[str] = "_V"
+    graph_id_field: Optional[str] = "graph_id"
+    node_id_field: Optional[str] = "node_id"
+
+
+class MetaEdge(dt.BaseModel):
+    """Class of edge_data in YAML. Internal use only."""
+
+    file_name: str
+    etype: Optional[List[str]] = ["_V", "_E", "_V"]
+    graph_id_field: Optional[str] = "graph_id"
+    src_id_field: Optional[str] = "src_id"
+    dst_id_field: Optional[str] = "dst_id"
+
+
+class MetaGraph(dt.BaseModel):
+    """Class of graph_data in YAML. Internal use only."""
+
+    file_name: str
+    graph_id_field: Optional[str] = "graph_id"
+
+
+class MetaYaml(dt.BaseModel):
+    """Class of YAML. Internal use only."""
+
+    version: Optional[str] = "1.0.0"
+    dataset_name: str
+    separator: Optional[str] = ","
+    node_data: List[MetaNode]
+    edge_data: List[MetaEdge]
+    graph_data: Optional[MetaGraph] = None
+
+
+def load_yaml_with_sanity_check(yaml_file):
+    """Load yaml and do sanity check. Internal use only."""
+    with open(yaml_file) as f:
+        yaml_data = yaml.load(f, Loader=yaml.loader.SafeLoader)
+        try:
+            meta_yaml = MetaYaml(**yaml_data)
+        except dt.ValidationError as e:
+            print("Details of pydantic.ValidationError:\n{}".format(e.json()))
+            raise DGLError(
+                "Validation Error for YAML fields. Details are shown above."
+            )
+        if meta_yaml.version != "1.0.0":
+            raise DGLError(
+                "Invalid CSVDataset version {}. Supported versions: '1.0.0'".format(
+                    meta_yaml.version
+                )
+            )
+        ntypes = [meta.ntype for meta in meta_yaml.node_data]
+        if len(ntypes) > len(set(ntypes)):
+            raise DGLError(
+                "Each node CSV file must have a unique node type name, but found duplicate node type: {}.".format(
+                    ntypes
+                )
+            )
+        etypes = [tuple(meta.etype) for meta in meta_yaml.edge_data]
+        if len(etypes) > len(set(etypes)):
+            raise DGLError(
+                "Each edge CSV file must have a unique edge type name, but found duplicate edge type: {}.".format(
+                    etypes
+                )
+            )
+        return meta_yaml
+
+
+def _validate_data_length(data_dict):
+    len_dict = {k: len(v) for k, v in data_dict.items()}
+    lst = list(len_dict.values())
+    res = lst.count(lst[0]) == len(lst)
+    if not res:
+        raise DGLError(
+            "All data are required to have same length while some of them does not. Length of data={}".format(
+                str(len_dict)
+            )
+        )
+
+
+def _tensor(data, dtype=None):
+    """Float32 is the default dtype for float tensor in DGL
+    so let's cast float64 into float32 to avoid dtype mismatch.
+    """
+    ret = F.tensor(data, dtype)
+    if F.dtype(ret) == F.float64:
+        ret = F.tensor(ret, dtype=F.float32)
+    return ret
+
+
+class BaseData:
+    """Class of base data which is inherited by Node/Edge/GraphData. Internal use only."""
+
+    @staticmethod
+    def read_csv(file_name, base_dir, separator):
+        csv_path = file_name
+        if base_dir is not None:
+            csv_path = os.path.join(base_dir, csv_path)
+        return pd.read_csv(csv_path, sep=separator)
+
+    @staticmethod
+    def pop_from_dataframe(df: pd.DataFrame, item: str):
+        ret = None
+        try:
+            ret = df.pop(item).to_numpy().squeeze()
+        except KeyError:
+            pass
+        return ret
+
+
+class NodeData(BaseData):
+    """Class of node data which is used for DGLGraph construction. Internal use only."""
+
+    def __init__(self, node_id, data, type=None, graph_id=None):
+        self.id = np.array(node_id)
+        self.data = data
+        self.type = type if type is not None else "_V"
+        self.graph_id = (
+            np.array(graph_id)
+            if graph_id is not None
+            else np.full(len(node_id), 0)
+        )
+        _validate_data_length(
+            {**{"id": self.id, "graph_id": self.graph_id}, **self.data}
+        )
+
+    @staticmethod
+    def load_from_csv(
+        meta: MetaNode, data_parser: Callable, base_dir=None, separator=","
+    ):
+        df = BaseData.read_csv(meta.file_name, base_dir, separator)
+        node_ids = BaseData.pop_from_dataframe(df, meta.node_id_field)
+        graph_ids = BaseData.pop_from_dataframe(df, meta.graph_id_field)
+        if node_ids is None:
+            raise DGLError(
+                "Missing node id field [{}] in file [{}].".format(
+                    meta.node_id_field, meta.file_name
+                )
+            )
+        ntype = meta.ntype
+        ndata = data_parser(df)
+        return NodeData(node_ids, ndata, type=ntype, graph_id=graph_ids)
+
+    @staticmethod
+    def to_dict(node_data: List["NodeData"]) -> dict:
+        # node_ids could be numeric or non-numeric values, but duplication is not allowed.
+        node_dict = {}
+        for n_data in node_data:
+            graph_ids = np.unique(n_data.graph_id)
+            for graph_id in graph_ids:
+                idx = n_data.graph_id == graph_id
+                ids = n_data.id[idx]
+                u_ids, u_indices, u_counts = np.unique(
+                    ids, return_index=True, return_counts=True
+                )
+                if len(ids) > len(u_ids):
+                    raise DGLError(
+                        "Node IDs are required to be unique but the following ids are duplicate: {}".format(
+                            u_ids[u_counts > 1]
+                        )
+                    )
+                if graph_id not in node_dict:
+                    node_dict[graph_id] = {}
+                node_dict[graph_id][n_data.type] = {
+                    "mapping": {
+                        index: i for i, index in enumerate(ids[u_indices])
+                    },
+                    "data": {
+                        k: _tensor(v[idx][u_indices])
+                        for k, v in n_data.data.items()
+                    },
+                    "dtype": ids.dtype,
+                }
+        return node_dict
+
+
+class EdgeData(BaseData):
+    """Class of edge data which is used for DGLGraph construction. Internal use only."""
+
+    def __init__(self, src_id, dst_id, data, type=None, graph_id=None):
+        self.src = np.array(src_id)
+        self.dst = np.array(dst_id)
+        self.data = data
+        self.type = type if type is not None else ("_V", "_E", "_V")
+        self.graph_id = (
+            np.array(graph_id)
+            if graph_id is not None
+            else np.full(len(src_id), 0)
+        )
+        _validate_data_length(
+            {
+                **{"src": self.src, "dst": self.dst, "graph_id": self.graph_id},
+                **self.data,
+            }
+        )
+
+    @staticmethod
+    def load_from_csv(
+        meta: MetaEdge, data_parser: Callable, base_dir=None, separator=","
+    ):
+        df = BaseData.read_csv(meta.file_name, base_dir, separator)
+        src_ids = BaseData.pop_from_dataframe(df, meta.src_id_field)
+        if src_ids is None:
+            raise DGLError(
+                "Missing src id field [{}] in file [{}].".format(
+                    meta.src_id_field, meta.file_name
+                )
+            )
+        dst_ids = BaseData.pop_from_dataframe(df, meta.dst_id_field)
+        if dst_ids is None:
+            raise DGLError(
+                "Missing dst id field [{}] in file [{}].".format(
+                    meta.dst_id_field, meta.file_name
+                )
+            )
+        graph_ids = BaseData.pop_from_dataframe(df, meta.graph_id_field)
+        etype = tuple(meta.etype)
+        edata = data_parser(df)
+        return EdgeData(src_ids, dst_ids, edata, type=etype, graph_id=graph_ids)
+
+    @staticmethod
+    def to_dict(edge_data: List["EdgeData"], node_dict: dict) -> dict:
+        edge_dict = {}
+        for e_data in edge_data:
+            (src_type, e_type, dst_type) = e_data.type
+            graph_ids = np.unique(e_data.graph_id)
+            for graph_id in graph_ids:
+                if graph_id in edge_dict and e_data.type in edge_dict[graph_id]:
+                    raise DGLError(
+                        f"Duplicate edge type[{e_data.type}] for same graph[{graph_id}], please place the same edge_type for same graph into single EdgeData."
+                    )
+                idx = e_data.graph_id == graph_id
+                src_mapping = node_dict[graph_id][src_type]["mapping"]
+                dst_mapping = node_dict[graph_id][dst_type]["mapping"]
+                orig_src_ids = e_data.src[idx].astype(
+                    node_dict[graph_id][src_type]["dtype"]
+                )
+                orig_dst_ids = e_data.dst[idx].astype(
+                    node_dict[graph_id][dst_type]["dtype"]
+                )
+                src_ids = [src_mapping[index] for index in orig_src_ids]
+                dst_ids = [dst_mapping[index] for index in orig_dst_ids]
+                if graph_id not in edge_dict:
+                    edge_dict[graph_id] = {}
+                edge_dict[graph_id][e_data.type] = {
+                    "edges": (_tensor(src_ids), _tensor(dst_ids)),
+                    "data": {
+                        k: _tensor(v[idx]) for k, v in e_data.data.items()
+                    },
+                }
+        return edge_dict
+
+
+class GraphData(BaseData):
+    """Class of graph data which is used for DGLGraph construction. Internal use only."""
+
+    def __init__(self, graph_id, data):
+        self.graph_id = np.array(graph_id)
+        self.data = data
+        _validate_data_length({**{"graph_id": self.graph_id}, **self.data})
+
+    @staticmethod
+    def load_from_csv(
+        meta: MetaGraph, data_parser: Callable, base_dir=None, separator=","
+    ):
+        df = BaseData.read_csv(meta.file_name, base_dir, separator)
+        graph_ids = BaseData.pop_from_dataframe(df, meta.graph_id_field)
+        if graph_ids is None:
+            raise DGLError(
+                "Missing graph id field [{}] in file [{}].".format(
+                    meta.graph_id_field, meta.file_name
+                )
+            )
+        gdata = data_parser(df)
+        return GraphData(graph_ids, gdata)
+
+    @staticmethod
+    def to_dict(graph_data: "GraphData", graphs_dict: dict) -> dict:
+        missing_ids = np.setdiff1d(
+            np.array(list(graphs_dict.keys())), graph_data.graph_id
+        )
+        if len(missing_ids) > 0:
+            raise DGLError(
+                "Found following graph ids in node/edge CSVs but not in graph CSV: {}.".format(
+                    missing_ids
+                )
+            )
+        graph_ids = graph_data.graph_id
+        graphs = []
+        for graph_id in graph_ids:
+            if graph_id not in graphs_dict:
+                graphs_dict[graph_id] = dgl_heterograph(
+                    {("_V", "_E", "_V"): ([], [])}
+                )
+        for graph_id in graph_ids:
+            graphs.append(graphs_dict[graph_id])
+        data = {
+            k: F.reshape(_tensor(v), (len(graphs), -1))
+            for k, v in graph_data.data.items()
+        }
+        return graphs, data
+
+
+class DGLGraphConstructor:
+    """Class for constructing DGLGraph from Node/Edge/Graph data. Internal use only."""
+
+    @staticmethod
+    def construct_graphs(node_data, edge_data, graph_data=None):
+        if not isinstance(node_data, list):
+            node_data = [node_data]
+        if not isinstance(edge_data, list):
+            edge_data = [edge_data]
+        node_dict = NodeData.to_dict(node_data)
+        edge_dict = EdgeData.to_dict(edge_data, node_dict)
+        graph_dict = DGLGraphConstructor._construct_graphs(node_dict, edge_dict)
+        if graph_data is None:
+            graph_data = GraphData(np.full(1, 0), {})
+        graphs, data = GraphData.to_dict(graph_data, graph_dict)
+        return graphs, data
+
+    @staticmethod
+    def _construct_graphs(node_dict, edge_dict):
+        graph_dict = {}
+        for graph_id in node_dict:
+            if graph_id not in edge_dict:
+                edge_dict[graph_id][("_V", "_E", "_V")] = {"edges": ([], [])}
+            graph = dgl_heterograph(
+                {
+                    etype: edata["edges"]
+                    for etype, edata in edge_dict[graph_id].items()
+                },
+                num_nodes_dict={
+                    ntype: len(ndata["mapping"])
+                    for ntype, ndata in node_dict[graph_id].items()
+                },
+            )
+
+            def assign_data(type, src_data, dst_data):
+                for key, value in src_data.items():
+                    dst_data[type].data[key] = value
+
+            for type, data in node_dict[graph_id].items():
+                assign_data(type, data["data"], graph.nodes)
+            for (type), data in edge_dict[graph_id].items():
+                assign_data(type, data["data"], graph.edges)
+            graph_dict[graph_id] = graph
+        return graph_dict
+
+
+class DefaultDataParser:
+    """Default data parser for CSVDataset. It
+    1. ignores any columns which does not have a header.
+    2. tries to convert to list of numeric values(generated by
+        np.array().tolist()) if cell data is a str separated by ','.
+    3. read data and infer data type directly, otherwise.
+    """
+
+    def __call__(self, df: pd.DataFrame):
+        data = {}
+        for header in df:
+            if "Unnamed" in header:
+                dgl_warning("Unnamed column is found. Ignored...")
+                continue
+            dt = df[header].to_numpy().squeeze()
+            if len(dt) > 0 and isinstance(dt[0], str):
+                # probably consists of list of numeric values
+                dt = np.array([ast.literal_eval(row) for row in dt])
+            data[header] = dt
+        return data

python/dgl/data/dgl_dataset.py

@@ -0,0 +1,349 @@
+"""Basic DGL Dataset
+"""
+
+from __future__ import absolute_import
+
+import abc
+import hashlib
+import os
+import traceback
+
+from ..utils import retry_method_with_fix
+from .utils import download, extract_archive, get_download_dir, makedirs
+
+
+class DGLDataset(object):
+    r"""The basic DGL dataset for creating graph datasets.
+    This class defines a basic template class for DGL Dataset.
+    The following steps will be executed automatically:
+
+      1. Check whether there is a dataset cache on disk
+         (already processed and stored on the disk) by
+         invoking ``has_cache()``. If true, goto 5.
+      2. Call ``download()`` to download the data if ``url`` is not None.
+      3. Call ``process()`` to process the data.
+      4. Call ``save()`` to save the processed dataset on disk and goto 6.
+      5. Call ``load()`` to load the processed dataset from disk.
+      6. Done.
+
+    Users can overwite these functions with their
+    own data processing logic.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset
+    url : str
+        Url to download the raw dataset. Default: None
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    save_dir : str
+        Directory to save the processed dataset.
+        Default: same as raw_dir
+    hash_key : tuple
+        A tuple of values as the input for the hash function.
+        Users can distinguish instances (and their caches on the disk)
+        from the same dataset class by comparing the hash values.
+        Default: (), the corresponding hash value is ``'f9065fa7'``.
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    url : str
+        The URL to download the dataset
+    name : str
+        The dataset name
+    raw_dir : str
+        Directory to store all the downloaded raw datasets.
+    raw_path : str
+        Path to the downloaded raw dataset folder. An alias for
+        ``os.path.join(self.raw_dir, self.name)``.
+    save_dir : str
+        Directory to save all the processed datasets.
+    save_path : str
+        Path to the processed dataset folder. An alias for
+        ``os.path.join(self.save_dir, self.name)``.
+    verbose : bool
+        Whether to print more runtime information.
+    hash : str
+        Hash value for the dataset and the setting.
+    """
+
+    def __init__(
+        self,
+        name,
+        url=None,
+        raw_dir=None,
+        save_dir=None,
+        hash_key=(),
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self._name = name
+        self._url = url
+        self._force_reload = force_reload
+        self._verbose = verbose
+        self._hash_key = hash_key
+        self._hash = self._get_hash()
+        self._transform = transform
+
+        # if no dir is provided, the default dgl download dir is used.
+        if raw_dir is None:
+            self._raw_dir = get_download_dir()
+        else:
+            self._raw_dir = raw_dir
+
+        if save_dir is None:
+            self._save_dir = self._raw_dir
+        else:
+            self._save_dir = save_dir
+
+        self._load()
+
+    def download(self):
+        r"""Overwite to realize your own logic of downloading data.
+
+        It is recommended to download the to the :obj:`self.raw_dir`
+        folder. Can be ignored if the dataset is
+        already in :obj:`self.raw_dir`.
+        """
+        pass
+
+    def save(self):
+        r"""Overwite to realize your own logic of
+        saving the processed dataset into files.
+
+        It is recommended to use ``dgl.data.utils.save_graphs``
+        to save dgl graph into files and use
+        ``dgl.data.utils.save_info`` to save extra
+        information into files.
+        """
+        pass
+
+    def load(self):
+        r"""Overwite to realize your own logic of
+        loading the saved dataset from files.
+
+        It is recommended to use ``dgl.data.utils.load_graphs``
+        to load dgl graph from files and use
+        ``dgl.data.utils.load_info`` to load extra information
+        into python dict object.
+        """
+        pass
+
+    @abc.abstractmethod
+    def process(self):
+        r"""Overwrite to realize your own logic of processing the input data."""
+        pass
+
+    def has_cache(self):
+        r"""Overwrite to realize your own logic of
+        deciding whether there exists a cached dataset.
+
+        By default False.
+        """
+        return False
+
+    @retry_method_with_fix(download)
+    def _download(self):
+        """Download dataset by calling ``self.download()``
+        if the dataset does not exists under ``self.raw_path``.
+
+        By default ``self.raw_path = os.path.join(self.raw_dir, self.name)``
+        One can overwrite ``raw_path()`` function to change the path.
+        """
+        if os.path.exists(self.raw_path):  # pragma: no cover
+            return
+
+        makedirs(self.raw_dir)
+        self.download()
+
+    def _load(self):
+        """Entry point from __init__ to load the dataset.
+
+        If cache exists:
+
+          - Load the dataset from saved dgl graph and information files.
+          - If loadin process fails, re-download and process the dataset.
+
+        else:
+
+          - Download the dataset if needed.
+          - Process the dataset and build the dgl graph.
+          - Save the processed dataset into files.
+        """
+        load_flag = not self._force_reload and self.has_cache()
+
+        if load_flag:
+            try:
+                self.load()
+                if self.verbose:
+                    print("Done loading data from cached files.")
+            except KeyboardInterrupt:
+                raise
+            except:
+                load_flag = False
+                if self.verbose:
+                    print(traceback.format_exc())
+                    print("Loading from cache failed, re-processing.")
+
+        if not load_flag:
+            self._download()
+            self.process()
+            self.save()
+            if self.verbose:
+                print("Done saving data into cached files.")
+
+    def _get_hash(self):
+        """Compute the hash of the input tuple
+
+        Example
+        -------
+        Assume `self._hash_key = (10, False, True)`
+
+        >>> hash_value = self._get_hash()
+        >>> hash_value
+        'a770b222'
+        """
+        hash_func = hashlib.sha1()
+        hash_func.update(str(self._hash_key).encode("utf-8"))
+        return hash_func.hexdigest()[:8]
+
+    def _get_hash_url_suffix(self):
+        """Get the suffix based on the hash value of the url."""
+        if self._url is None:
+            return ""
+        else:
+            hash_func = hashlib.sha1()
+            hash_func.update(str(self._url).encode("utf-8"))
+            return "_" + hash_func.hexdigest()[:8]
+
+    @property
+    def url(self):
+        r"""Get url to download the raw dataset."""
+        return self._url
+
+    @property
+    def name(self):
+        r"""Name of the dataset."""
+        return self._name
+
+    @property
+    def raw_dir(self):
+        r"""Raw file directory contains the input data folder."""
+        return self._raw_dir
+
+    @property
+    def raw_path(self):
+        r"""Directory contains the input data files.
+        By default raw_path = os.path.join(self.raw_dir, self.name)
+        """
+        return os.path.join(
+            self.raw_dir, self.name + self._get_hash_url_suffix()
+        )
+
+    @property
+    def save_dir(self):
+        r"""Directory to save the processed dataset."""
+        return self._save_dir
+
+    @property
+    def save_path(self):
+        r"""Path to save the processed dataset."""
+        return os.path.join(
+            self.save_dir, self.name + self._get_hash_url_suffix()
+        )
+
+    @property
+    def verbose(self):
+        r"""Whether to print information."""
+        return self._verbose
+
+    @property
+    def hash(self):
+        r"""Hash value for the dataset and the setting."""
+        return self._hash
+
+    @abc.abstractmethod
+    def __getitem__(self, idx):
+        r"""Gets the data object at index."""
+        pass
+
+    @abc.abstractmethod
+    def __len__(self):
+        r"""The number of examples in the dataset."""
+        pass
+
+    def __repr__(self):
+        return (
+            f'Dataset("{self.name}", num_graphs={len(self)},'
+            + f" save_path={self.save_path})"
+        )
+
+
+class DGLBuiltinDataset(DGLDataset):
+    r"""The Basic DGL Builtin Dataset.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset.
+    url : str
+        Url to download the raw dataset.
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    hash_key : tuple
+        A tuple of values as the input for the hash function.
+        Users can distinguish instances (and their caches on the disk)
+        from the same dataset class by comparing the hash values.
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(
+        self,
+        name,
+        url,
+        raw_dir=None,
+        hash_key=(),
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        super(DGLBuiltinDataset, self).__init__(
+            name,
+            url=url,
+            raw_dir=raw_dir,
+            save_dir=None,
+            hash_key=hash_key,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        r"""Automatically download data and extract it."""
+        if self.url is not None:
+            zip_file_path = os.path.join(self.raw_dir, self.name + ".zip")
+            download(self.url, path=zip_file_path)
+            extract_archive(zip_file_path, self.raw_path)

python/dgl/data/fakenews.py

@@ -0,0 +1,255 @@
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs, load_info, save_graphs, save_info
+
+
+class FakeNewsDataset(DGLBuiltinDataset):
+    r"""Fake News Graph Classification dataset.
+
+    The dataset is composed of two sets of tree-structured fake/real
+    news propagation graphs extracted from Twitter. Different from
+    most of the benchmark datasets for the graph classification task,
+    the graphs in this dataset are directed tree-structured graphs where
+    the root node represents the news, the leaf nodes are Twitter users
+    who retweeted the root news. Besides, the node features are encoded
+    user historical tweets using different pretrained language models:
+
+    - bert: the 768-dimensional node feature composed of Twitter user historical tweets encoded by the bert-as-service
+    - content: the 310-dimensional node feature composed of a 300-dimensional “spacy” vector plus a 10-dimensional “profile” vector
+    - profile: the 10-dimensional node feature composed of ten Twitter user profile attributes.
+    - spacy: the 300-dimensional node feature composed of Twitter user historical tweets encoded by the spaCy word2vec encoder.
+
+    Reference: <https://github.com/safe-graph/GNN-FakeNews>
+
+    Note: this dataset is for academic use only, and commercial use is prohibited.
+
+    Statistics:
+
+        Politifact:
+
+        - Graphs: 314
+        - Nodes: 41,054
+        - Edges: 40,740
+        - Classes:
+
+            - Fake: 157
+            - Real: 157
+
+        - Node feature size:
+
+            - bert: 768
+            - content: 310
+            - profile: 10
+            - spacy: 300
+
+        Gossipcop:
+
+        - Graphs: 5,464
+        - Nodes: 314,262
+        - Edges: 308,798
+        - Classes:
+
+            - Fake: 2,732
+            - Real: 2,732
+
+        - Node feature size:
+
+            - bert: 768
+            - content: 310
+            - profile: 10
+            - spacy: 300
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset (gossipcop, or politifact)
+    feature_name : str
+        Name of the feature (bert, content, profile, or spacy)
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    name : str
+        Name of the dataset (gossipcop, or politifact)
+    num_classes : int
+        Number of label classes
+    num_graphs : int
+        Number of graphs
+    graphs : list
+        A list of DGLGraph objects
+    labels : Tensor
+        Graph labels
+    feature_name : str
+        Name of the feature (bert, content, profile, or spacy)
+    feature : Tensor
+        Node features
+    train_mask : Tensor
+        Mask of training set
+    val_mask : Tensor
+        Mask of validation set
+    test_mask : Tensor
+        Mask of testing set
+
+    Examples
+    --------
+    >>> dataset = FakeNewsDataset('gossipcop', 'bert')
+    >>> graph, label = dataset[0]
+    >>> num_classes = dataset.num_classes
+    >>> feat = dataset.feature
+    >>> labels = dataset.labels
+    """
+    file_urls = {
+        "gossipcop": "dataset/FakeNewsGOS.zip",
+        "politifact": "dataset/FakeNewsPOL.zip",
+    }
+
+    def __init__(self, name, feature_name, raw_dir=None, transform=None):
+        assert name in [
+            "gossipcop",
+            "politifact",
+        ], "Only supports 'gossipcop' or 'politifact'."
+        url = _get_dgl_url(self.file_urls[name])
+
+        assert feature_name in [
+            "bert",
+            "content",
+            "profile",
+            "spacy",
+        ], "Only supports 'bert', 'content', 'profile', or 'spacy'"
+        self.feature_name = feature_name
+        super(FakeNewsDataset, self).__init__(
+            name=name, url=url, raw_dir=raw_dir, transform=transform
+        )
+
+    def process(self):
+        """process raw data to graph, labels and masks"""
+        self.labels = F.tensor(
+            np.load(os.path.join(self.raw_path, "graph_labels.npy"))
+        )
+        num_graphs = self.labels.shape[0]
+
+        node_graph_id = np.load(
+            os.path.join(self.raw_path, "node_graph_id.npy")
+        )
+        edges = np.genfromtxt(
+            os.path.join(self.raw_path, "A.txt"), delimiter=",", dtype=int
+        )
+        src = edges[:, 0]
+        dst = edges[:, 1]
+        g = graph((src, dst))
+
+        node_idx_list = []
+        for idx in range(np.max(node_graph_id) + 1):
+            node_idx = np.where(node_graph_id == idx)
+            node_idx_list.append(node_idx[0])
+
+        self.graphs = [g.subgraph(node_idx) for node_idx in node_idx_list]
+
+        train_idx = np.load(os.path.join(self.raw_path, "train_idx.npy"))
+        val_idx = np.load(os.path.join(self.raw_path, "val_idx.npy"))
+        test_idx = np.load(os.path.join(self.raw_path, "test_idx.npy"))
+        train_mask = np.zeros(num_graphs, dtype=np.bool_)
+        val_mask = np.zeros(num_graphs, dtype=np.bool_)
+        test_mask = np.zeros(num_graphs, dtype=np.bool_)
+        train_mask[train_idx] = True
+        val_mask[val_idx] = True
+        test_mask[test_idx] = True
+        self.train_mask = F.tensor(train_mask)
+        self.val_mask = F.tensor(val_mask)
+        self.test_mask = F.tensor(test_mask)
+
+        feature_file = "new_" + self.feature_name + "_feature.npz"
+        self.feature = F.tensor(
+            sp.load_npz(os.path.join(self.raw_path, feature_file)).todense()
+        )
+
+    def save(self):
+        """save the graph list and the labels"""
+        save_graphs(str(self.graph_path), self.graphs)
+        save_info(
+            self.info_path,
+            {
+                "label": self.labels,
+                "feature": self.feature,
+                "train_mask": self.train_mask,
+                "val_mask": self.val_mask,
+                "test_mask": self.test_mask,
+            },
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.name + "_dgl_graph.bin")
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, self.name + "_dgl_graph.pkl")
+
+    def has_cache(self):
+        """check whether there are processed data in `self.save_path`"""
+        return os.path.exists(self.graph_path) and os.path.exists(
+            self.info_path
+        )
+
+    def load(self):
+        """load processed data from directory `self.save_path`"""
+        graphs, _ = load_graphs(str(self.graph_path))
+        info = load_info(str(self.info_path))
+        self.graphs = graphs
+        self.labels = info["label"]
+        self.feature = info["feature"]
+
+        self.train_mask = info["train_mask"]
+        self.val_mask = info["val_mask"]
+        self.test_mask = info["test_mask"]
+
+    @property
+    def num_classes(self):
+        """Number of classes for each graph, i.e. number of prediction tasks."""
+        return 2
+
+    @property
+    def num_graphs(self):
+        """Number of graphs."""
+        return self.labels.shape[0]
+
+    def __getitem__(self, i):
+        r"""Get graph and label by index
+
+        Parameters
+        ----------
+        i : int
+            Item index
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+        """
+        if self._transform is None:
+            g = self.graphs[i]
+        else:
+            g = self._transform(self.graphs[i])
+        return g, self.labels[i]
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return len(self.graphs)

python/dgl/data/flickr.py

@@ -0,0 +1,178 @@
+"""Flickr Dataset"""
+import json
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import from_scipy
+from ..transforms import reorder_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, generate_mask_tensor, load_graphs, save_graphs
+
+
+class FlickrDataset(DGLBuiltinDataset):
+    r"""Flickr dataset for node classification from `GraphSAINT: Graph Sampling Based Inductive
+    Learning Method <https://arxiv.org/abs/1907.04931>`_
+
+    The task of this dataset is categorizing types of images based on the descriptions and common
+    properties of online images.
+
+    Flickr dataset statistics:
+
+    - Nodes: 89,250
+    - Edges: 899,756
+    - Number of classes: 7
+    - Node feature size: 500
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`.
+        Default: False.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+    >>> from dgl.data import FlickrDataset
+    >>> dataset = FlickrDataset()
+    >>> dataset.num_classes
+    7
+    >>> g = dataset[0]
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>> # get node labels
+    >>> labels = g.ndata['label']
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+        reorder=False,
+    ):
+        _url = _get_dgl_url("dataset/flickr.zip")
+        self._reorder = reorder
+        super(FlickrDataset, self).__init__(
+            name="flickr",
+            raw_dir=raw_dir,
+            url=_url,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """process raw data to graph, labels and masks"""
+        coo_adj = sp.load_npz(os.path.join(self.raw_path, "adj_full.npz"))
+        g = from_scipy(coo_adj)
+
+        features = np.load(os.path.join(self.raw_path, "feats.npy"))
+        features = F.tensor(features, dtype=F.float32)
+
+        y = [-1] * features.shape[0]
+        with open(os.path.join(self.raw_path, "class_map.json")) as f:
+            class_map = json.load(f)
+            for key, item in class_map.items():
+                y[int(key)] = item
+        labels = F.tensor(np.array(y), dtype=F.int64)
+
+        with open(os.path.join(self.raw_path, "role.json")) as f:
+            role = json.load(f)
+
+        train_mask = np.zeros(features.shape[0], dtype=bool)
+        train_mask[role["tr"]] = True
+
+        val_mask = np.zeros(features.shape[0], dtype=bool)
+        val_mask[role["va"]] = True
+
+        test_mask = np.zeros(features.shape[0], dtype=bool)
+        test_mask[role["te"]] = True
+
+        g.ndata["feat"] = features
+        g.ndata["label"] = labels
+        g.ndata["train_mask"] = generate_mask_tensor(train_mask)
+        g.ndata["val_mask"] = generate_mask_tensor(val_mask)
+        g.ndata["test_mask"] = generate_mask_tensor(test_mask)
+
+        if self._reorder:
+            self._graph = reorder_graph(
+                g,
+                node_permute_algo="rcmk",
+                edge_permute_algo="dst",
+                store_ids=False,
+            )
+        else:
+            self._graph = g
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        return os.path.exists(graph_path)
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        g, _ = load_graphs(graph_path)
+        self._graph = g[0]
+
+    @property
+    def num_classes(self):
+        return 7
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return 1
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index, FlickrDataset has only one graph object
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``ndata['label']``: node label
+            - ``ndata['feat']``: node feature
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)

python/dgl/data/fraud.py

@@ -0,0 +1,415 @@
+"""Fraud Dataset
+"""
+import os
+
+import numpy as np
+from scipy import io
+
+from .. import backend as F
+from ..convert import heterograph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs, save_graphs
+
+
+class FraudDataset(DGLBuiltinDataset):
+    r"""Fraud node prediction dataset.
+
+    The dataset includes two multi-relational graphs extracted from Yelp and Amazon
+    where nodes represent fraudulent reviews or fraudulent reviewers.
+
+    It was first proposed in a CIKM'20 paper <https://arxiv.org/pdf/2008.08692.pdf> and
+    has been used by a recent WWW'21 paper <https://ponderly.github.io/pub/PCGNN_WWW2021.pdf>
+    as a benchmark. Another paper <https://arxiv.org/pdf/2104.01404.pdf> also takes
+    the dataset as an example to study the non-homophilous graphs. This dataset is built
+    upon industrial data and has rich relational information and unique properties like
+    class-imbalance and feature inconsistency, which makes the dataset be a good instance
+    to investigate how GNNs perform on real-world noisy graphs. These graphs are bidirected
+    and not self connected.
+
+    Reference: <https://github.com/YingtongDou/CARE-GNN>
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    random_seed : int
+        Specifying the random seed in splitting the dataset.
+        Default: 717
+    train_size : float
+        training set size of the dataset.
+        Default: 0.7
+    val_size : float
+        validation set size of the dataset, and the
+        size of testing set is (1 - train_size - val_size)
+        Default: 0.1
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of label classes
+    graph : dgl.DGLGraph
+        Graph structure, etc.
+    seed : int
+        Random seed in splitting the dataset.
+    train_size : float
+        Training set size of the dataset.
+    val_size : float
+        Validation set size of the dataset
+
+    Examples
+    --------
+    >>> dataset = FraudDataset('yelp')
+    >>> graph = dataset[0]
+    >>> num_classes = dataset.num_classes
+    >>> feat = graph.ndata['feature']
+    >>> label = graph.ndata['label']
+    """
+    file_urls = {
+        "yelp": "dataset/FraudYelp.zip",
+        "amazon": "dataset/FraudAmazon.zip",
+    }
+    relations = {
+        "yelp": ["net_rsr", "net_rtr", "net_rur"],
+        "amazon": ["net_upu", "net_usu", "net_uvu"],
+    }
+    file_names = {"yelp": "YelpChi.mat", "amazon": "Amazon.mat"}
+    node_name = {"yelp": "review", "amazon": "user"}
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        random_seed=717,
+        train_size=0.7,
+        val_size=0.1,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        assert name in ["yelp", "amazon"], "only supports 'yelp', or 'amazon'"
+        url = _get_dgl_url(self.file_urls[name])
+        self.seed = random_seed
+        self.train_size = train_size
+        self.val_size = val_size
+        super(FraudDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            hash_key=(random_seed, train_size, val_size),
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """process raw data to graph, labels, splitting masks"""
+        file_path = os.path.join(self.raw_path, self.file_names[self.name])
+
+        data = io.loadmat(file_path)
+        node_features = data["features"].todense()
+        # remove additional dimension of length 1 in raw .mat file
+        node_labels = data["label"].squeeze()
+
+        graph_data = {}
+        for relation in self.relations[self.name]:
+            adj = data[relation].tocoo()
+            row, col = adj.row, adj.col
+            graph_data[
+                (self.node_name[self.name], relation, self.node_name[self.name])
+            ] = (row, col)
+        g = heterograph(graph_data)
+
+        g.ndata["feature"] = F.tensor(
+            node_features, dtype=F.data_type_dict["float32"]
+        )
+        g.ndata["label"] = F.tensor(
+            node_labels, dtype=F.data_type_dict["int64"]
+        )
+        self.graph = g
+
+        self._random_split(
+            g.ndata["feature"], self.seed, self.train_size, self.val_size
+        )
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node features, node labels and masks
+
+            - ``ndata['feature']``: node features
+            - ``ndata['label']``: node labels
+            - ``ndata['train_mask']``: mask of training set
+            - ``ndata['val_mask']``: mask of validation set
+            - ``ndata['test_mask']``: mask of testing set
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self.graph
+        else:
+            return self._transform(self.graph)
+
+    def __len__(self):
+        """number of data examples"""
+        return len(self.graph)
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 2
+
+    def save(self):
+        """save processed data to directory `self.save_path`"""
+        graph_path = os.path.join(
+            self.save_path, self.name + "_dgl_graph_{}.bin".format(self.hash)
+        )
+        save_graphs(str(graph_path), self.graph)
+
+    def load(self):
+        """load processed data from directory `self.save_path`"""
+        graph_path = os.path.join(
+            self.save_path, self.name + "_dgl_graph_{}.bin".format(self.hash)
+        )
+        graph_list, _ = load_graphs(str(graph_path))
+        g = graph_list[0]
+        self.graph = g
+
+    def has_cache(self):
+        """check whether there are processed data in `self.save_path`"""
+        graph_path = os.path.join(
+            self.save_path, self.name + "_dgl_graph_{}.bin".format(self.hash)
+        )
+        return os.path.exists(graph_path)
+
+    def _random_split(self, x, seed=717, train_size=0.7, val_size=0.1):
+        """split the dataset into training set, validation set and testing set"""
+
+        assert 0 <= train_size + val_size <= 1, (
+            "The sum of valid training set size and validation set size "
+            "must between 0 and 1 (inclusive)."
+        )
+
+        N = x.shape[0]
+        index = np.arange(N)
+        if self.name == "amazon":
+            # 0-3304 are unlabeled nodes
+            index = np.arange(3305, N)
+
+        index = np.random.RandomState(seed).permutation(index)
+        train_idx = index[: int(train_size * len(index))]
+        val_idx = index[len(index) - int(val_size * len(index)) :]
+        test_idx = index[
+            int(train_size * len(index)) : len(index)
+            - int(val_size * len(index))
+        ]
+        train_mask = np.zeros(N, dtype=np.bool_)
+        val_mask = np.zeros(N, dtype=np.bool_)
+        test_mask = np.zeros(N, dtype=np.bool_)
+        train_mask[train_idx] = True
+        val_mask[val_idx] = True
+        test_mask[test_idx] = True
+        self.graph.ndata["train_mask"] = F.tensor(train_mask)
+        self.graph.ndata["val_mask"] = F.tensor(val_mask)
+        self.graph.ndata["test_mask"] = F.tensor(test_mask)
+
+
+class FraudYelpDataset(FraudDataset):
+    r"""Fraud Yelp Dataset
+
+    The Yelp dataset includes hotel and restaurant reviews filtered (spam) and recommended
+    (legitimate) by Yelp. A spam review detection task can be conducted, which is a binary
+    classification task. 32 handcrafted features from <http://dx.doi.org/10.1145/2783258.2783370>
+    are taken as the raw node features. Reviews are nodes in the graph, and three relations are:
+
+        1. R-U-R: it connects reviews posted by the same user
+        2. R-S-R: it connects reviews under the same product with the same star rating (1-5 stars)
+        3. R-T-R: it connects two reviews under the same product posted in the same month.
+
+    Statistics:
+
+    - Nodes: 45,954
+    - Edges:
+
+        - R-U-R: 98,630
+        - R-T-R: 1,147,232
+        - R-S-R: 6,805,486
+
+    - Classes:
+
+        - Positive (spam): 6,677
+        - Negative (legitimate): 39,277
+
+    - Positive-Negative ratio: 1 : 5.9
+    - Node feature size: 32
+
+    Parameters
+    ----------
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    random_seed : int
+        Specifying the random seed in splitting the dataset.
+        Default: 717
+    train_size : float
+        training set size of the dataset.
+        Default: 0.7
+    val_size : float
+        validation set size of the dataset, and the
+        size of testing set is (1 - train_size - val_size)
+        Default: 0.1
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Examples
+    --------
+    >>> dataset = FraudYelpDataset()
+    >>> graph = dataset[0]
+    >>> num_classes = dataset.num_classes
+    >>> feat = graph.ndata['feature']
+    >>> label = graph.ndata['label']
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        random_seed=717,
+        train_size=0.7,
+        val_size=0.1,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        super(FraudYelpDataset, self).__init__(
+            name="yelp",
+            raw_dir=raw_dir,
+            random_seed=random_seed,
+            train_size=train_size,
+            val_size=val_size,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class FraudAmazonDataset(FraudDataset):
+    r"""Fraud Amazon Dataset
+
+    The Amazon dataset includes product reviews under the Musical Instruments category.
+    Users with more than 80% helpful votes are labelled as benign entities and users with
+    less than 20% helpful votes are labelled as fraudulent entities. A fraudulent user
+    detection task can be conducted on the Amazon dataset, which is a binary classification
+    task. 25 handcrafted features from <https://arxiv.org/pdf/2005.10150.pdf> are taken as
+    the raw node features .
+
+    Users are nodes in the graph, and three relations are:
+    1. U-P-U : it connects users reviewing at least one same product
+    2. U-S-U : it connects users having at least one same star rating within one week
+    3. U-V-U : it connects users with top 5% mutual review text similarities (measured by
+    TF-IDF) among all users.
+
+    Statistics:
+
+    - Nodes: 11,944
+    - Edges:
+
+        - U-P-U: 351,216
+        - U-S-U: 7,132,958
+        - U-V-U: 2,073,474
+
+    - Classes:
+
+        - Positive (fraudulent): 821
+        - Negative (benign): 7,818
+        - Unlabeled: 3,305
+
+    - Positive-Negative ratio: 1 : 10.5
+    - Node feature size: 25
+
+    Parameters
+    ----------
+    raw_dir : str
+        Specifying the directory that will store the
+        downloaded data or the directory that
+        already stores the input data.
+        Default: ~/.dgl/
+    random_seed : int
+        Specifying the random seed in splitting the dataset.
+        Default: 717
+    train_size : float
+        training set size of the dataset.
+        Default: 0.7
+    val_size : float
+        validation set size of the dataset, and the
+        size of testing set is (1 - train_size - val_size)
+        Default: 0.1
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Examples
+    --------
+    >>> dataset = FraudAmazonDataset()
+    >>> graph = dataset[0]
+    >>> num_classes = dataset.num_classes
+    >>> feat = graph.ndata['feature']
+    >>> label = graph.ndata['label']
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        random_seed=717,
+        train_size=0.7,
+        val_size=0.1,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        super(FraudAmazonDataset, self).__init__(
+            name="amazon",
+            raw_dir=raw_dir,
+            random_seed=random_seed,
+            train_size=train_size,
+            val_size=val_size,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )

python/dgl/data/gdelt.py

@@ -0,0 +1,203 @@
+""" GDELT dataset for temporal graph """
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_info, loadtxt, save_info
+
+
+class GDELTDataset(DGLBuiltinDataset):
+    r"""GDELT dataset for event-based temporal graph
+
+    The Global Database of Events, Language, and Tone (GDELT) dataset.
+    This contains events happend all over the world (ie every protest held
+    anywhere in Russia on a given day is collapsed to a single entry).
+    This Dataset consists ofevents collected from 1/1/2018 to 1/31/2018
+    (15 minutes time granularity).
+
+    Reference:
+
+        - `Recurrent Event Network for Reasoning over Temporal Knowledge Graphs <https://arxiv.org/abs/1904.05530>`_
+        - `The Global Database of Events, Language, and Tone (GDELT) <https://dataverse.harvard.edu/dataset.xhtml?persistentId=doi:10.7910/DVN/28075>`_
+
+    Statistics:
+
+    - Train examples: 2,304
+    - Valid examples: 288
+    - Test examples: 384
+
+    Parameters
+    ----------
+    mode : str
+        Must be one of ('train', 'valid', 'test'). Default: 'train'
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    start_time : int
+        Start time of the temporal graph
+    end_time : int
+        End time of the temporal graph
+    is_temporal : bool
+        Does the dataset contain temporal graphs
+
+    Examples
+    ----------
+    >>> # get train, valid, test dataset
+    >>> train_data = GDELTDataset()
+    >>> valid_data = GDELTDataset(mode='valid')
+    >>> test_data = GDELTDataset(mode='test')
+    >>>
+    >>> # length of train set
+    >>> train_size = len(train_data)
+    >>>
+    >>> for g in train_data:
+    ....    e_feat = g.edata['rel_type']
+    ....    # your code here
+    ....
+    >>>
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        mode = mode.lower()
+        assert mode in ["train", "valid", "test"], "Mode not valid."
+        self.mode = mode
+        self.num_nodes = 23033
+        _url = _get_dgl_url("dataset/gdelt.zip")
+        super(GDELTDataset, self).__init__(
+            name="GDELT",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        file_path = os.path.join(self.raw_path, self.mode + ".txt")
+        self.data = loadtxt(file_path, delimiter="\t").astype(np.int64)
+
+        # The source code is not released, but the paper indicates there're
+        # totally 137 samples. The cutoff below has exactly 137 samples.
+        self.time_index = np.floor(self.data[:, 3] / 15).astype(np.int64)
+        self._start_time = self.time_index.min()
+        self._end_time = self.time_index.max()
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, self.mode + "_info.pkl")
+
+    def has_cache(self):
+        return os.path.exists(self.info_path)
+
+    def save(self):
+        save_info(
+            self.info_path,
+            {
+                "data": self.data,
+                "time_index": self.time_index,
+                "start_time": self.start_time,
+                "end_time": self.end_time,
+            },
+        )
+
+    def load(self):
+        info = load_info(self.info_path)
+        self.data, self.time_index, self._start_time, self._end_time = (
+            info["data"],
+            info["time_index"],
+            info["start_time"],
+            info["end_time"],
+        )
+
+    @property
+    def start_time(self):
+        r"""Start time of events in the temporal graph
+
+        Returns
+        -------
+        int
+        """
+        return self._start_time
+
+    @property
+    def end_time(self):
+        r"""End time of events in the temporal graph
+
+        Returns
+        -------
+        int
+        """
+        return self._end_time
+
+    def __getitem__(self, t):
+        r"""Get graph by with events before time `t + self.start_time`
+
+        Parameters
+        ----------
+        t : int
+            Time, its value must be in range [0, `self.end_time` - `self.start_time`]
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``edata['rel_type']``: edge type
+        """
+        if t >= len(self) or t < 0:
+            raise IndexError("Index out of range")
+        i = t + self.start_time
+        row_mask = self.time_index <= i
+        edges = self.data[row_mask][:, [0, 2]]
+        rate = self.data[row_mask][:, 1]
+        g = dgl_graph((edges[:, 0], edges[:, 1]))
+        g.edata["rel_type"] = F.tensor(
+            rate.reshape(-1, 1), dtype=F.data_type_dict["int64"]
+        )
+        if self._transform is not None:
+            g = self._transform(g)
+        return g
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return self._end_time - self._start_time + 1
+
+    @property
+    def is_temporal(self):
+        r"""Does the dataset contain temporal graphs
+
+        Returns
+        -------
+        bool
+        """
+        return True
+
+
+GDELT = GDELTDataset

python/dgl/data/geom_gcn.py

@@ -0,0 +1,483 @@
+"""Datasets introduced in the Geom-GCN paper."""
+import os
+
+import numpy as np
+
+from ..convert import graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url
+
+
+class GeomGCNDataset(DGLBuiltinDataset):
+    r"""Datasets introduced in
+    `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset.
+    raw_dir : str
+        Raw file directory to store the processed data.
+    force_reload : bool
+        Whether to re-download the data source.
+    verbose : bool
+        Whether to print progress information.
+    transform : callable
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(self, name, raw_dir, force_reload, verbose, transform):
+        url = _get_dgl_url(f"dataset/{name}.zip")
+        super(GeomGCNDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """Load and process the data."""
+        try:
+            import torch
+        except ImportError:
+            raise ModuleNotFoundError(
+                "This dataset requires PyTorch to be the backend."
+            )
+
+        # Process node features and labels.
+        with open(f"{self.raw_path}/out1_node_feature_label.txt", "r") as f:
+            data = f.read().split("\n")[1:-1]
+        features = [
+            [float(v) for v in r.split("\t")[1].split(",")] for r in data
+        ]
+        features = torch.tensor(features, dtype=torch.float)
+        labels = [int(r.split("\t")[2]) for r in data]
+        self._num_classes = max(labels) + 1
+        labels = torch.tensor(labels, dtype=torch.long)
+
+        # Process graph structure.
+        with open(f"{self.raw_path}/out1_graph_edges.txt", "r") as f:
+            data = f.read().split("\n")[1:-1]
+            data = [[int(v) for v in r.split("\t")] for r in data]
+        dst, src = torch.tensor(data, dtype=torch.long).t().contiguous()
+
+        self._g = graph((src, dst), num_nodes=features.size(0))
+        self._g.ndata["feat"] = features
+        self._g.ndata["label"] = labels
+
+        # Process 10 train/val/test node splits.
+        train_masks, val_masks, test_masks = [], [], []
+        for i in range(10):
+            filepath = f"{self.raw_path}/{self.name}_split_0.6_0.2_{i}.npz"
+            f = np.load(filepath)
+            train_masks += [torch.from_numpy(f["train_mask"])]
+            val_masks += [torch.from_numpy(f["val_mask"])]
+            test_masks += [torch.from_numpy(f["test_mask"])]
+        self._g.ndata["train_mask"] = torch.stack(train_masks, dim=1).bool()
+        self._g.ndata["val_mask"] = torch.stack(val_masks, dim=1).bool()
+        self._g.ndata["test_mask"] = torch.stack(test_masks, dim=1).bool()
+
+    def has_cache(self):
+        return os.path.exists(self.raw_path)
+
+    def load(self):
+        self.process()
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return self._num_classes
+
+
+class ChameleonDataset(GeomGCNDataset):
+    r"""Wikipedia page-page network on chameleons from `Multi-scale Attributed
+    Node Embedding <https://arxiv.org/abs/1909.13021>`__ and later modified by
+    `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent articles from the English Wikipedia, edges reflect mutual
+    links between them. Node features indicate the presence of particular nouns
+    in the articles. The nodes were classified into 5 classes in terms of their
+    average monthly traffic.
+
+    Statistics:
+
+    - Nodes: 2277
+    - Edges: 36101
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 1092
+        - Val: 729
+        - Test: 456
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import ChameleonDataset
+    >>> dataset = ChameleonDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(ChameleonDataset, self).__init__(
+            name="chameleon",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class SquirrelDataset(GeomGCNDataset):
+    r"""Wikipedia page-page network on squirrels from `Multi-scale Attributed
+    Node Embedding <https://arxiv.org/abs/1909.13021>`__ and later modified by
+    `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent articles from the English Wikipedia, edges reflect mutual
+    links between them. Node features indicate the presence of particular nouns
+    in the articles. The nodes were classified into 5 classes in terms of their
+    average monthly traffic.
+
+    Statistics:
+
+    - Nodes: 5201
+    - Edges: 217073
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 2496
+        - Val: 1664
+        - Test: 1041
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import SquirrelDataset
+    >>> dataset = SquirrelDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(SquirrelDataset, self).__init__(
+            name="squirrel",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class CornellDataset(GeomGCNDataset):
+    r"""Cornell subset of
+    `WebKB <http://www.cs.cmu.edu/afs/cs.cmu.edu/project/theo-11/www/wwkb/>`__,
+    later modified by `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent web pages. Edges represent hyperlinks between them. Node
+    features are the bag-of-words representation of web pages. The web pages
+    are manually classified into the five categories, student, project, course,
+    staff, and faculty.
+
+    Statistics:
+
+    - Nodes: 183
+    - Edges: 298
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 87
+        - Val: 59
+        - Test: 37
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import CornellDataset
+    >>> dataset = CornellDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(CornellDataset, self).__init__(
+            name="cornell",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class TexasDataset(GeomGCNDataset):
+    r"""Texas subset of
+    `WebKB <http://www.cs.cmu.edu/afs/cs.cmu.edu/project/theo-11/www/wwkb/>`__,
+    later modified by `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent web pages. Edges represent hyperlinks between them. Node
+    features are the bag-of-words representation of web pages. The web pages
+    are manually classified into the five categories, student, project, course,
+    staff, and faculty.
+
+    Statistics:
+
+    - Nodes: 183
+    - Edges: 325
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 87
+        - Val: 59
+        - Test: 37
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import TexasDataset
+    >>> dataset = TexasDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(TexasDataset, self).__init__(
+            name="texas",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class WisconsinDataset(GeomGCNDataset):
+    r"""Wisconsin subset of
+    `WebKB <http://www.cs.cmu.edu/afs/cs.cmu.edu/project/theo-11/www/wwkb/>`__,
+    later modified by `Geom-GCN: Geometric Graph Convolutional Networks
+    <https://arxiv.org/abs/2002.05287>`__
+
+    Nodes represent web pages. Edges represent hyperlinks between them. Node
+    features are the bag-of-words representation of web pages. The web pages
+    are manually classified into the five categories, student, project, course,
+    staff, and faculty.
+
+    Statistics:
+
+    - Nodes: 251
+    - Edges: 515
+    - Number of Classes: 5
+    - 10 train/val/test splits
+
+        - Train: 120
+        - Val: 80
+        - Test: 51
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Notes
+    -----
+    The graph does not come with edges for both directions.
+
+    Examples
+    --------
+
+    >>> from dgl.data import WisconsinDataset
+    >>> dataset = WisconsinDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get data split
+    >>> train_mask = g.ndata["train_mask"]
+    >>> val_mask = g.ndata["val_mask"]
+    >>> test_mask = g.ndata["test_mask"]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(WisconsinDataset, self).__init__(
+            name="wisconsin",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )

python/dgl/data/gindt.py

@@ -0,0 +1,420 @@
+"""Datasets used in How Powerful Are Graph Neural Networks?
+(chen jun)
+Datasets include:
+MUTAG, COLLAB, IMDBBINARY, IMDBMULTI, NCI1, PROTEINS, PTC, REDDITBINARY, REDDITMULTI5K
+https://github.com/weihua916/powerful-gnns/blob/master/dataset.zip
+"""
+
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from ..utils import retry_method_with_fix
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    download,
+    extract_archive,
+    load_graphs,
+    load_info,
+    loadtxt,
+    save_graphs,
+    save_info,
+)
+
+
+class GINDataset(DGLBuiltinDataset):
+    """Dataset Class for `How Powerful Are Graph Neural Networks? <https://arxiv.org/abs/1810.00826>`_.
+
+    This is adapted from `<https://github.com/weihua916/powerful-gnns/blob/master/dataset.zip>`_.
+
+    The class provides an interface for nine datasets used in the paper along with the paper-specific
+    settings. The datasets are ``'MUTAG'``, ``'COLLAB'``, ``'IMDBBINARY'``, ``'IMDBMULTI'``,
+    ``'NCI1'``, ``'PROTEINS'``, ``'PTC'``, ``'REDDITBINARY'``, ``'REDDITMULTI5K'``.
+
+    If ``degree_as_nlabel`` is set to ``False``, then ``ndata['label']`` stores the provided node label,
+    otherwise ``ndata['label']`` stores the node in-degrees.
+
+    For graphs that have node attributes, ``ndata['attr']`` stores the node attributes.
+    For graphs that have no attribute, ``ndata['attr']`` stores the corresponding one-hot encoding
+    of ``ndata['label']``.
+
+    Parameters
+    ---------
+    name: str
+        dataset name, one of
+        (``'MUTAG'``, ``'COLLAB'``, \
+        ``'IMDBBINARY'``, ``'IMDBMULTI'``, \
+        ``'NCI1'``, ``'PROTEINS'``, ``'PTC'``, \
+        ``'REDDITBINARY'``, ``'REDDITMULTI5K'``)
+    self_loop: bool
+        add self to self edge if true
+    degree_as_nlabel: bool
+        take node degree as label and feature if true
+    transform: callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for multiclass classification
+
+    Examples
+    --------
+    >>> data = GINDataset(name='MUTAG', self_loop=False)
+
+    The dataset instance is an iterable
+
+    >>> len(data)
+    188
+    >>> g, label = data[128]
+    >>> g
+    Graph(num_nodes=13, num_edges=26,
+          ndata_schemes={'label': Scheme(shape=(), dtype=torch.int64), 'attr': Scheme(shape=(7,), dtype=torch.float32)}
+          edata_schemes={})
+    >>> label
+    tensor(1)
+
+    Batch the graphs and labels for mini-batch training
+
+    >>> graphs, labels = zip(*[data[i] for i in range(16)])
+    >>> batched_graphs = dgl.batch(graphs)
+    >>> batched_labels = torch.tensor(labels)
+    >>> batched_graphs
+    Graph(num_nodes=330, num_edges=748,
+          ndata_schemes={'label': Scheme(shape=(), dtype=torch.int64), 'attr': Scheme(shape=(7,), dtype=torch.float32)}
+          edata_schemes={})
+    """
+
+    def __init__(
+        self,
+        name,
+        self_loop,
+        degree_as_nlabel=False,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self._name = name  # MUTAG
+        gin_url = "https://raw.githubusercontent.com/weihua916/powerful-gnns/master/dataset.zip"
+        self.ds_name = "nig"
+
+        self.self_loop = self_loop
+        self.graphs = []
+        self.labels = []
+
+        # relabel
+        self.glabel_dict = {}
+        self.nlabel_dict = {}
+        self.elabel_dict = {}
+        self.ndegree_dict = {}
+
+        # global num
+        self.N = 0  # total graphs number
+        self.n = 0  # total nodes number
+        self.m = 0  # total edges number
+
+        # global num of classes
+        self.gclasses = 0
+        self.nclasses = 0
+        self.eclasses = 0
+        self.dim_nfeats = 0
+
+        # flags
+        self.degree_as_nlabel = degree_as_nlabel
+        self.nattrs_flag = False
+        self.nlabels_flag = False
+
+        super(GINDataset, self).__init__(
+            name=name,
+            url=gin_url,
+            hash_key=(name, self_loop, degree_as_nlabel),
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def raw_path(self):
+        return os.path.join(self.raw_dir, "GINDataset")
+
+    def download(self):
+        r"""Automatically download data and extract it."""
+        zip_file_path = os.path.join(self.raw_dir, "GINDataset.zip")
+        download(self.url, path=zip_file_path)
+        extract_archive(zip_file_path, self.raw_path)
+
+    def __len__(self):
+        """Return the number of graphs in the dataset."""
+        return len(self.graphs)
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (:class:`dgl.Graph`, Tensor)
+            The graph and its label.
+        """
+        if self._transform is None:
+            g = self.graphs[idx]
+        else:
+            g = self._transform(self.graphs[idx])
+        return g, self.labels[idx]
+
+    def _file_path(self):
+        return os.path.join(
+            self.raw_dir,
+            "GINDataset",
+            "dataset",
+            self.name,
+            "{}.txt".format(self.name),
+        )
+
+    def process(self):
+        """Loads input dataset from dataset/NAME/NAME.txt file"""
+        if self.verbose:
+            print("loading data...")
+        self.file = self._file_path()
+        with open(self.file, "r") as f:
+            # line_1 == N, total number of graphs
+            self.N = int(f.readline().strip())
+
+            for i in range(self.N):
+                if (i + 1) % 10 == 0 and self.verbose is True:
+                    print("processing graph {}...".format(i + 1))
+
+                grow = f.readline().strip().split()
+                # line_2 == [n_nodes, l] is equal to
+                # [node number of a graph, class label of a graph]
+                n_nodes, glabel = [int(w) for w in grow]
+
+                # relabel graphs
+                if glabel not in self.glabel_dict:
+                    mapped = len(self.glabel_dict)
+                    self.glabel_dict[glabel] = mapped
+
+                self.labels.append(self.glabel_dict[glabel])
+
+                g = dgl_graph(([], []))
+                g.add_nodes(n_nodes)
+
+                nlabels = []  # node labels
+                nattrs = []  # node attributes if it has
+                m_edges = 0
+
+                for j in range(n_nodes):
+                    nrow = f.readline().strip().split()
+
+                    # handle edges and attributes(if has)
+                    tmp = int(nrow[1]) + 2  # tmp == 2 + #edges
+                    if tmp == len(nrow):
+                        # no node attributes
+                        nrow = [int(w) for w in nrow]
+                    elif tmp > len(nrow):
+                        nrow = [int(w) for w in nrow[:tmp]]
+                        nattr = [float(w) for w in nrow[tmp:]]
+                        nattrs.append(nattr)
+                    else:
+                        raise Exception("edge number is incorrect!")
+
+                    # relabel nodes if it has labels
+                    # if it doesn't have node labels, then every nrow[0]==0
+                    if not nrow[0] in self.nlabel_dict:
+                        mapped = len(self.nlabel_dict)
+                        self.nlabel_dict[nrow[0]] = mapped
+
+                    nlabels.append(self.nlabel_dict[nrow[0]])
+
+                    m_edges += nrow[1]
+                    g.add_edges(j, nrow[2:])
+
+                    # add self loop
+                    if self.self_loop:
+                        m_edges += 1
+                        g.add_edges(j, j)
+
+                    if (j + 1) % 10 == 0 and self.verbose is True:
+                        print(
+                            "processing node {} of graph {}...".format(
+                                j + 1, i + 1
+                            )
+                        )
+                        print("this node has {} edgs.".format(nrow[1]))
+
+                if nattrs != []:
+                    nattrs = np.stack(nattrs)
+                    g.ndata["attr"] = F.tensor(nattrs, F.float32)
+                    self.nattrs_flag = True
+
+                g.ndata["label"] = F.tensor(nlabels)
+                if len(self.nlabel_dict) > 1:
+                    self.nlabels_flag = True
+
+                assert g.num_nodes() == n_nodes
+
+                # update statistics of graphs
+                self.n += n_nodes
+                self.m += m_edges
+
+                self.graphs.append(g)
+
+        self.labels = F.tensor(self.labels)
+        # if no attr
+        if not self.nattrs_flag:
+            if self.verbose:
+                print("there are no node features in this dataset!")
+            # generate node attr by node degree
+            if self.degree_as_nlabel:
+                if self.verbose:
+                    print("generate node features by node degree...")
+                for g in self.graphs:
+                    # actually this label shouldn't be updated
+                    # in case users want to keep it
+                    # but usually no features means no labels, fine.
+                    g.ndata["label"] = g.in_degrees()
+                    # extracting unique node labels
+
+            # in case the labels/degrees are not continuous number
+            nlabel_set = set([])
+            for g in self.graphs:
+                nlabel_set = nlabel_set.union(
+                    set([F.as_scalar(nl) for nl in g.ndata["label"]])
+                )
+            nlabel_set = list(nlabel_set)
+            is_label_valid = all(
+                [label in self.nlabel_dict for label in nlabel_set]
+            )
+            if (
+                is_label_valid
+                and len(nlabel_set) == np.max(nlabel_set) + 1
+                and np.min(nlabel_set) == 0
+            ):
+                # Note this is different from the author's implementation. In weihua916's implementation,
+                # the labels are relabeled anyway. But here we didn't relabel it if the labels are contiguous
+                # to make it consistent with the original dataset
+                label2idx = self.nlabel_dict
+            else:
+                label2idx = {nlabel_set[i]: i for i in range(len(nlabel_set))}
+            # generate node attr by node label
+            for g in self.graphs:
+                attr = np.zeros((g.num_nodes(), len(label2idx)))
+                attr[
+                    range(g.num_nodes()),
+                    [
+                        label2idx[nl]
+                        for nl in F.asnumpy(g.ndata["label"]).tolist()
+                    ],
+                ] = 1
+                g.ndata["attr"] = F.tensor(attr, F.float32)
+
+        # after load, get the #classes and #dim
+        self.gclasses = len(self.glabel_dict)
+        self.nclasses = len(self.nlabel_dict)
+        self.eclasses = len(self.elabel_dict)
+        self.dim_nfeats = len(self.graphs[0].ndata["attr"][0])
+
+        if self.verbose:
+            print("Done.")
+            print(
+                """
+                -------- Data Statistics --------'
+                #Graphs: %d
+                #Graph Classes: %d
+                #Nodes: %d
+                #Node Classes: %d
+                #Node Features Dim: %d
+                #Edges: %d
+                #Edge Classes: %d
+                Avg. of #Nodes: %.2f
+                Avg. of #Edges: %.2f
+                Graph Relabeled: %s
+                Node Relabeled: %s
+                Degree Relabeled(If degree_as_nlabel=True): %s \n """
+                % (
+                    self.N,
+                    self.gclasses,
+                    self.n,
+                    self.nclasses,
+                    self.dim_nfeats,
+                    self.m,
+                    self.eclasses,
+                    self.n / self.N,
+                    self.m / self.N,
+                    self.glabel_dict,
+                    self.nlabel_dict,
+                    self.ndegree_dict,
+                )
+            )
+
+    def save(self):
+        label_dict = {"labels": self.labels}
+        info_dict = {
+            "N": self.N,
+            "n": self.n,
+            "m": self.m,
+            "self_loop": self.self_loop,
+            "gclasses": self.gclasses,
+            "nclasses": self.nclasses,
+            "eclasses": self.eclasses,
+            "dim_nfeats": self.dim_nfeats,
+            "degree_as_nlabel": self.degree_as_nlabel,
+            "glabel_dict": self.glabel_dict,
+            "nlabel_dict": self.nlabel_dict,
+            "elabel_dict": self.elabel_dict,
+            "ndegree_dict": self.ndegree_dict,
+        }
+        save_graphs(str(self.graph_path), self.graphs, label_dict)
+        save_info(str(self.info_path), info_dict)
+
+    def load(self):
+        graphs, label_dict = load_graphs(str(self.graph_path))
+        info_dict = load_info(str(self.info_path))
+
+        self.graphs = graphs
+        self.labels = label_dict["labels"]
+
+        self.N = info_dict["N"]
+        self.n = info_dict["n"]
+        self.m = info_dict["m"]
+        self.self_loop = info_dict["self_loop"]
+        self.gclasses = info_dict["gclasses"]
+        self.nclasses = info_dict["nclasses"]
+        self.eclasses = info_dict["eclasses"]
+        self.dim_nfeats = info_dict["dim_nfeats"]
+        self.glabel_dict = info_dict["glabel_dict"]
+        self.nlabel_dict = info_dict["nlabel_dict"]
+        self.elabel_dict = info_dict["elabel_dict"]
+        self.ndegree_dict = info_dict["ndegree_dict"]
+        self.degree_as_nlabel = info_dict["degree_as_nlabel"]
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "gin_{}_{}.bin".format(self.name, self.hash)
+        )
+
+    @property
+    def info_path(self):
+        return os.path.join(
+            self.save_path, "gin_{}_{}.pkl".format(self.name, self.hash)
+        )
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+        return False
+
+    @property
+    def num_classes(self):
+        return self.gclasses

python/dgl/data/gnn_benchmark.py

@@ -0,0 +1,544 @@
+"""GNN Benchmark datasets for node classification."""
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F, transforms
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    _get_dgl_url,
+    deprecate_class,
+    deprecate_property,
+    load_graphs,
+    save_graphs,
+)
+
+__all__ = [
+    "AmazonCoBuyComputerDataset",
+    "AmazonCoBuyPhotoDataset",
+    "CoauthorPhysicsDataset",
+    "CoauthorCSDataset",
+    "CoraFullDataset",
+    "AmazonCoBuy",
+    "Coauthor",
+    "CoraFull",
+]
+
+
+def eliminate_self_loops(A):
+    """Remove self-loops from the adjacency matrix."""
+    A = A.tolil()
+    A.setdiag(0)
+    A = A.tocsr()
+    A.eliminate_zeros()
+    return A
+
+
+class GNNBenchmarkDataset(DGLBuiltinDataset):
+    r"""Base Class for GNN Benchmark dataset
+
+    Reference: https://github.com/shchur/gnn-benchmark#datasets
+    """
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        _url = _get_dgl_url("dataset/" + name + ".zip")
+        super(GNNBenchmarkDataset, self).__init__(
+            name=name,
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        npz_path = os.path.join(self.raw_path, self.name + ".npz")
+        g = self._load_npz(npz_path)
+        g = transforms.reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+        self._graph = g
+        self._data = [g]
+        self._print_info()
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph_v1.bin")
+        if os.path.exists(graph_path):
+            return True
+        return False
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph_v1.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph_v1.bin")
+        graphs, _ = load_graphs(graph_path)
+        self._graph = graphs[0]
+        self._data = [graphs[0]]
+        self._print_info()
+
+    def _print_info(self):
+        if self.verbose:
+            print("  NumNodes: {}".format(self._graph.num_nodes()))
+            print("  NumEdges: {}".format(self._graph.num_edges()))
+            print("  NumFeats: {}".format(self._graph.ndata["feat"].shape[-1]))
+            print("  NumbClasses: {}".format(self.num_classes))
+
+    def _load_npz(self, file_name):
+        with np.load(file_name, allow_pickle=True) as loader:
+            loader = dict(loader)
+            num_nodes = loader["adj_shape"][0]
+            adj_matrix = sp.csr_matrix(
+                (
+                    loader["adj_data"],
+                    loader["adj_indices"],
+                    loader["adj_indptr"],
+                ),
+                shape=loader["adj_shape"],
+            ).tocoo()
+
+            if "attr_data" in loader:
+                # Attributes are stored as a sparse CSR matrix
+                attr_matrix = sp.csr_matrix(
+                    (
+                        loader["attr_data"],
+                        loader["attr_indices"],
+                        loader["attr_indptr"],
+                    ),
+                    shape=loader["attr_shape"],
+                ).todense()
+            elif "attr_matrix" in loader:
+                # Attributes are stored as a (dense) np.ndarray
+                attr_matrix = loader["attr_matrix"]
+            else:
+                attr_matrix = None
+
+            if "labels_data" in loader:
+                # Labels are stored as a CSR matrix
+                labels = sp.csr_matrix(
+                    (
+                        loader["labels_data"],
+                        loader["labels_indices"],
+                        loader["labels_indptr"],
+                    ),
+                    shape=loader["labels_shape"],
+                ).todense()
+            elif "labels" in loader:
+                # Labels are stored as a numpy array
+                labels = loader["labels"]
+            else:
+                labels = None
+        g = dgl_graph((adj_matrix.row, adj_matrix.col))
+        g = transforms.to_bidirected(g)
+        g.ndata["feat"] = F.tensor(attr_matrix, F.data_type_dict["float32"])
+        g.ndata["label"] = F.tensor(labels, F.data_type_dict["int64"])
+        return g
+
+    @property
+    def num_classes(self):
+        """Number of classes."""
+        raise NotImplementedError
+
+    def __getitem__(self, idx):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        r"""Number of graphs in the dataset"""
+        return 1
+
+
+class CoraFullDataset(GNNBenchmarkDataset):
+    r"""CORA-Full dataset for node classification task.
+
+    Extended Cora dataset. Nodes represent paper and edges represent citations.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics:
+
+    - Nodes: 19,793
+    - Edges: 126,842 (note that the original dataset has 65,311 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of Classes: 70
+    - Node feature size: 8,710
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = CoraFullDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(CoraFullDataset, self).__init__(
+            name="cora_full",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 70
+
+
+class CoauthorCSDataset(GNNBenchmarkDataset):
+    r"""'Computer Science (CS)' part of the Coauthor dataset for node classification task.
+
+    Coauthor CS and Coauthor Physics are co-authorship graphs based on the Microsoft Academic Graph
+    from the KDD Cup 2016 challenge. Here, nodes are authors, that are connected by an edge if they
+    co-authored a paper; node features represent paper keywords for each author’s papers, and class
+    labels indicate most active fields of study for each author.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics:
+
+    - Nodes: 18,333
+    - Edges: 163,788 (note that the original dataset has 81,894 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of classes: 15
+    - Node feature size: 6,805
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = CoauthorCSDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(CoauthorCSDataset, self).__init__(
+            name="coauthor_cs",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 15
+
+
+class CoauthorPhysicsDataset(GNNBenchmarkDataset):
+    r"""'Physics' part of the Coauthor dataset for node classification task.
+
+    Coauthor CS and Coauthor Physics are co-authorship graphs based on the Microsoft Academic Graph
+    from the KDD Cup 2016 challenge. Here, nodes are authors, that are connected by an edge if they
+    co-authored a paper; node features represent paper keywords for each author’s papers, and class
+    labels indicate most active fields of study for each author.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics
+
+    - Nodes: 34,493
+    - Edges: 495,924 (note that the original dataset has 247,962 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of classes: 5
+    - Node feature size: 8,415
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = CoauthorPhysicsDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(CoauthorPhysicsDataset, self).__init__(
+            name="coauthor_physics",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 5
+
+
+class AmazonCoBuyComputerDataset(GNNBenchmarkDataset):
+    r"""'Computer' part of the AmazonCoBuy dataset for node classification task.
+
+    Amazon Computers and Amazon Photo are segments of the Amazon co-purchase graph [McAuley et al., 2015],
+    where nodes represent goods, edges indicate that two goods are frequently bought together, node
+    features are bag-of-words encoded product reviews, and class labels are given by the product category.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics:
+
+    - Nodes: 13,752
+    - Edges: 491,722 (note that the original dataset has 245,778 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of classes: 10
+    - Node feature size: 767
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = AmazonCoBuyComputerDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(AmazonCoBuyComputerDataset, self).__init__(
+            name="amazon_co_buy_computer",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 10
+
+
+class AmazonCoBuyPhotoDataset(GNNBenchmarkDataset):
+    r"""AmazonCoBuy dataset for node classification task.
+
+    Amazon Computers and Amazon Photo are segments of the Amazon co-purchase graph [McAuley et al., 2015],
+    where nodes represent goods, edges indicate that two goods are frequently bought together, node
+    features are bag-of-words encoded product reviews, and class labels are given by the product category.
+
+    Reference: `<https://github.com/shchur/gnn-benchmark#datasets>`_
+
+    Statistics
+
+    - Nodes: 7,650
+    - Edges: 238,163 (note that the original dataset has 119,043 edges but DGL adds
+      the reverse edges and remove the duplicates, hence with a different number)
+    - Number of classes: 8
+    - Node feature size: 745
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> data = AmazonCoBuyPhotoDataset()
+    >>> g = data[0]
+    >>> num_class = data.num_classes
+    >>> feat = g.ndata['feat']  # get node feature
+    >>> label = g.ndata['label']  # get node labels
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(AmazonCoBuyPhotoDataset, self).__init__(
+            name="amazon_co_buy_photo",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def num_classes(self):
+        """Number of classes.
+
+        Return
+        -------
+        int
+        """
+        return 8
+
+
+class CoraFull(CoraFullDataset):
+    def __init__(self, **kwargs):
+        deprecate_class("CoraFull", "CoraFullDataset")
+        super(CoraFull, self).__init__(**kwargs)
+
+
+def AmazonCoBuy(name):
+    if name == "computers":
+        deprecate_class("AmazonCoBuy", "AmazonCoBuyComputerDataset")
+        return AmazonCoBuyComputerDataset()
+    elif name == "photo":
+        deprecate_class("AmazonCoBuy", "AmazonCoBuyPhotoDataset")
+        return AmazonCoBuyPhotoDataset()
+    else:
+        raise ValueError('Dataset name should be "computers" or "photo".')
+
+
+def Coauthor(name):
+    if name == "cs":
+        deprecate_class("Coauthor", "CoauthorCSDataset")
+        return CoauthorCSDataset()
+    elif name == "physics":
+        deprecate_class("Coauthor", "CoauthorPhysicsDataset")
+        return CoauthorPhysicsDataset()
+    else:
+        raise ValueError('Dataset name should be "cs" or "physics".')

python/dgl/data/graph_serialize.py

@@ -0,0 +1,272 @@
+"""For Graph Serialization"""
+from __future__ import absolute_import
+
+import os
+
+from .. import backend as F
+from .._ffi.function import _init_api
+from .._ffi.object import ObjectBase, register_object
+from ..base import dgl_warning, DGLError
+from ..heterograph import DGLGraph
+from .heterograph_serialize import save_heterographs
+
+_init_api("dgl.data.graph_serialize")
+
+__all__ = ["save_graphs", "load_graphs", "load_labels"]
+
+
+@register_object("graph_serialize.StorageMetaData")
+class StorageMetaData(ObjectBase):
+    """StorageMetaData Object
+    attributes available:
+      num_graph [int]: return numbers of graphs
+      nodes_num_list Value of NDArray: return number of nodes for each graph
+      edges_num_list Value of NDArray: return number of edges for each graph
+      labels [dict of backend tensors]: return dict of labels
+      graph_data [list of GraphData]: return list of GraphData Object
+    """
+
+
+def is_local_path(filepath):
+    return not (
+        filepath.startswith("hdfs://")
+        or filepath.startswith("viewfs://")
+        or filepath.startswith("s3://")
+    )
+
+
+def check_local_file_exists(filename):
+    if is_local_path(filename) and not os.path.exists(filename):
+        raise DGLError("File {} does not exist.".format(filename))
+
+
+@register_object("graph_serialize.GraphData")
+class GraphData(ObjectBase):
+    """GraphData Object"""
+
+    @staticmethod
+    def create(g):
+        """Create GraphData"""
+        # TODO(zihao): support serialize batched graph in the future.
+        assert (
+            g.batch_size == 1
+        ), "Batched DGLGraph is not supported for serialization"
+        ghandle = g._graph
+        if len(g.ndata) != 0:
+            node_tensors = dict()
+            for key, value in g.ndata.items():
+                node_tensors[key] = F.zerocopy_to_dgl_ndarray(value)
+        else:
+            node_tensors = None
+        if len(g.edata) != 0:
+            edge_tensors = dict()
+            for key, value in g.edata.items():
+                edge_tensors[key] = F.zerocopy_to_dgl_ndarray(value)
+        else:
+            edge_tensors = None
+        return _CAPI_MakeGraphData(ghandle, node_tensors, edge_tensors)
+
+    def get_graph(self):
+        """Get DGLGraph from GraphData"""
+        ghandle = _CAPI_GDataGraphHandle(self)
+        hgi = _CAPI_DGLAsHeteroGraph(ghandle)
+        g = DGLGraph(hgi, ["_U"], ["_E"])
+        node_tensors_items = _CAPI_GDataNodeTensors(self).items()
+        edge_tensors_items = _CAPI_GDataEdgeTensors(self).items()
+        for k, v in node_tensors_items:
+            g.ndata[k] = F.zerocopy_from_dgl_ndarray(v)
+        for k, v in edge_tensors_items:
+            g.edata[k] = F.zerocopy_from_dgl_ndarray(v)
+        return g
+
+
+def save_graphs(filename, g_list, labels=None, formats=None):
+    r"""Save graphs and optionally their labels to file.
+
+    Besides saving to local files, DGL supports writing the graphs directly
+    to S3 (by providing a ``"s3://..."`` path) or to HDFS (by providing
+    ``"hdfs://..."`` a path).
+
+    The function saves both the graph structure and node/edge features to file
+    in DGL's own binary format. For graph-level features, pass them via
+    the :attr:`labels` argument.
+
+    Parameters
+    ----------
+    filename : str
+        The file name to store the graphs and labels.
+    g_list: list
+        The graphs to be saved.
+    labels: dict[str, Tensor]
+        labels should be dict of tensors, with str as keys
+    formats: str or list[str]
+        Save graph in specified formats. It could be any combination of
+        ``coo``, ``csc`` and ``csr``. If not specified, save one format
+        only according to what format is available. If multiple formats
+        are available, selection priority from high to low is ``coo``,
+        ``csc``, ``csr``.
+
+    Examples
+    ----------
+    >>> import dgl
+    >>> import torch as th
+
+    Create :class:`DGLGraph` objects and initialize node
+    and edge features.
+
+    >>> g1 = dgl.graph(([0, 1, 2], [1, 2, 3]))
+    >>> g2 = dgl.graph(([0, 2], [2, 3]))
+    >>> g2.edata["e"] = th.ones(2, 4)
+
+    Save Graphs into file
+
+    >>> from dgl.data.utils import save_graphs
+    >>> graph_labels = {"glabel": th.tensor([0, 1])}
+    >>> save_graphs("./data.bin", [g1, g2], graph_labels)
+
+    See Also
+    --------
+    load_graphs
+    """
+    # if it is local file, do some sanity check
+    if is_local_path(filename):
+        if os.path.isdir(filename):
+            raise DGLError(
+                "Filename {} is an existing directory.".format(filename)
+            )
+        f_path = os.path.dirname(filename)
+        if f_path and not os.path.exists(f_path):
+            os.makedirs(f_path)
+    g_sample = g_list[0] if isinstance(g_list, list) else g_list
+    if type(g_sample) == DGLGraph:  # Doesn't support DGLGraph's derived class
+        save_heterographs(filename, g_list, labels, formats)
+    else:
+        raise DGLError(
+            "Invalid argument g_list. Must be a DGLGraph or a list of DGLGraphs."
+        )
+
+
+def load_graphs(filename, idx_list=None):
+    """Load graphs and optionally their labels from file saved by :func:`save_graphs`.
+
+    Besides loading from local files, DGL supports loading the graphs directly
+    from S3 (by providing a ``"s3://..."`` path) or from HDFS (by providing
+    ``"hdfs://..."`` a path).
+
+    Parameters
+    ----------
+    filename: str
+        The file name to load graphs from.
+    idx_list: list[int], optional
+        The indices of the graphs to be loaded if the file contains multiple graphs.
+        Default is loading all the graphs stored in the file.
+
+    Returns
+    --------
+    graph_list: list[DGLGraph]
+        The loaded graphs.
+    labels: dict[str, Tensor]
+        The graph labels stored in file. If no label is stored, the dictionary is empty.
+        Regardless of whether the ``idx_list`` argument is given or not,
+        the returned dictionary always contains the labels of all the graphs.
+
+    Examples
+    ----------
+    Following the example in :func:`save_graphs`.
+
+    >>> from dgl.data.utils import load_graphs
+    >>> glist, label_dict = load_graphs("./data.bin") # glist will be [g1, g2]
+    >>> glist, label_dict = load_graphs("./data.bin", [0]) # glist will be [g1]
+
+    See Also
+    --------
+    save_graphs
+    """
+    # if it is local file, do some sanity check
+    check_local_file_exists(filename)
+    version = _CAPI_GetFileVersion(filename)
+    if version == 1:
+        dgl_warning(
+            "You are loading a graph file saved by old version of dgl.  \
+            Please consider saving it again with the current format."
+        )
+        return load_graph_v1(filename, idx_list)
+    elif version == 2:
+        return load_graph_v2(filename, idx_list)
+    else:
+        raise DGLError("Invalid DGL Version Number.")
+
+
+def load_graph_v2(filename, idx_list=None):
+    """Internal functions for loading DGLGraphs."""
+    if idx_list is None:
+        idx_list = []
+    assert isinstance(idx_list, list)
+    heterograph_list = _CAPI_LoadGraphFiles_V2(filename, idx_list)
+    label_dict = load_labels_v2(filename)
+    return [gdata.get_graph() for gdata in heterograph_list], label_dict
+
+
+def load_graph_v1(filename, idx_list=None):
+    """ "Internal functions for loading DGLGraphs (V0)."""
+    if idx_list is None:
+        idx_list = []
+    assert isinstance(idx_list, list)
+    metadata = _CAPI_LoadGraphFiles_V1(filename, idx_list, False)
+    label_dict = {}
+    for k, v in metadata.labels.items():
+        label_dict[k] = F.zerocopy_from_dgl_ndarray(v)
+    return [gdata.get_graph() for gdata in metadata.graph_data], label_dict
+
+
+def load_labels(filename):
+    """
+    Load label dict from file
+
+    Parameters
+    ----------
+    filename: str
+        filename to load DGLGraphs
+
+    Returns
+    ----------
+    labels: dict
+        dict of labels stored in file (empty dict returned if no
+        label stored)
+
+    Examples
+    ----------
+    Following the example in save_graphs.
+
+    >>> from dgl.data.utils import load_labels
+    >>> label_dict = load_graphs("./data.bin")
+
+    """
+    # if it is local file, do some sanity check
+    check_local_file_exists(filename)
+
+    version = _CAPI_GetFileVersion(filename)
+    if version == 1:
+        return load_labels_v1(filename)
+    elif version == 2:
+        return load_labels_v2(filename)
+    else:
+        raise Exception("Invalid DGL Version Number")
+
+
+def load_labels_v2(filename):
+    """Internal functions for loading labels from V2 format"""
+    label_dict = {}
+    nd_dict = _CAPI_LoadLabels_V2(filename)
+    for k, v in nd_dict.items():
+        label_dict[k] = F.zerocopy_from_dgl_ndarray(v)
+    return label_dict
+
+
+def load_labels_v1(filename):
+    """Internal functions for loading labels from V1 format"""
+    metadata = _CAPI_LoadGraphFiles_V1(filename, [], True)
+    label_dict = {}
+    for k, v in metadata.labels.items():
+        label_dict[k] = F.zerocopy_from_dgl_ndarray(v)
+    return label_dict

python/dgl/data/heterograph_serialize.py

@@ -0,0 +1,79 @@
+"""For HeteroGraph Serialization"""
+from __future__ import absolute_import
+
+from .. import backend as F
+from .._ffi.function import _init_api
+from .._ffi.object import ObjectBase, register_object
+from ..container import convert_to_strmap
+from ..frame import Frame
+from ..heterograph import DGLGraph
+
+_init_api("dgl.data.heterograph_serialize")
+
+
+def tensor_dict_to_ndarray_dict(tensor_dict):
+    """Convert dict[str, tensor] to StrMap[NDArray]"""
+    ndarray_dict = {}
+    for key, value in tensor_dict.items():
+        ndarray_dict[key] = F.zerocopy_to_dgl_ndarray(value)
+    return convert_to_strmap(ndarray_dict)
+
+
+def save_heterographs(filename, g_list, labels, formats):
+    """Save heterographs into file"""
+    if labels is None:
+        labels = {}
+    if isinstance(g_list, DGLGraph):
+        g_list = [g_list]
+    assert all(
+        [type(g) == DGLGraph for g in g_list]
+    ), "Invalid DGLGraph in g_list argument"
+    gdata_list = [HeteroGraphData.create(g) for g in g_list]
+    if formats is None:
+        formats = []
+    elif isinstance(formats, str):
+        formats = [formats]
+    _CAPI_SaveHeteroGraphData(
+        filename, gdata_list, tensor_dict_to_ndarray_dict(labels), formats
+    )
+
+
+@register_object("heterograph_serialize.HeteroGraphData")
+class HeteroGraphData(ObjectBase):
+    """Object to hold the data to be stored for DGLGraph"""
+
+    @staticmethod
+    def create(g):
+        edata_list = []
+        ndata_list = []
+        for etype in g.canonical_etypes:
+            edata_list.append(tensor_dict_to_ndarray_dict(g.edges[etype].data))
+        for ntype in g.ntypes:
+            ndata_list.append(tensor_dict_to_ndarray_dict(g.nodes[ntype].data))
+        return _CAPI_MakeHeteroGraphData(
+            g._graph, ndata_list, edata_list, g.ntypes, g.etypes
+        )
+
+    def get_graph(self):
+        ntensor_list = list(_CAPI_GetNDataFromHeteroGraphData(self))
+        etensor_list = list(_CAPI_GetEDataFromHeteroGraphData(self))
+        ntype_names = list(_CAPI_GetNtypesFromHeteroGraphData(self))
+        etype_names = list(_CAPI_GetEtypesFromHeteroGraphData(self))
+        gidx = _CAPI_GetGindexFromHeteroGraphData(self)
+        nframes = []
+        eframes = []
+        for ntid, ntensor in enumerate(ntensor_list):
+            ndict = {
+                ntensor[i]: F.zerocopy_from_dgl_ndarray(ntensor[i + 1])
+                for i in range(0, len(ntensor), 2)
+            }
+            nframes.append(Frame(ndict, num_rows=gidx.num_nodes(ntid)))
+
+        for etid, etensor in enumerate(etensor_list):
+            edict = {
+                etensor[i]: F.zerocopy_from_dgl_ndarray(etensor[i + 1])
+                for i in range(0, len(etensor), 2)
+            }
+            eframes.append(Frame(edict, num_rows=gidx.num_edges(etid)))
+
+        return DGLGraph(gidx, ntype_names, etype_names, nframes, eframes)

python/dgl/data/heterophilous_graphs.py

@@ -0,0 +1,456 @@
+"""
+Datasets introduced in the 'A Critical Look at the Evaluation of GNNs under Heterophily: Are We
+Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+"""
+import os
+
+import numpy as np
+
+from ..convert import graph
+from ..transforms.functional import to_bidirected
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import download
+
+
+class HeterophilousGraphDataset(DGLBuiltinDataset):
+    r"""Datasets introduced in the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset. One of 'roman-empire', 'amazon-ratings', 'minesweeper', 'tolokers',
+        'questions'.
+    raw_dir : str
+        Raw file directory to store the processed data.
+    force_reload : bool
+        Whether to re-download the data source.
+    verbose : bool
+        Whether to print progress information.
+    transform : callable
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        name = name.lower().replace("-", "_")
+        url = f"https://github.com/yandex-research/heterophilous-graphs/raw/main/data/{name}.npz"
+        super(HeterophilousGraphDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        download(
+            url=self.url, path=os.path.join(self.raw_path, f"{self.name}.npz")
+        )
+
+    def process(self):
+        """Load and process the data."""
+        try:
+            import torch
+        except ImportError:
+            raise ModuleNotFoundError(
+                "This dataset requires PyTorch to be the backend."
+            )
+
+        data = np.load(os.path.join(self.raw_path, f"{self.name}.npz"))
+        src = torch.from_numpy(data["edges"][:, 0])
+        dst = torch.from_numpy(data["edges"][:, 1])
+        features = torch.from_numpy(data["node_features"])
+        labels = torch.from_numpy(data["node_labels"])
+        train_masks = torch.from_numpy(data["train_masks"].T)
+        val_masks = torch.from_numpy(data["val_masks"].T)
+        test_masks = torch.from_numpy(data["test_masks"].T)
+        num_nodes = len(labels)
+        num_classes = len(labels.unique())
+
+        self._num_classes = num_classes
+
+        self._g = to_bidirected(graph((src, dst), num_nodes=num_nodes))
+        self._g.ndata["feat"] = features
+        self._g.ndata["label"] = labels
+        self._g.ndata["train_mask"] = train_masks
+        self._g.ndata["val_mask"] = val_masks
+        self._g.ndata["test_mask"] = test_masks
+
+    def has_cache(self):
+        return os.path.exists(self.raw_path)
+
+    def load(self):
+        self.process()
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return self._num_classes
+
+
+class RomanEmpireDataset(HeterophilousGraphDataset):
+    r"""Roman-empire dataset from the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is based on the Roman Empire article from English Wikipedia, which was selected
+    since it is one of the longest articles on Wikipedia. Each node in the graph corresponds to one
+    (non-unique) word in the text. Thus, the number of nodes in the graph is equal to the article’s
+    length. Two words are connected with an edge if at least one of the following two conditions
+    holds: either these words follow each other in the text, or these words are connected in the
+    dependency tree of the sentence (one word depends on the other). Thus, the graph is a chain
+    graph with additional shortcut edges corresponding to syntactic dependencies between words. The
+    class of a node is its syntactic role (17 most frequent roles were selected as unique classes
+    and all the other roles were grouped into the 18th class). Node features are word embeddings.
+
+    Statistics:
+
+    - Nodes: 22662
+    - Edges: 65854
+    - Classes: 18
+    - Node features: 300
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import RomanEmpireDataset
+    >>> dataset = RomanEmpireDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(RomanEmpireDataset, self).__init__(
+            name="roman-empire",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class AmazonRatingsDataset(HeterophilousGraphDataset):
+    r"""Amazon-ratings dataset from the 'A Critical Look at the Evaluation of GNNs under
+    Heterophily: Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is based on the Amazon product co-purchasing data. Nodes are products (books, music
+    CDs, DVDs, VHS video tapes), and edges connect products that are frequently bought together. The
+    task is to predict the average rating given to a product by reviewers. All possible rating
+    values were grouped into five classes. Node features are the mean of word embeddings for words
+    in the product description.
+
+    Statistics:
+
+    - Nodes: 24492
+    - Edges: 186100
+    - Classes: 5
+    - Node features: 300
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import AmazonRatingsDataset
+    >>> dataset = AmazonRatingsDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(AmazonRatingsDataset, self).__init__(
+            name="amazon-ratings",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class MinesweeperDataset(HeterophilousGraphDataset):
+    r"""Minesweeper dataset from the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is inspired by the Minesweeper game. The graph is a regular 100x100 grid where each
+    node (cell) is connected to eight neighboring nodes (with the exception of nodes at the edge of
+    the grid, which have fewer neighbors). 20% of the nodes are randomly selected as mines. The task
+    is to predict which nodes are mines. The node features are one-hot-encoded numbers of
+    neighboring mines. However, for randomly selected 50% of the nodes, the features are unknown,
+    which is indicated by a separate binary feature.
+
+    Statistics:
+
+    - Nodes: 10000
+    - Edges: 78804
+    - Classes: 2
+    - Node features: 7
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import MinesweeperDataset
+    >>> dataset = MinesweeperDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(MinesweeperDataset, self).__init__(
+            name="minesweeper",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class TolokersDataset(HeterophilousGraphDataset):
+    r"""Tolokers dataset from the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is based on data from the Toloka crowdsourcing platform. The nodes represent
+    tolokers (workers). An edge connects two tolokers if they have worked on the same task. The goal
+    is to predict which tolokers have been banned in one of the projects. Node features are based on
+    the worker’s profile information and task performance statistics.
+
+    Statistics:
+
+    - Nodes: 11758
+    - Edges: 1038000
+    - Classes: 2
+    - Node features: 10
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import TolokersDataset
+    >>> dataset = TolokersDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(TolokersDataset, self).__init__(
+            name="tolokers",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class QuestionsDataset(HeterophilousGraphDataset):
+    r"""Questions dataset from the 'A Critical Look at the Evaluation of GNNs under Heterophily:
+    Are We Really Making Progress? <https://arxiv.org/abs/2302.11640>'__ paper.
+
+    This dataset is based on data from the question-answering website Yandex Q. Nodes are users, and
+    an edge connects two nodes if one user answered the other user’s question. The task is to
+    predict which users remained active on the website (were not deleted or blocked). Node features
+    are the mean of word embeddings for words in the user description. Users that do not have
+    description are indicated by a separate binary feature.
+
+    Statistics:
+
+    - Nodes: 48921
+    - Edges: 307080
+    - Classes: 2
+    - Node features: 301
+    - 10 train/val/test splits
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download the data source. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+
+    Examples
+    --------
+
+    >>> from dgl.data import QuestionsDataset
+    >>> dataset = QuestionsDataset()
+    >>> g = dataset[0]
+    >>> num_classes = dataset.num_classes
+
+    >>> # get node features
+    >>> feat = g.ndata["feat"]
+
+    >>> # get the first data split
+    >>> train_mask = g.ndata["train_mask"][:, 0]
+    >>> val_mask = g.ndata["val_mask"][:, 0]
+    >>> test_mask = g.ndata["test_mask"][:, 0]
+
+    >>> # get labels
+    >>> label = g.ndata['label']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(QuestionsDataset, self).__init__(
+            name="questions",
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )

python/dgl/data/icews18.py

@@ -0,0 +1,172 @@
+"""ICEWS18 dataset for temporal graph"""
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs, loadtxt, save_graphs
+
+
+class ICEWS18Dataset(DGLBuiltinDataset):
+    r"""ICEWS18 dataset for temporal graph
+
+    Integrated Crisis Early Warning System (ICEWS18)
+
+    Event data consists of coded interactions between socio-political
+    actors (i.e., cooperative or hostile actions between individuals,
+    groups, sectors and nation states). This Dataset consists of events
+    from 1/1/2018 to 10/31/2018 (24 hours time granularity).
+
+    Reference:
+
+        - `Recurrent Event Network for Reasoning over Temporal Knowledge Graphs <https://arxiv.org/abs/1904.05530>`_
+        - `ICEWS Coded Event Data <https://dataverse.harvard.edu/dataset.xhtml?persistentId=doi:10.7910/DVN/28075>`_
+
+    Statistics：
+
+    - Train examples: 240
+    - Valid examples: 30
+    - Test examples: 34
+    - Nodes per graph: 23033
+
+    Parameters
+    ----------
+    mode: str
+        Load train/valid/test data. Has to be one of ['train', 'valid', 'test']
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    -------
+    is_temporal : bool
+        Is the dataset contains temporal graphs
+
+    Examples
+    --------
+    >>> # get train, valid, test set
+    >>> train_data = ICEWS18Dataset()
+    >>> valid_data = ICEWS18Dataset(mode='valid')
+    >>> test_data = ICEWS18Dataset(mode='test')
+    >>>
+    >>> train_size = len(train_data)
+    >>> for g in train_data:
+    ....    e_feat = g.edata['rel_type']
+    ....    # your code here
+    ....
+    >>>
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        mode = mode.lower()
+        assert mode in ["train", "valid", "test"], "Mode not valid"
+        self.mode = mode
+        _url = _get_dgl_url("dataset/icews18.zip")
+        super(ICEWS18Dataset, self).__init__(
+            name="ICEWS18",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        data = loadtxt(
+            os.path.join(self.save_path, "{}.txt".format(self.mode)),
+            delimiter="\t",
+        ).astype(np.int64)
+        num_nodes = 23033
+        # The source code is not released, but the paper indicates there're
+        # totally 137 samples. The cutoff below has exactly 137 samples.
+        time_index = np.floor(data[:, 3] / 24).astype(np.int64)
+        start_time = time_index[time_index != -1].min()
+        end_time = time_index.max()
+        self._graphs = []
+        for i in range(start_time, end_time + 1):
+            row_mask = time_index <= i
+            edges = data[row_mask][:, [0, 2]]
+            rate = data[row_mask][:, 1]
+            g = dgl_graph((edges[:, 0], edges[:, 1]))
+            g.edata["rel_type"] = F.tensor(
+                rate.reshape(-1, 1), dtype=F.data_type_dict["int64"]
+            )
+            self._graphs.append(g)
+
+    def has_cache(self):
+        graph_path = os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.mode)
+        )
+        return os.path.exists(graph_path)
+
+    def save(self):
+        graph_path = os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.mode)
+        )
+        save_graphs(graph_path, self._graphs)
+
+    def load(self):
+        graph_path = os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.mode)
+        )
+        self._graphs = load_graphs(graph_path)[0]
+
+    def __getitem__(self, idx):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``edata['rel_type']``: edge type
+        """
+        if self._transform is None:
+            return self._graphs[idx]
+        else:
+            return self._transform(self._graphs[idx])
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return len(self._graphs)
+
+    @property
+    def is_temporal(self):
+        r"""Is the dataset contains temporal graphs
+
+        Returns
+        -------
+        bool
+        """
+        return True
+
+
+ICEWS18 = ICEWS18Dataset

python/dgl/data/karate.py

@@ -0,0 +1,98 @@
+"""KarateClub Dataset
+"""
+import networkx as nx
+import numpy as np
+
+from .. import backend as F
+from ..convert import from_networkx
+from .dgl_dataset import DGLDataset
+from .utils import deprecate_property
+
+__all__ = ["KarateClubDataset", "KarateClub"]
+
+
+class KarateClubDataset(DGLDataset):
+    r"""Karate Club dataset for Node Classification
+
+    Zachary's karate club is a social network of a university
+    karate club, described in the paper "An Information Flow
+    Model for Conflict and Fission in Small Groups" by Wayne W. Zachary.
+    The network became a popular example of community structure in
+    networks after its use by Michelle Girvan and Mark Newman in 2002.
+    Official website: `<http://konect.cc/networks/ucidata-zachary/>`_
+
+    Karate Club dataset statistics:
+
+    - Nodes: 34
+    - Edges: 156
+    - Number of Classes: 2
+
+    Parameters
+    ----------
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+    >>> dataset = KarateClubDataset()
+    >>> num_classes = dataset.num_classes
+    >>> g = dataset[0]
+    >>> labels = g.ndata['label']
+    """
+
+    def __init__(self, transform=None):
+        super(KarateClubDataset, self).__init__(
+            name="karate_club", transform=transform
+        )
+
+    def process(self):
+        kc_graph = nx.karate_club_graph()
+        label = np.asarray(
+            [kc_graph.nodes[i]["club"] != "Mr. Hi" for i in kc_graph.nodes]
+        ).astype(np.int64)
+        label = F.tensor(label)
+        g = from_networkx(kc_graph)
+        g.ndata["label"] = label
+        self._graph = g
+        self._data = [g]
+
+    @property
+    def num_classes(self):
+        """Number of classes."""
+        return 2
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index, KarateClubDataset has only one graph object
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            graph structure and labels.
+
+            - ``ndata['label']``: ground truth labels
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return 1
+
+
+KarateClub = KarateClubDataset

python/dgl/data/knowledge_graph.py

@@ -0,0 +1,779 @@
+from __future__ import absolute_import
+
+import os, sys
+import pickle as pkl
+
+import networkx as nx
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from ..utils import retry_method_with_fix
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    _get_dgl_url,
+    deprecate_function,
+    deprecate_property,
+    download,
+    extract_archive,
+    generate_mask_tensor,
+    get_download_dir,
+    load_graphs,
+    load_info,
+    makedirs,
+    save_graphs,
+    save_info,
+)
+
+
+class KnowledgeGraphDataset(DGLBuiltinDataset):
+    """KnowledgeGraph link prediction dataset
+
+    The dataset contains a graph depicting the connectivity of a knowledge
+    base. Currently, the knowledge bases from the
+    `RGCN paper <https://arxiv.org/pdf/1703.06103.pdf>`_ supported are
+    FB15k-237, FB15k, wn18
+
+    Parameters
+    -----------
+    name : str
+        Name can be 'FB15k-237', 'FB15k' or 'wn18'.
+    reverse : bool
+        Whether add reverse edges. Default: True.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    """
+
+    def __init__(
+        self,
+        name,
+        reverse=True,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self._name = name
+        self.reverse = reverse
+        url = _get_dgl_url("dataset/") + "{}.tgz".format(name)
+        super(KnowledgeGraphDataset, self).__init__(
+            name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        r"""Automatically download data and extract it."""
+        tgz_path = os.path.join(self.raw_dir, self.name + ".tgz")
+        download(self.url, path=tgz_path)
+        extract_archive(tgz_path, self.raw_path)
+
+    def process(self):
+        """
+        The original knowledge base is stored in triplets.
+        This function will parse these triplets and build the DGLGraph.
+        """
+        root_path = self.raw_path
+        entity_path = os.path.join(root_path, "entities.dict")
+        relation_path = os.path.join(root_path, "relations.dict")
+        train_path = os.path.join(root_path, "train.txt")
+        valid_path = os.path.join(root_path, "valid.txt")
+        test_path = os.path.join(root_path, "test.txt")
+        entity_dict = _read_dictionary(entity_path)
+        relation_dict = _read_dictionary(relation_path)
+        train = np.asarray(
+            _read_triplets_as_list(train_path, entity_dict, relation_dict)
+        )
+        valid = np.asarray(
+            _read_triplets_as_list(valid_path, entity_dict, relation_dict)
+        )
+        test = np.asarray(
+            _read_triplets_as_list(test_path, entity_dict, relation_dict)
+        )
+        num_nodes = len(entity_dict)
+        num_rels = len(relation_dict)
+        if self.verbose:
+            print("# entities: {}".format(num_nodes))
+            print("# relations: {}".format(num_rels))
+            print("# training edges: {}".format(train.shape[0]))
+            print("# validation edges: {}".format(valid.shape[0]))
+            print("# testing edges: {}".format(test.shape[0]))
+
+        # for compatability
+        self._train = train
+        self._valid = valid
+        self._test = test
+
+        self._num_nodes = num_nodes
+        self._num_rels = num_rels
+        # build graph
+        g, data = build_knowledge_graph(
+            num_nodes, num_rels, train, valid, test, reverse=self.reverse
+        )
+        (
+            etype,
+            ntype,
+            train_edge_mask,
+            valid_edge_mask,
+            test_edge_mask,
+            train_mask,
+            val_mask,
+            test_mask,
+        ) = data
+        g.edata["train_edge_mask"] = train_edge_mask
+        g.edata["valid_edge_mask"] = valid_edge_mask
+        g.edata["test_edge_mask"] = test_edge_mask
+        g.edata["train_mask"] = train_mask
+        g.edata["val_mask"] = val_mask
+        g.edata["test_mask"] = test_mask
+        g.edata["etype"] = etype
+        g.ndata["ntype"] = ntype
+        self._g = g
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.save_name + ".bin")
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, self.save_name + ".pkl")
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+
+        return False
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._g
+        else:
+            return self._transform(self._g)
+
+    def __len__(self):
+        return 1
+
+    def save(self):
+        """save the graph list and the labels"""
+        save_graphs(str(self.graph_path), self._g)
+        save_info(
+            str(self.info_path),
+            {"num_nodes": self.num_nodes, "num_rels": self.num_rels},
+        )
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+
+        info = load_info(str(self.info_path))
+        self._num_nodes = info["num_nodes"]
+        self._num_rels = info["num_rels"]
+        self._g = graphs[0]
+        train_mask = self._g.edata["train_edge_mask"].numpy()
+        val_mask = self._g.edata["valid_edge_mask"].numpy()
+        test_mask = self._g.edata["test_edge_mask"].numpy()
+
+        # convert mask tensor into bool tensor if possible
+        self._g.edata["train_edge_mask"] = generate_mask_tensor(
+            self._g.edata["train_edge_mask"].numpy()
+        )
+        self._g.edata["valid_edge_mask"] = generate_mask_tensor(
+            self._g.edata["valid_edge_mask"].numpy()
+        )
+        self._g.edata["test_edge_mask"] = generate_mask_tensor(
+            self._g.edata["test_edge_mask"].numpy()
+        )
+        self._g.edata["train_mask"] = generate_mask_tensor(
+            self._g.edata["train_mask"].numpy()
+        )
+        self._g.edata["val_mask"] = generate_mask_tensor(
+            self._g.edata["val_mask"].numpy()
+        )
+        self._g.edata["test_mask"] = generate_mask_tensor(
+            self._g.edata["test_mask"].numpy()
+        )
+
+        # for compatability (with 0.4.x) generate train_idx, valid_idx and test_idx
+        etype = self._g.edata["etype"].numpy()
+        self._etype = etype
+        u, v = self._g.all_edges(form="uv")
+        u = u.numpy()
+        v = v.numpy()
+        train_idx = np.nonzero(train_mask == 1)
+        self._train = np.column_stack(
+            (u[train_idx], etype[train_idx], v[train_idx])
+        )
+        valid_idx = np.nonzero(val_mask == 1)
+        self._valid = np.column_stack(
+            (u[valid_idx], etype[valid_idx], v[valid_idx])
+        )
+        test_idx = np.nonzero(test_mask == 1)
+        self._test = np.column_stack(
+            (u[test_idx], etype[test_idx], v[test_idx])
+        )
+
+        if self.verbose:
+            print("# entities: {}".format(self.num_nodes))
+            print("# relations: {}".format(self.num_rels))
+            print("# training edges: {}".format(self._train.shape[0]))
+            print("# validation edges: {}".format(self._valid.shape[0]))
+            print("# testing edges: {}".format(self._test.shape[0]))
+
+    @property
+    def num_nodes(self):
+        return self._num_nodes
+
+    @property
+    def num_rels(self):
+        return self._num_rels
+
+    @property
+    def save_name(self):
+        return self.name + "_dgl_graph"
+
+
+def _read_dictionary(filename):
+    d = {}
+    with open(filename, "r+") as f:
+        for line in f:
+            line = line.strip().split("\t")
+            d[line[1]] = int(line[0])
+    return d
+
+
+def _read_triplets(filename):
+    with open(filename, "r+") as f:
+        for line in f:
+            processed_line = line.strip().split("\t")
+            yield processed_line
+
+
+def _read_triplets_as_list(filename, entity_dict, relation_dict):
+    l = []
+    for triplet in _read_triplets(filename):
+        s = entity_dict[triplet[0]]
+        r = relation_dict[triplet[1]]
+        o = entity_dict[triplet[2]]
+        l.append([s, r, o])
+    return l
+
+
+def build_knowledge_graph(
+    num_nodes, num_rels, train, valid, test, reverse=True
+):
+    """Create a DGL Homogeneous graph with heterograph info stored as node or edge features."""
+    src = []
+    rel = []
+    dst = []
+    raw_subg = {}
+    raw_subg_eset = {}
+    raw_subg_etype = {}
+    raw_reverse_sugb = {}
+    raw_reverse_subg_eset = {}
+    raw_reverse_subg_etype = {}
+
+    # here there is noly one node type
+    s_type = "node"
+    d_type = "node"
+
+    def add_edge(s, r, d, reverse, edge_set):
+        r_type = str(r)
+        e_type = (s_type, r_type, d_type)
+        if raw_subg.get(e_type, None) is None:
+            raw_subg[e_type] = ([], [])
+            raw_subg_eset[e_type] = []
+            raw_subg_etype[e_type] = []
+        raw_subg[e_type][0].append(s)
+        raw_subg[e_type][1].append(d)
+        raw_subg_eset[e_type].append(edge_set)
+        raw_subg_etype[e_type].append(r)
+
+        if reverse is True:
+            r_type = str(r + num_rels)
+            re_type = (d_type, r_type, s_type)
+            if raw_reverse_sugb.get(re_type, None) is None:
+                raw_reverse_sugb[re_type] = ([], [])
+                raw_reverse_subg_etype[re_type] = []
+                raw_reverse_subg_eset[re_type] = []
+            raw_reverse_sugb[re_type][0].append(d)
+            raw_reverse_sugb[re_type][1].append(s)
+            raw_reverse_subg_eset[re_type].append(edge_set)
+            raw_reverse_subg_etype[re_type].append(r + num_rels)
+
+    for edge in train:
+        s, r, d = edge
+        assert r < num_rels
+        add_edge(s, r, d, reverse, 1)  # train set
+
+    for edge in valid:
+        s, r, d = edge
+        assert r < num_rels
+        add_edge(s, r, d, reverse, 2)  # valid set
+
+    for edge in test:
+        s, r, d = edge
+        assert r < num_rels
+        add_edge(s, r, d, reverse, 3)  # test set
+
+    subg = []
+    fg_s = []
+    fg_d = []
+    fg_etype = []
+    fg_settype = []
+    for e_type, val in raw_subg.items():
+        s, d = val
+        s = np.asarray(s)
+        d = np.asarray(d)
+        etype = raw_subg_etype[e_type]
+        etype = np.asarray(etype)
+        settype = raw_subg_eset[e_type]
+        settype = np.asarray(settype)
+
+        fg_s.append(s)
+        fg_d.append(d)
+        fg_etype.append(etype)
+        fg_settype.append(settype)
+
+    settype = np.concatenate(fg_settype)
+    if reverse is True:
+        settype = np.concatenate([settype, np.full((settype.shape[0]), 0)])
+    train_edge_mask = generate_mask_tensor(settype == 1)
+    valid_edge_mask = generate_mask_tensor(settype == 2)
+    test_edge_mask = generate_mask_tensor(settype == 3)
+
+    for e_type, val in raw_reverse_sugb.items():
+        s, d = val
+        s = np.asarray(s)
+        d = np.asarray(d)
+        etype = raw_reverse_subg_etype[e_type]
+        etype = np.asarray(etype)
+        settype = raw_reverse_subg_eset[e_type]
+        settype = np.asarray(settype)
+
+        fg_s.append(s)
+        fg_d.append(d)
+        fg_etype.append(etype)
+        fg_settype.append(settype)
+
+    s = np.concatenate(fg_s)
+    d = np.concatenate(fg_d)
+    g = dgl_graph((s, d), num_nodes=num_nodes)
+    etype = np.concatenate(fg_etype)
+    settype = np.concatenate(fg_settype)
+    etype = F.tensor(etype, dtype=F.data_type_dict["int64"])
+    train_edge_mask = train_edge_mask
+    valid_edge_mask = valid_edge_mask
+    test_edge_mask = test_edge_mask
+    train_mask = (
+        generate_mask_tensor(settype == 1)
+        if reverse is True
+        else train_edge_mask
+    )
+    valid_mask = (
+        generate_mask_tensor(settype == 2)
+        if reverse is True
+        else valid_edge_mask
+    )
+    test_mask = (
+        generate_mask_tensor(settype == 3)
+        if reverse is True
+        else test_edge_mask
+    )
+    ntype = F.full_1d(
+        num_nodes, 0, dtype=F.data_type_dict["int64"], ctx=F.cpu()
+    )
+
+    return g, (
+        etype,
+        ntype,
+        train_edge_mask,
+        valid_edge_mask,
+        test_edge_mask,
+        train_mask,
+        valid_mask,
+        test_mask,
+    )
+
+
+class FB15k237Dataset(KnowledgeGraphDataset):
+    r"""FB15k237 link prediction dataset.
+
+    FB15k-237 is a subset of FB15k where inverse
+    relations are removed. When creating the dataset,
+    a reverse edge with reversed relation types are
+    created for each edge by default.
+
+    FB15k237 dataset statistics:
+
+    - Nodes: 14541
+    - Number of relation types: 237
+    - Number of reversed relation types: 237
+    - Label Split:
+
+        - Train: 272115
+        - Valid: 17535
+        - Test: 20466
+
+    Parameters
+    ----------
+    reverse : bool
+        Whether to add reverse edge. Default True.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_nodes: int
+        Number of nodes
+    num_rels: int
+        Number of relation types
+
+    Examples
+    ----------
+    >>> dataset = FB15k237Dataset()
+    >>> g = dataset.graph
+    >>> e_type = g.edata['e_type']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.edata['train_mask']
+    >>> val_mask = g.edata['val_mask']
+    >>> test_mask = g.edata['test_mask']
+    >>>
+    >>> train_set = th.arange(g.num_edges())[train_mask]
+    >>> val_set = th.arange(g.num_edges())[val_mask]
+    >>>
+    >>> # build train_g
+    >>> train_edges = train_set
+    >>> train_g = g.edge_subgraph(train_edges,
+                                  relabel_nodes=False)
+    >>> train_g.edata['e_type'] = e_type[train_edges];
+    >>>
+    >>> # build val_g
+    >>> val_edges = th.cat([train_edges, val_edges])
+    >>> val_g = g.edge_subgraph(val_edges,
+                                relabel_nodes=False)
+    >>> val_g.edata['e_type'] = e_type[val_edges];
+    >>>
+    >>> # Train, Validation and Test
+    """
+
+    def __init__(
+        self,
+        reverse=True,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        name = "FB15k-237"
+        super(FB15k237Dataset, self).__init__(
+            name, reverse, raw_dir, force_reload, verbose, transform
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, FB15k237Dataset has only one graph object
+
+        Return
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains
+
+            - ``edata['e_type']``: edge relation type
+            - ``edata['train_edge_mask']``: positive training edge mask
+            - ``edata['val_edge_mask']``: positive validation edge mask
+            - ``edata['test_edge_mask']``: positive testing edge mask
+            - ``edata['train_mask']``: training edge set mask (include reversed training edges)
+            - ``edata['val_mask']``: validation edge set mask (include reversed validation edges)
+            - ``edata['test_mask']``: testing edge set mask (include reversed testing edges)
+            - ``ndata['ntype']``: node type. All 0 in this dataset
+        """
+        return super(FB15k237Dataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(FB15k237Dataset, self).__len__()
+
+
+class FB15kDataset(KnowledgeGraphDataset):
+    r"""FB15k link prediction dataset.
+
+    The FB15K dataset was introduced in `Translating Embeddings for Modeling
+    Multi-relational Data <http://papers.nips.cc/paper/5071-translating-embeddings-for-modeling-multi-relational-data.pdf>`_.
+    It is a subset of Freebase which contains about
+    14,951 entities with 1,345 different relations.
+    When creating the dataset, a reverse edge with
+    reversed relation types are created for each edge
+    by default.
+
+    FB15k dataset statistics:
+
+    - Nodes: 14,951
+    - Number of relation types: 1,345
+    - Number of reversed relation types: 1,345
+    - Label Split:
+
+        - Train: 483142
+        - Valid: 50000
+        - Test: 59071
+
+    Parameters
+    ----------
+    reverse : bool
+        Whether to add reverse edge. Default True.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_nodes: int
+        Number of nodes
+    num_rels: int
+        Number of relation types
+
+    Examples
+    ----------
+    >>> dataset = FB15kDataset()
+    >>> g = dataset.graph
+    >>> e_type = g.edata['e_type']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.edata['train_mask']
+    >>> val_mask = g.edata['val_mask']
+    >>>
+    >>> train_set = th.arange(g.num_edges())[train_mask]
+    >>> val_set = th.arange(g.num_edges())[val_mask]
+    >>>
+    >>> # build train_g
+    >>> train_edges = train_set
+    >>> train_g = g.edge_subgraph(train_edges,
+                                  relabel_nodes=False)
+    >>> train_g.edata['e_type'] = e_type[train_edges];
+    >>>
+    >>> # build val_g
+    >>> val_edges = th.cat([train_edges, val_edges])
+    >>> val_g = g.edge_subgraph(val_edges,
+                                relabel_nodes=False)
+    >>> val_g.edata['e_type'] = e_type[val_edges];
+    >>>
+    >>> # Train, Validation and Test
+    >>>
+    """
+
+    def __init__(
+        self,
+        reverse=True,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        name = "FB15k"
+        super(FB15kDataset, self).__init__(
+            name, reverse, raw_dir, force_reload, verbose, transform
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, FB15kDataset has only one graph object
+
+        Return
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains
+
+            - ``edata['e_type']``: edge relation type
+            - ``edata['train_edge_mask']``: positive training edge mask
+            - ``edata['val_edge_mask']``: positive validation edge mask
+            - ``edata['test_edge_mask']``: positive testing edge mask
+            - ``edata['train_mask']``: training edge set mask (include reversed training edges)
+            - ``edata['val_mask']``: validation edge set mask (include reversed validation edges)
+            - ``edata['test_mask']``: testing edge set mask (include reversed testing edges)
+            - ``ndata['ntype']``: node type. All 0 in this dataset
+        """
+        return super(FB15kDataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(FB15kDataset, self).__len__()
+
+
+class WN18Dataset(KnowledgeGraphDataset):
+    r"""WN18 link prediction dataset.
+
+    The WN18 dataset was introduced in `Translating Embeddings for Modeling
+    Multi-relational Data <http://papers.nips.cc/paper/5071-translating-embeddings-for-modeling-multi-relational-data.pdf>`_.
+    It included the full 18 relations scraped from
+    WordNet for roughly 41,000 synsets. When creating
+    the dataset, a reverse edge with reversed relation
+    types are created for each edge by default.
+
+    WN18 dataset statistics:
+
+    - Nodes: 40943
+    - Number of relation types: 18
+    - Number of reversed relation types: 18
+    - Label Split:
+
+        - Train: 141442
+        - Valid: 5000
+        - Test: 5000
+
+    Parameters
+    ----------
+    reverse : bool
+        Whether to add reverse edge. Default True.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_nodes: int
+        Number of nodes
+    num_rels: int
+        Number of relation types
+
+    Examples
+    ----------
+    >>> dataset = WN18Dataset()
+    >>> g = dataset.graph
+    >>> e_type = g.edata['e_type']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.edata['train_mask']
+    >>> val_mask = g.edata['val_mask']
+    >>>
+    >>> train_set = th.arange(g.num_edges())[train_mask]
+    >>> val_set = th.arange(g.num_edges())[val_mask]
+    >>>
+    >>> # build train_g
+    >>> train_edges = train_set
+    >>> train_g = g.edge_subgraph(train_edges,
+                                  relabel_nodes=False)
+    >>> train_g.edata['e_type'] = e_type[train_edges];
+    >>>
+    >>> # build val_g
+    >>> val_edges = th.cat([train_edges, val_edges])
+    >>> val_g = g.edge_subgraph(val_edges,
+                                relabel_nodes=False)
+    >>> val_g.edata['e_type'] = e_type[val_edges];
+    >>>
+    >>> # Train, Validation and Test
+    >>>
+    """
+
+    def __init__(
+        self,
+        reverse=True,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        name = "wn18"
+        super(WN18Dataset, self).__init__(
+            name, reverse, raw_dir, force_reload, verbose, transform
+        )
+
+    def __getitem__(self, idx):
+        r"""Gets the graph object
+
+        Parameters
+        -----------
+        idx: int
+            Item index, WN18Dataset has only one graph object
+
+        Return
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains
+
+            - ``edata['e_type']``: edge relation type
+            - ``edata['train_edge_mask']``: positive training edge mask
+            - ``edata['val_edge_mask']``: positive validation edge mask
+            - ``edata['test_edge_mask']``: positive testing edge mask
+            - ``edata['train_mask']``: training edge set mask (include reversed training edges)
+            - ``edata['val_mask']``: validation edge set mask (include reversed validation edges)
+            - ``edata['test_mask']``: testing edge set mask (include reversed testing edges)
+            - ``ndata['ntype']``: node type. All 0 in this dataset
+        """
+        return super(WN18Dataset, self).__getitem__(idx)
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return super(WN18Dataset, self).__len__()
+
+
+def load_data(dataset):
+    r"""Load knowledge graph dataset for RGCN link prediction tasks
+
+    It supports three datasets: wn18, FB15k and FB15k-237
+
+    Parameters
+    ----------
+    dataset: str
+        The name of the dataset to load.
+
+    Return
+    ------
+    The dataset object.
+    """
+    if dataset == "wn18":
+        return WN18Dataset()
+    elif dataset == "FB15k":
+        return FB15kDataset()
+    elif dataset == "FB15k-237":
+        return FB15k237Dataset()

python/dgl/data/minigc.py

@@ -0,0 +1,248 @@
+"""A mini synthetic dataset for graph classification benchmark."""
+import math
+import os
+
+import networkx as nx
+import numpy as np
+
+from .. import backend as F
+from ..convert import from_networkx
+from ..transforms import add_self_loop
+from .dgl_dataset import DGLDataset
+from .utils import load_graphs, makedirs, save_graphs
+
+__all__ = ["MiniGCDataset"]
+
+
+class MiniGCDataset(DGLDataset):
+    """The synthetic graph classification dataset class.
+
+    The datset contains 8 different types of graphs.
+
+    - class 0 : cycle graph
+    - class 1 : star graph
+    - class 2 : wheel graph
+    - class 3 : lollipop graph
+    - class 4 : hypercube graph
+    - class 5 : grid graph
+    - class 6 : clique graph
+    - class 7 : circular ladder graph
+
+    Parameters
+    ----------
+    num_graphs: int
+        Number of graphs in this dataset.
+    min_num_v: int
+        Minimum number of nodes for graphs
+    max_num_v: int
+        Maximum number of nodes for graphs
+    seed: int, default is 0
+        Random seed for data generation
+    transform: callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_graphs : int
+        Number of graphs
+    min_num_v : int
+        The minimum number of nodes
+    max_num_v : int
+        The maximum number of nodes
+    num_classes : int
+        The number of classes
+
+    Examples
+    --------
+    >>> data = MiniGCDataset(100, 16, 32, seed=0)
+
+    The dataset instance is an iterable
+
+    >>> len(data)
+    100
+    >>> g, label = data[64]
+    >>> g
+    Graph(num_nodes=20, num_edges=82,
+          ndata_schemes={}
+          edata_schemes={})
+    >>> label
+    tensor(5)
+
+    Batch the graphs and labels for mini-batch training
+
+    >>> graphs, labels = zip(*[data[i] for i in range(16)])
+    >>> batched_graphs = dgl.batch(graphs)
+    >>> batched_labels = torch.tensor(labels)
+    >>> batched_graphs
+    Graph(num_nodes=356, num_edges=1060,
+          ndata_schemes={}
+          edata_schemes={})
+    """
+
+    def __init__(
+        self,
+        num_graphs,
+        min_num_v,
+        max_num_v,
+        seed=0,
+        save_graph=True,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self.num_graphs = num_graphs
+        self.min_num_v = min_num_v
+        self.max_num_v = max_num_v
+        self.seed = seed
+        self.save_graph = save_graph
+
+        super(MiniGCDataset, self).__init__(
+            name="minigc",
+            hash_key=(num_graphs, min_num_v, max_num_v, seed),
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.graphs = []
+        self.labels = []
+        self._generate(self.seed)
+
+    def __len__(self):
+        """Return the number of graphs in the dataset."""
+        return len(self.graphs)
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (:class:`dgl.Graph`, Tensor)
+            The graph and its label.
+        """
+        if self._transform is None:
+            g = self.graphs[idx]
+        else:
+            g = self._transform(self.graphs[idx])
+        return g, self.labels[idx]
+
+    def has_cache(self):
+        graph_path = os.path.join(
+            self.save_path, "dgl_graph_{}.bin".format(self.hash)
+        )
+        if os.path.exists(graph_path):
+            return True
+
+        return False
+
+    def save(self):
+        """save the graph list and the labels"""
+        if self.save_graph:
+            graph_path = os.path.join(
+                self.save_path, "dgl_graph_{}.bin".format(self.hash)
+            )
+            save_graphs(str(graph_path), self.graphs, {"labels": self.labels})
+
+    def load(self):
+        graphs, label_dict = load_graphs(
+            os.path.join(self.save_path, "dgl_graph_{}.bin".format(self.hash))
+        )
+        self.graphs = graphs
+        self.labels = label_dict["labels"]
+
+    @property
+    def num_classes(self):
+        """Number of classes."""
+        return 8
+
+    def _generate(self, seed):
+        if seed is not None:
+            np.random.seed(seed)
+        self._gen_cycle(self.num_graphs // 8)
+        self._gen_star(self.num_graphs // 8)
+        self._gen_wheel(self.num_graphs // 8)
+        self._gen_lollipop(self.num_graphs // 8)
+        self._gen_hypercube(self.num_graphs // 8)
+        self._gen_grid(self.num_graphs // 8)
+        self._gen_clique(self.num_graphs // 8)
+        self._gen_circular_ladder(self.num_graphs - len(self.graphs))
+        # preprocess
+        for i in range(self.num_graphs):
+            # convert to DGLGraph, and add self loops
+            self.graphs[i] = add_self_loop(from_networkx(self.graphs[i]))
+        self.labels = F.tensor(np.array(self.labels).astype(np.int64))
+
+    def _gen_cycle(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.cycle_graph(num_v)
+            self.graphs.append(g)
+            self.labels.append(0)
+
+    def _gen_star(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            # nx.star_graph(N) gives a star graph with N+1 nodes
+            g = nx.star_graph(num_v - 1)
+            self.graphs.append(g)
+            self.labels.append(1)
+
+    def _gen_wheel(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.wheel_graph(num_v)
+            self.graphs.append(g)
+            self.labels.append(2)
+
+    def _gen_lollipop(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            path_len = np.random.randint(2, num_v // 2)
+            g = nx.lollipop_graph(m=num_v - path_len, n=path_len)
+            self.graphs.append(g)
+            self.labels.append(3)
+
+    def _gen_hypercube(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.hypercube_graph(int(math.log(num_v, 2)))
+            g = nx.convert_node_labels_to_integers(g)
+            self.graphs.append(g)
+            self.labels.append(4)
+
+    def _gen_grid(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            assert num_v >= 4, (
+                "We require a grid graph to contain at least two "
+                "rows and two columns, thus 4 nodes, got {:d} "
+                "nodes".format(num_v)
+            )
+            n_rows = np.random.randint(2, num_v // 2)
+            n_cols = num_v // n_rows
+            g = nx.grid_graph([n_rows, n_cols])
+            g = nx.convert_node_labels_to_integers(g)
+            self.graphs.append(g)
+            self.labels.append(5)
+
+    def _gen_clique(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.complete_graph(num_v)
+            self.graphs.append(g)
+            self.labels.append(6)
+
+    def _gen_circular_ladder(self, n):
+        for _ in range(n):
+            num_v = np.random.randint(self.min_num_v, self.max_num_v)
+            g = nx.circular_ladder_graph(num_v // 2)
+            self.graphs.append(g)
+            self.labels.append(7)

python/dgl/data/movielens.py

@@ -0,0 +1,646 @@
+"""MovieLens dataset"""
+import os
+
+import numpy as np
+import pandas as pd
+
+from torch import LongTensor, Tensor
+
+from ..base import dgl_warning
+from ..convert import heterograph
+from .dgl_dataset import DGLDataset
+
+from .utils import (
+    _get_dgl_url,
+    download,
+    extract_archive,
+    load_graphs,
+    load_info,
+    save_graphs,
+    save_info,
+    split_dataset,
+)
+
+GENRES_ML_100K = [
+    "unknown",
+    "Action",
+    "Adventure",
+    "Animation",
+    "Children",
+    "Comedy",
+    "Crime",
+    "Documentary",
+    "Drama",
+    "Fantasy",
+    "Film-Noir",
+    "Horror",
+    "Musical",
+    "Mystery",
+    "Romance",
+    "Sci-Fi",
+    "Thriller",
+    "War",
+    "Western",
+]
+GENRES_ML_1M = GENRES_ML_100K[1:]
+GENRES_ML_10M = GENRES_ML_100K + ["IMAX"]
+
+try:
+    import torch
+except ImportError:
+    HAS_TORCH = False
+else:
+    HAS_TORCH = True
+
+
+def check_pytorch():
+    """Check if PyTorch is the backend."""
+    if not HAS_TORCH:
+        raise ModuleNotFoundError(
+            "MovieLensDataset requires PyTorch to be the backend."
+        )
+
+
+class MovieLensDataset(DGLDataset):
+    r"""MovieLens dataset for edge prediction tasks. The raw datasets are extracted from
+    `MovieLens <https://grouplens.org/datasets/movielens/>`, introduced by
+    `Movielens unplugged: experiences with an occasionally connected recommender system <https://dl.acm.org/doi/10.1145/604045.604094>`.
+
+    The datasets consist of user ratings for movies and incorporate additional user/movie information in the form of features.
+    The nodes represent users and movies, and the edges store ratings that users assign to movies.
+
+    Statistics:
+
+    MovieLens-100K (ml-100k)
+
+    - Users: 943
+    - Movies: 1,682
+    - Ratings: 100,000 (1, 2, 3, 4, 5)
+
+    MovieLens-1M (ml-1m)
+
+    - Users: 6,040
+    - Movies: 3,706
+    - Ratings: 1,000,209 (1, 2, 3, 4, 5)
+
+    MovieLens-10M (ml-10m)
+
+    - Users: 69,878
+    - Movies: 10,677
+    - Ratings: 10,000,054 (0.5, 1, 1.5, ..., 4.5, 5.0)
+
+    Parameters
+    ----------
+    name: str
+        Dataset name. (:obj:`"ml-100k"`, :obj:`"ml-1m"`, :obj:`"ml-10m"`).
+    valid_ratio: int
+        Ratio of validation samples out of the whole dataset. Should be in (0.0, 1.0).
+    test_ratio: int, optional
+        Ratio of testing samples out of the whole dataset. Should be in (0.0, 1.0). And its sum with
+        :obj:`valid_ratio` should be in (0.0, 1.0) as well. This parameter is invalid
+        when :obj:`name` is :obj:`"ml-100k"`, since its testing samples are pre-specified.
+        Default: None
+    raw_dir : str, optional
+        Raw file directory to download/store the data.
+        Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to re-download(if the dataset has not been downloaded) and re-process the dataset.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    random_state : int, optional
+        Random seed used for random dataset split. Default: 0
+
+    Notes
+    -----
+    - When :obj:`name` is :obj:`"ml-100k"`, the :obj:`test_ratio` is invalid, and the training ratio is equal to 1-:obj:`valid_ratio`.
+    When :obj:`name` is :obj:`"ml-1m"` or :obj:`"ml-10m"`, the :obj:`test_ratio` is valid,
+    and the training ratio is equal to 1-:obj:`valid_ratio`-:obj:`test_ratio`.
+    - The number of edges is doubled to form an undirected(bidirected) graph structure.
+
+    Examples
+    --------
+    >>> from dgl.data import MovieLensDataset
+    >>> dataset = MovieLensDataset(name='ml-100k', valid_ratio=0.2)
+    >>> g = dataset[0]
+    >>> g
+    Graph(num_nodes={'movie': 1682, 'user': 943},
+          num_edges={('movie', 'movie-user', 'user'): 100000, ('user', 'user-movie', 'movie'): 100000},
+          metagraph=[('movie', 'user', 'movie-user'), ('user', 'movie', 'user-movie')])
+
+    >>> # get ratings of edges in the training graph.
+    >>> rate = g.edges['user-movie'].data['rate'] # or rate = g.edges['movie-user'].data['rate']
+    >>> rate
+    tensor([5., 5., 3.,  ..., 3., 3., 5.])
+
+    >>> # get train, valid and test mask of edges
+    >>> train_mask = g.edges['user-movie'].data['train_mask']
+    >>> valid_mask = g.edges['user-movie'].data['valid_mask']
+    >>> test_mask = g.edges['user-movie'].data['test_mask']
+
+    >>> # get train, valid and test ratings
+    >>> train_ratings = rate[train_mask]
+    >>> valid_ratings = rate[valid_mask]
+    >>> test_ratings = rate[test_mask]
+
+    >>> # get input features of users
+    >>> g.nodes["user"].data["feat"] # or g.nodes["movie"].data["feat"] for movie nodes
+    tensor([[0.4800, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
+            [1.0600, 1.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
+            [0.4600, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
+            ...,
+            [0.4000, 0.0000, 1.0000,  ..., 0.0000, 0.0000, 0.0000],
+            [0.9600, 1.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
+            [0.4400, 0.0000, 1.0000,  ..., 0.0000, 0.0000, 0.0000]])
+
+    """
+
+    _url = {
+        "ml-100k": "dataset/ml-100k.zip",
+        "ml-1m": "dataset/ml-1m.zip",
+        "ml-10m": "dataset/ml-10m.zip",
+    }
+
+    def __init__(
+        self,
+        name,
+        valid_ratio,
+        test_ratio=None,
+        raw_dir=None,
+        force_reload=None,
+        verbose=None,
+        transform=None,
+        random_state=0,
+    ):
+        check_pytorch()
+        assert name in [
+            "ml-100k",
+            "ml-1m",
+            "ml-10m",
+        ], f"currently movielens does not support {name}"
+
+        # test regarding valid and test split ratio
+        assert (
+            valid_ratio > 0.0 and valid_ratio < 1.0
+        ), f"valid_ratio {valid_ratio} must be in (0.0, 1.0)"
+
+        if name in ["ml-1m", "ml-10m"]:
+            assert (
+                test_ratio is not None and test_ratio > 0.0 and test_ratio < 1.0
+            ), f"test_ratio({test_ratio}) must be set to a value in (0.0, 1.0) when using ml-1m and ml-10m"
+            assert (
+                test_ratio + valid_ratio > 0.0
+                and test_ratio + valid_ratio < 1.0
+            ), f"test_ratio({test_ratio}) + valid_ratio({valid_ratio}) must be set to (0.0, 1.0) when using ml-1m and ml-10m"
+
+        if name == "ml-100k" and test_ratio is not None:
+            dgl_warning(
+                f"test_ratio ({test_ratio}) is not set to None for ml-100k. "
+                "Note that dataset split would not be affected by the test_ratio since "
+                "testing samples of ml-100k have been pre-specified."
+            )
+
+        self.valid_ratio = valid_ratio
+        self.test_ratio = test_ratio
+        self.random_state = random_state
+
+        if name == "ml-100k":
+            self.genres = GENRES_ML_100K
+        elif name == "ml-1m":
+            self.genres = GENRES_ML_1M
+        elif name == "ml-10m":
+            self.genres = GENRES_ML_10M
+        else:
+            raise NotImplementedError
+
+        super(MovieLensDataset, self).__init__(
+            name=name,
+            url=_get_dgl_url(self._url[name]),
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def check_version(self):
+        valid_ratio, test_ratio = load_info(self.version_path)
+        if self.valid_ratio == valid_ratio and (
+            self.test_ratio == test_ratio if self.name != "ml-100k" else True
+        ):
+            return True
+        else:
+            if self.name == "ml-100k":
+                print(
+                    f"The current valid ratio ({self.valid_ratio}) "
+                    "is not the same as the last setting "
+                    f"(valid: {valid_ratio}). "
+                    f"MovieLens {self.name} will be re-processed with the new dataset split setting."
+                )
+            else:
+                print(
+                    f"At least one of current valid ({self.valid_ratio}) and test ({self.test_ratio}) ratio "
+                    "are not the same as the last setting "
+                    f"(valid: {valid_ratio}, test: {test_ratio}). "
+                    f"MovieLens {self.name} will be re-processed with the new dataset split setting."
+                )
+            return False
+
+    def download(self):
+        zip_file_path = os.path.join(self.raw_dir, self.name + ".zip")
+        download(self.url, path=zip_file_path)
+        extract_archive(zip_file_path, self.raw_dir, overwrite=True)
+
+    def process(self):
+        print(f"Starting processing {self.name} ...")
+
+        # 0. loading movie features
+        movie_feat = load_info(
+            os.path.join(self.raw_path, "movie_feat.pkl")
+        ).to(torch.float)
+        # 1. dataset split: train + (valid + ) test
+        if self.name == "ml-100k":
+            train_rating_data = self._load_raw_rates(
+                os.path.join(self.raw_path, "u1.base"), "\t"
+            )
+            test_rating_data = self._load_raw_rates(
+                os.path.join(self.raw_path, "u1.test"), "\t"
+            )
+            indices = np.arange(len(train_rating_data))
+            train, valid, _ = split_dataset(
+                indices,
+                [1 - self.valid_ratio, self.valid_ratio, 0.0],
+                shuffle=True,
+                random_state=self.random_state,
+            )
+            train_rating_data, valid_rating_data = (
+                train_rating_data.iloc[train.indices],
+                train_rating_data.iloc[valid.indices],
+            )
+            all_rating_data = pd.concat(
+                [train_rating_data, valid_rating_data, test_rating_data]
+            )
+
+        elif self.name == "ml-1m" or self.name == "ml-10m":
+            all_rating_data = self._load_raw_rates(
+                os.path.join(self.raw_path, "ratings.dat"), "::"
+            )
+            indices = np.arange(len(all_rating_data))
+            train, valid, test = split_dataset(
+                indices,
+                [
+                    1 - self.valid_ratio - self.test_ratio,
+                    self.valid_ratio,
+                    self.test_ratio,
+                ],
+                shuffle=True,
+                random_state=self.random_state,
+            )
+            train_rating_data, valid_rating_data, test_rating_data = (
+                all_rating_data.iloc[train.indices],
+                all_rating_data.iloc[valid.indices],
+                all_rating_data.iloc[test.indices],
+            )
+
+        # 2. load user and movie data, and drop those unseen in rating_data
+        user_data = self._load_raw_user_data()
+        movie_data = self._load_raw_movie_data()
+        user_data = self._drop_unseen_nodes(
+            data_df=user_data,
+            col_name="id",
+            reserved_ids_set=set(all_rating_data["user_id"].values),
+        )
+        movie_data = self._drop_unseen_nodes(
+            data_df=movie_data,
+            col_name="id",
+            reserved_ids_set=set(all_rating_data["movie_id"].values),
+        )
+
+        user_feat = Tensor(self._process_user_feat(user_data))
+
+        # 3. generate rating pairs
+        # Map user/movie to the global id
+        self._global_user_id_map = {
+            ele: i for i, ele in enumerate(user_data["id"])
+        }
+        self._global_movie_id_map = {
+            ele: i for i, ele in enumerate(movie_data["id"])
+        }
+
+        # pair value is idx rather than id
+        u_indices, v_indices, labels = self._generate_pair_value(
+            all_rating_data
+        )
+        all_rating_pairs = (
+            LongTensor(u_indices),
+            LongTensor(v_indices),
+        )
+        all_rating_values = Tensor(labels)
+
+        graph = self.construct_g(
+            all_rating_pairs, all_rating_values, user_feat, movie_feat
+        )
+        self.graph = self.add_masks(
+            graph, train_rating_data, valid_rating_data, test_rating_data
+        )
+
+        print(f"End processing {self.name} ...")
+
+    def construct_g(self, rate_pairs, rate_values, user_feat, movie_feat):
+        g = heterograph(
+            {
+                ("user", "user-movie", "movie"): (rate_pairs[0], rate_pairs[1]),
+                ("movie", "movie-user", "user"): (rate_pairs[1], rate_pairs[0]),
+            }
+        )
+        ndata = {"user": user_feat, "movie": movie_feat}
+        edata = {"user-movie": rate_values, "movie-user": rate_values}
+        g.ndata["feat"] = ndata
+        g.edata["rate"] = edata
+        return g
+
+    def add_masks(
+        self, g, train_rating_data, valid_rating_data, test_rating_data
+    ):
+        train_u_indices, train_v_indices, _ = self._generate_pair_value(
+            train_rating_data
+        )
+        valid_u_indices, valid_v_indices, _ = self._generate_pair_value(
+            valid_rating_data
+        )
+        test_u_indices, test_v_indices, _ = self._generate_pair_value(
+            test_rating_data
+        )
+
+        # user-movie
+        train_mask = torch.zeros((g.num_edges("user-movie"),), dtype=torch.bool)
+        train_mask[
+            g.edge_ids(train_u_indices, train_v_indices, etype="user-movie")
+        ] = True
+        valid_mask = torch.zeros((g.num_edges("user-movie"),), dtype=torch.bool)
+        valid_mask[
+            g.edge_ids(valid_u_indices, valid_v_indices, etype="user-movie")
+        ] = True
+        test_mask = torch.zeros((g.num_edges("user-movie"),), dtype=torch.bool)
+        test_mask[
+            g.edge_ids(test_u_indices, test_v_indices, etype="user-movie")
+        ] = True
+
+        g.edges["user-movie"].data["train_mask"] = train_mask
+        g.edges["user-movie"].data["valid_mask"] = valid_mask
+        g.edges["user-movie"].data["test_mask"] = test_mask
+
+        # movie-user
+        train_mask_rev = torch.zeros(
+            (g.num_edges("movie-user"),), dtype=torch.bool
+        )
+        train_mask_rev[
+            g.edge_ids(train_v_indices, train_u_indices, etype="movie-user")
+        ] = True
+        valid_mask_rev = torch.zeros(
+            (g.num_edges("movie-user"),), dtype=torch.bool
+        )
+        valid_mask_rev[
+            g.edge_ids(valid_v_indices, valid_u_indices, etype="movie-user")
+        ] = True
+        test_mask_rev = torch.zeros(
+            (g.num_edges("movie-user"),), dtype=torch.bool
+        )
+        test_mask_rev[
+            g.edge_ids(test_v_indices, test_u_indices, etype="movie-user")
+        ] = True
+
+        g.edges["movie-user"].data["train_mask"] = train_mask_rev
+        g.edges["movie-user"].data["valid_mask"] = valid_mask_rev
+        g.edges["movie-user"].data["test_mask"] = test_mask_rev
+
+        return g
+
+    def has_cache(self):
+        if (
+            os.path.exists(self.graph_path)
+            and os.path.exists(self.version_path)
+            and self.check_version()
+        ):
+            return True
+        return False
+
+    def save(self):
+        save_graphs(self.graph_path, [self.graph])
+        save_info(self.version_path, [self.valid_ratio, self.test_ratio])
+        if self.verbose:
+            print(f"Done saving data into {self.raw_path}.")
+
+    def load(self):
+        g_list, _ = load_graphs(self.graph_path)
+        self.graph = g_list[0]
+
+        """
+        To avoid the problem each time loading boolean tensor from the disk, boolean values
+        would be automatically converted into torch.uint8 types, and a deprecation warning would
+        be raised for using torch.uint8
+        """
+        for e in self.graph.etypes:
+            self.graph.edges[e].data["train_mask"] = (
+                self.graph.edges[e].data["train_mask"].to(torch.bool)
+            )
+            self.graph.edges[e].data["valid_mask"] = (
+                self.graph.edges[e].data["valid_mask"].to(torch.bool)
+            )
+            self.graph.edges[e].data["test_mask"] = (
+                self.graph.edges[e].data["test_mask"].to(torch.bool)
+            )
+
+    def __getitem__(self, idx):
+        assert (
+            idx == 0
+        ), "This dataset has only one set of training, validation and testing graph"
+        if self._transform is None:
+            return self.graph
+        else:
+            return self._transform(self.graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def raw_path(self):
+        return os.path.join(self.raw_dir, self.name)
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.raw_path, self.name + ".bin")
+
+    @property
+    def version_path(self):
+        return os.path.join(self.raw_path, self.name + "_version.pkl")
+
+    def _process_user_feat(self, user_data):
+        if self.name == "ml-100k" or self.name == "ml-1m":
+            ages = user_data["age"].values.astype(np.float32)
+            gender = (user_data["gender"] == "F").values.astype(np.float32)
+            all_occupations = set(user_data["occupation"])
+            occupation_map = {ele: i for i, ele in enumerate(all_occupations)}
+            occupation_one_hot = np.zeros(
+                shape=(user_data.shape[0], len(all_occupations)),
+                dtype=np.float32,
+            )
+            occupation_one_hot[
+                np.arange(user_data.shape[0]),
+                np.array(
+                    [occupation_map[ele] for ele in user_data["occupation"]]
+                ),
+            ] = 1
+            user_features = np.concatenate(
+                [
+                    ages.reshape((user_data.shape[0], 1)) / 50.0,
+                    gender.reshape((user_data.shape[0], 1)),
+                    occupation_one_hot,
+                ],
+                axis=1,
+            )
+        elif self.name == "ml-10m":
+            user_features = np.zeros(
+                shape=(user_data.shape[0], 1), dtype=np.float32
+            )
+        else:
+            raise NotImplementedError
+        return user_features
+
+    def _load_raw_user_data(self):
+        if self.name == "ml-100k":
+            user_data = pd.read_csv(
+                os.path.join(self.raw_path, "u.user"),
+                sep="|",
+                header=None,
+                names=["id", "age", "gender", "occupation", "zip_code"],
+                engine="python",
+            )
+        elif self.name == "ml-1m":
+            user_data = pd.read_csv(
+                os.path.join(self.raw_path, "users.dat"),
+                sep="::",
+                header=None,
+                names=["id", "gender", "age", "occupation", "zip_code"],
+                engine="python",
+            )
+        elif self.name == "ml-10m":
+            rating_info = pd.read_csv(
+                os.path.join(self.raw_path, "ratings.dat"),
+                sep="::",
+                header=None,
+                names=["user_id", "movie_id", "rating", "timestamp"],
+                dtype={
+                    "user_id": np.int32,
+                    "movie_id": np.int32,
+                    "ratings": np.float32,
+                    "timestamp": np.int64,
+                },
+                engine="python",
+            )
+            user_data = pd.DataFrame(
+                np.unique(rating_info["user_id"].values.astype(np.int32)),
+                columns=["id"],
+            )
+        else:
+            raise NotImplementedError
+        return user_data
+
+    def _load_raw_movie_data(self):
+        file_path = os.path.join(self.raw_path, "u.item")
+        if self.name == "ml-100k":
+            movie_data = pd.read_csv(
+                file_path,
+                sep="|",
+                header=None,
+                names=[
+                    "id",
+                    "title",
+                    "release_date",
+                    "video_release_date",
+                    "url",
+                ]
+                + GENRES_ML_100K,
+                engine="python",
+                encoding="ISO-8859-1",
+            )
+        elif self.name == "ml-1m" or self.name == "ml-10m":
+            file_path = os.path.join(self.raw_path, "movies.dat")
+            movie_data = pd.read_csv(
+                file_path,
+                sep="::",
+                header=None,
+                names=["id", "title", "genres"],
+                encoding="iso-8859-1",
+                engine="python",
+            )
+            genre_map = {ele: i for i, ele in enumerate(self.genres)}
+            genre_map["Children's"] = genre_map["Children"]
+            genre_map["Childrens"] = genre_map["Children"]
+            movie_genres = np.zeros(
+                shape=(movie_data.shape[0], len(self.genres)), dtype=np.float32
+            )
+            for i, genres in enumerate(movie_data["genres"]):
+                for ele in genres.split("|"):
+                    if ele in genre_map:
+                        movie_genres[i, genre_map[ele]] = 1.0
+                    else:
+                        movie_genres[i, genre_map["unknown"]] = 1.0
+            for idx, genre_name in enumerate(self.genres):
+                movie_data[genre_name] = movie_genres[:, idx]
+            movie_data = movie_data.drop(columns=["genres"])
+        else:
+            raise NotImplementedError
+
+        return movie_data
+
+    def _load_raw_rates(self, file_path, sep):
+        rating_data = pd.read_csv(
+            file_path,
+            sep=sep,
+            header=None,
+            names=["user_id", "movie_id", "rating", "timestamp"],
+            dtype={
+                "user_id": np.int32,
+                "movie_id": np.int32,
+                "ratings": np.float32,
+                "timestamp": np.int64,
+            },
+            engine="python",
+        )
+        rating_data = rating_data.reset_index(drop=True)
+        return rating_data
+
+    def _drop_unseen_nodes(self, data_df, col_name, reserved_ids_set):
+        data_df = data_df[data_df[col_name].isin(reserved_ids_set)]
+        data_df.reset_index(drop=True, inplace=True)
+        return data_df
+
+    def _generate_pair_value(self, rating_data):
+        rating_pairs = (
+            np.array(
+                [
+                    self._global_user_id_map[ele]
+                    for ele in rating_data["user_id"]
+                ],
+                dtype=np.int32,
+            ),
+            np.array(
+                [
+                    self._global_movie_id_map[ele]
+                    for ele in rating_data["movie_id"]
+                ],
+                dtype=np.int32,
+            ),
+        )
+        rating_values = rating_data["rating"].values.astype(np.float32)
+        return rating_pairs[0], rating_pairs[1], rating_values
+
+    def __repr__(self):
+        return (
+            f'Dataset("{self.name}", num_graphs={len(self)},'
+            + f" save_path={self.raw_path}), valid_ratio={self.valid_ratio}, test_ratio={self.test_ratio}"
+        )

python/dgl/data/pattern.py

@@ -0,0 +1,130 @@
+""" PATTERNDataset for inductive learning. """
+import os
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs
+
+
+class PATTERNDataset(DGLBuiltinDataset):
+    r"""PATTERN dataset for graph pattern recognition task.
+
+    Each graph G contains 5 communities with sizes randomly selected between [5, 35].
+    The SBM of each community is p = 0.5, q = 0.35, and the node features on G are
+    generated with a uniform random distribution with a vocabulary of size 3, i.e. {0, 1, 2}.
+    Then randomly generate 100 patterns P composed of 20 nodes with intra-probability :math:`p_P` = 0.5
+    and extra-probability :math:`q_P` = 0.5 (i.e. 50% of nodes in P are connected to G). The node features
+    for P are also generated as a random signal with values {0, 1, 2}. The graphs are of sizes
+    44-188 nodes. The output node labels have value 1 if the node belongs to P and value 0 if it is in G.
+
+    Reference `<https://arxiv.org/pdf/2003.00982.pdf>`_
+
+    Statistics:
+
+    - Train examples: 10,000
+    - Valid examples: 2,000
+    - Test examples: 2,000
+    - Number of classes for each node: 2
+
+    Parameters
+    ----------
+    mode : str
+        Must be one of ('train', 'valid', 'test').
+        Default: 'train'
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node.
+
+    Examples
+    --------
+    >>> from dgl.data import PATTERNDataset
+    >>> data = PATTERNDataset(mode='train')
+    >>> data.num_classes
+    2
+    >>> len(trainset)
+    10000
+    >>> data[0]
+    Graph(num_nodes=108, num_edges=4884, ndata_schemes={'feat': Scheme(shape=(), dtype=torch.int64), 'label': Scheme(shape=(), dtype=torch.int16)}
+    edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)})
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        assert mode in ["train", "valid", "test"]
+        self.mode = mode
+        _url = _get_dgl_url("dataset/SBM_PATTERN.zip")
+
+        super(PATTERNDataset, self).__init__(
+            name="pattern",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.load()
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "SBM_PATTERN_{}.bin".format(self.mode)
+        )
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        self._graphs, _ = load_graphs(self.graph_path)
+
+    @property
+    def num_classes(self):
+        r"""Number of classes for each node."""
+        return 2
+
+    def __len__(self):
+        r"""The number of examples in the dataset."""
+        return len(self._graphs)
+
+    def __getitem__(self, idx):
+        r"""Get the idx^th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node features, node labels and edge features.
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+            - ``edata['feat']``: edge features
+        """
+        if self._transform is None:
+            return self._graphs[idx]
+        else:
+            return self._transform(self._graphs[idx])

python/dgl/data/ppi.py

@@ -0,0 +1,226 @@
+""" PPIDataset for inductive learning. """
+import json
+import os
+
+import networkx as nx
+import numpy as np
+from networkx.readwrite import json_graph
+
+from .. import backend as F
+from ..convert import from_networkx
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs, load_info, save_graphs, save_info
+
+
+class PPIDataset(DGLBuiltinDataset):
+    r"""Protein-Protein Interaction dataset for inductive node classification
+
+    A toy Protein-Protein Interaction network dataset. The dataset contains
+    24 graphs. The average number of nodes per graph is 2372. Each node has
+    50 features and 121 labels. 20 graphs for training, 2 for validation
+    and 2 for testing.
+
+    Reference: `<http://snap.stanford.edu/graphsage/>`_
+
+    Statistics:
+
+    - Train examples: 20
+    - Valid examples: 2
+    - Test examples: 2
+
+    Parameters
+    ----------
+    mode : str
+        Must be one of ('train', 'valid', 'test').
+        Default: 'train'
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_labels : int
+        Number of labels for each node
+    labels : Tensor
+        Node labels
+    features : Tensor
+        Node features
+
+    Examples
+    --------
+    >>> dataset = PPIDataset(mode='valid')
+    >>> num_classes = dataset.num_classes
+    >>> for g in dataset:
+    ....    feat = g.ndata['feat']
+    ....    label = g.ndata['label']
+    ....    # your code here
+    >>>
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        assert mode in ["train", "valid", "test"]
+        self.mode = mode
+        _url = _get_dgl_url("dataset/ppi.zip")
+        super(PPIDataset, self).__init__(
+            name="ppi",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        graph_file = os.path.join(
+            self.save_path, "{}_graph.json".format(self.mode)
+        )
+        label_file = os.path.join(
+            self.save_path, "{}_labels.npy".format(self.mode)
+        )
+        feat_file = os.path.join(
+            self.save_path, "{}_feats.npy".format(self.mode)
+        )
+        graph_id_file = os.path.join(
+            self.save_path, "{}_graph_id.npy".format(self.mode)
+        )
+
+        g_data = json.load(open(graph_file))
+        self._labels = np.load(label_file)
+        self._feats = np.load(feat_file)
+        self.graph = from_networkx(
+            nx.DiGraph(json_graph.node_link_graph(g_data))
+        )
+        graph_id = np.load(graph_id_file)
+
+        # lo, hi means the range of graph ids for different portion of the dataset,
+        # 20 graphs for training, 2 for validation and 2 for testing.
+        lo, hi = 1, 21
+        if self.mode == "valid":
+            lo, hi = 21, 23
+        elif self.mode == "test":
+            lo, hi = 23, 25
+
+        graph_masks = []
+        self.graphs = []
+        for g_id in range(lo, hi):
+            g_mask = np.where(graph_id == g_id)[0]
+            graph_masks.append(g_mask)
+            g = self.graph.subgraph(g_mask)
+            g.ndata["feat"] = F.tensor(
+                self._feats[g_mask], dtype=F.data_type_dict["float32"]
+            )
+            g.ndata["label"] = F.tensor(
+                self._labels[g_mask], dtype=F.data_type_dict["float32"]
+            )
+            self.graphs.append(g)
+
+    @property
+    def graph_list_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph_list.bin".format(self.mode)
+        )
+
+    @property
+    def g_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.mode)
+        )
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, "{}_info.pkl".format(self.mode))
+
+    def has_cache(self):
+        return (
+            os.path.exists(self.graph_list_path)
+            and os.path.exists(self.g_path)
+            and os.path.exists(self.info_path)
+        )
+
+    def save(self):
+        save_graphs(self.graph_list_path, self.graphs)
+        save_graphs(self.g_path, self.graph)
+        save_info(
+            self.info_path, {"labels": self._labels, "feats": self._feats}
+        )
+
+    def load(self):
+        self.graphs = load_graphs(self.graph_list_path)[0]
+        g, _ = load_graphs(self.g_path)
+        self.graph = g[0]
+        info = load_info(self.info_path)
+        self._labels = info["labels"]
+        self._feats = info["feats"]
+
+    @property
+    def num_labels(self):
+        return 121
+
+    @property
+    def num_classes(self):
+        return 121
+
+    def __len__(self):
+        """Return number of samples in this dataset."""
+        return len(self.graphs)
+
+    def __getitem__(self, item):
+        """Get the item^th sample.
+
+        Parameters
+        ---------
+        item : int
+            The sample index.
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node features and node labels.
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+        """
+        if self._transform is None:
+            return self.graphs[item]
+        else:
+            return self._transform(self.graphs[item])
+
+
+class LegacyPPIDataset(PPIDataset):
+    """Legacy version of PPI Dataset"""
+
+    def __getitem__(self, item):
+        """Get the item^th sample.
+
+        Paramters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (dgl.DGLGraph, Tensor, Tensor)
+            The graph, features and its label.
+        """
+        if self._transform is None:
+            g = self.graphs[item]
+        else:
+            g = self._transform(self.graphs[item])
+        return g, g.ndata["feat"], g.ndata["label"]

python/dgl/data/qm7b.py

@@ -0,0 +1,177 @@
+"""QM7b dataset for graph property prediction (regression)."""
+import os
+
+from scipy import io
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLDataset
+from .utils import check_sha1, download, load_graphs, save_graphs
+
+
+class QM7bDataset(DGLDataset):
+    r"""QM7b dataset for graph property prediction (regression)
+
+    This dataset consists of 7,211 molecules with 14 regression targets.
+    Nodes means atoms and edges means bonds. Edge data 'h' means
+    the entry of Coulomb matrix.
+
+    Reference: `<http://quantum-machine.org/datasets/>`_
+
+    Statistics:
+
+    - Number of graphs: 7,211
+    - Number of regression targets: 14
+    - Average number of nodes: 15
+    - Average number of edges: 245
+    - Edge feature size: 1
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_tasks : int
+        Number of prediction tasks
+    num_labels : int
+        (DEPRECATED, use num_tasks instead) Number of prediction tasks
+
+    Raises
+    ------
+    UserWarning
+        If the raw data is changed in the remote server by the author.
+
+    Examples
+    --------
+    >>> data = QM7bDataset()
+    >>> data.num_tasks
+    14
+    >>>
+    >>> # iterate over the dataset
+    >>> for g, label in data:
+    ...     edge_feat = g.edata['h']  # get edge feature
+    ...     # your code here...
+    ...
+    >>>
+    """
+
+    _url = (
+        "http://deepchem.io.s3-website-us-west-1.amazonaws.com/"
+        "datasets/qm7b.mat"
+    )
+    _sha1_str = "4102c744bb9d6fd7b40ac67a300e49cd87e28392"
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        super(QM7bDataset, self).__init__(
+            name="qm7b",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        mat_path = os.path.join(self.raw_dir, self.name + ".mat")
+        self.graphs, self.label = self._load_graph(mat_path)
+
+    def _load_graph(self, filename):
+        data = io.loadmat(filename)
+        labels = F.tensor(data["T"], dtype=F.data_type_dict["float32"])
+        feats = data["X"]
+        num_graphs = labels.shape[0]
+        graphs = []
+        for i in range(num_graphs):
+            edge_list = feats[i].nonzero()
+            g = dgl_graph(edge_list)
+            g.edata["h"] = F.tensor(
+                feats[i][edge_list[0], edge_list[1]].reshape(-1, 1),
+                dtype=F.data_type_dict["float32"],
+            )
+            graphs.append(g)
+        return graphs, labels
+
+    def save(self):
+        """save the graph list and the labels"""
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(str(graph_path), self.graphs, {"labels": self.label})
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        return os.path.exists(graph_path)
+
+    def load(self):
+        graphs, label_dict = load_graphs(
+            os.path.join(self.save_path, "dgl_graph.bin")
+        )
+        self.graphs = graphs
+        self.label = label_dict["labels"]
+
+    def download(self):
+        file_path = os.path.join(self.raw_dir, self.name + ".mat")
+        download(self.url, path=file_path)
+        if not check_sha1(file_path, self._sha1_str):
+            raise UserWarning(
+                "File {} is downloaded but the content hash does not match."
+                "The repo may be outdated or download may be incomplete. "
+                "Otherwise you can create an issue for it.".format(self.name)
+            )
+
+    @property
+    def num_tasks(self):
+        """Number of prediction tasks."""
+        return self.num_labels
+
+    @property
+    def num_labels(self):
+        """Number of prediction tasks."""
+        return 14
+
+    @property
+    def num_classes(self):
+        """Number of prediction tasks."""
+        return 14
+
+    def __getitem__(self, idx):
+        r"""Get graph and label by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+        """
+        if self._transform is None:
+            g = self.graphs[idx]
+        else:
+            g = self._transform(self.graphs[idx])
+        return g, self.label[idx]
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return len(self.graphs)
+
+
+QM7b = QM7bDataset

python/dgl/data/qm9.py

@@ -0,0 +1,231 @@
+"""QM9 dataset for graph property prediction (regression)."""
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+from ..transforms import to_bidirected
+
+from .dgl_dataset import DGLDataset
+from .utils import _get_dgl_url, download
+
+
+class QM9Dataset(DGLDataset):
+    r"""QM9 dataset for graph property prediction (regression)
+
+    This dataset consists of 130,831 molecules with 12 regression targets.
+    Nodes correspond to atoms and edges correspond to close atom pairs.
+
+    This dataset differs from :class:`~dgl.data.QM9EdgeDataset` in the following aspects:
+        1. Edges in this dataset are purely distance-based.
+        2. It only provides atoms' coordinates and atomic numbers as node features
+        3. It only provides 12 regression targets.
+
+    Reference:
+
+    - `"Quantum-Machine.org" <http://quantum-machine.org/datasets/>`_,
+    - `"Directional Message Passing for Molecular Graphs" <https://arxiv.org/abs/2003.03123>`_
+
+    Statistics:
+
+    - Number of graphs: 130,831
+    - Number of regression targets: 12
+
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | Keys   | Property                         | Description                                                                       | Unit                                        |
+    +========+==================================+===================================================================================+=============================================+
+    | mu     | :math:`\mu`                      | Dipole moment                                                                     | :math:`\textrm{D}`                          |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | alpha  | :math:`\alpha`                   | Isotropic polarizability                                                          | :math:`{a_0}^3`                             |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | homo   | :math:`\epsilon_{\textrm{HOMO}}` | Highest occupied molecular orbital energy                                         | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | lumo   | :math:`\epsilon_{\textrm{LUMO}}` | Lowest unoccupied molecular orbital energy                                        | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | gap    | :math:`\Delta \epsilon`          | Gap between :math:`\epsilon_{\textrm{HOMO}}` and :math:`\epsilon_{\textrm{LUMO}}` | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | r2     | :math:`\langle R^2 \rangle`      | Electronic spatial extent                                                         | :math:`{a_0}^2`                             |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | zpve   | :math:`\textrm{ZPVE}`            | Zero point vibrational energy                                                     | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U0     | :math:`U_0`                      | Internal energy at 0K                                                             | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U      | :math:`U`                        | Internal energy at 298.15K                                                        | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | H      | :math:`H`                        | Enthalpy at 298.15K                                                               | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | G      | :math:`G`                        | Free energy at 298.15K                                                            | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | Cv     | :math:`c_{\textrm{v}}`           | Heat capavity at 298.15K                                                          | :math:`\frac{\textrm{cal}}{\textrm{mol K}}` |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+
+    Parameters
+    ----------
+    label_keys : list
+        Names of the regression property, which should be a subset of the keys in the table above.
+    cutoff : float
+        Cutoff distance for interatomic interactions, i.e. two atoms are connected in the corresponding graph if the distance between them is no larger than this.
+        Default: 5.0 Angstrom
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_tasks : int
+        Number of prediction tasks
+    num_labels : int
+        (DEPRECATED, use num_tasks instead) Number of prediction tasks
+
+    Raises
+    ------
+    UserWarning
+        If the raw data is changed in the remote server by the author.
+
+    Examples
+    --------
+    >>> data = QM9Dataset(label_keys=['mu', 'gap'], cutoff=5.0)
+    >>> data.num_tasks
+    2
+    >>>
+    >>> # iterate over the dataset
+    >>> for g, label in data:
+    ...     R = g.ndata['R'] # get coordinates of each atom
+    ...     Z = g.ndata['Z'] # get atomic numbers of each atom
+    ...     # your code here...
+    >>>
+    """
+
+    def __init__(
+        self,
+        label_keys,
+        cutoff=5.0,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self.cutoff = cutoff
+        self.label_keys = label_keys
+        self._url = _get_dgl_url("dataset/qm9_eV.npz")
+
+        super(QM9Dataset, self).__init__(
+            name="qm9",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        npz_path = f"{self.raw_dir}/qm9_eV.npz"
+        data_dict = np.load(npz_path, allow_pickle=True)
+        # data_dict['N'] contains the number of atoms in each molecule.
+        # Atomic properties (Z and R) of all molecules are concatenated as single tensors,
+        # so you need this value to select the correct atoms for each molecule.
+        self.N = data_dict["N"]
+        self.R = data_dict["R"]
+        self.Z = data_dict["Z"]
+        self.label = np.stack(
+            [data_dict[key] for key in self.label_keys], axis=1
+        )
+        self.N_cumsum = np.concatenate([[0], np.cumsum(self.N)])
+
+    def download(self):
+        file_path = f"{self.raw_dir}/qm9_eV.npz"
+        if not os.path.exists(file_path):
+            download(self._url, path=file_path)
+
+    @property
+    def num_labels(self):
+        r"""
+        Returns
+        --------
+        int
+            Number of prediction tasks.
+        """
+        return self.label.shape[1]
+
+    @property
+    def num_classes(self):
+        r"""
+        Returns
+        --------
+        int
+            Number of prediction tasks.
+        """
+        return self.label.shape[1]
+
+    @property
+    def num_tasks(self):
+        r"""
+        Returns
+        --------
+        int
+            Number of prediction tasks.
+        """
+        return self.label.shape[1]
+
+    def __getitem__(self, idx):
+        r"""Get graph and label by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        dgl.DGLGraph
+            The graph contains:
+
+            - ``ndata['R']``: the coordinates of each atom
+            - ``ndata['Z']``: the atomic number
+
+        Tensor
+            Property values of molecular graphs
+        """
+        label = F.tensor(self.label[idx], dtype=F.data_type_dict["float32"])
+        n_atoms = self.N[idx]
+        R = self.R[self.N_cumsum[idx] : self.N_cumsum[idx + 1]]
+        dist = np.linalg.norm(R[:, None, :] - R[None, :, :], axis=-1)
+        adj = sp.csr_matrix(dist <= self.cutoff) - sp.eye(
+            n_atoms, dtype=np.bool_
+        )
+        adj = adj.tocoo()
+        u, v = F.tensor(adj.row), F.tensor(adj.col)
+        g = dgl_graph((u, v))
+        g = to_bidirected(g)
+        g.ndata["R"] = F.tensor(R, dtype=F.data_type_dict["float32"])
+        g.ndata["Z"] = F.tensor(
+            self.Z[self.N_cumsum[idx] : self.N_cumsum[idx + 1]],
+            dtype=F.data_type_dict["int64"],
+        )
+
+        if self._transform is not None:
+            g = self._transform(g)
+
+        return g, label
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Return
+        -------
+        int
+        """
+        return self.label.shape[0]
+
+
+QM9 = QM9Dataset

python/dgl/data/qm9_edge.py

@@ -0,0 +1,296 @@
+""" QM9 dataset for graph property prediction (regression) """
+
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLDataset
+from .utils import _get_dgl_url, download, extract_archive
+
+
+class QM9EdgeDataset(DGLDataset):
+    r"""QM9Edge dataset for graph property prediction (regression)
+
+    This dataset consists of 130,831 molecules with 19 regression targets.
+    Nodes correspond to atoms and edges correspond to bonds.
+
+    This dataset differs from :class:`~dgl.data.QM9Dataset` in the following aspects:
+        1. It includes the bonds in a molecule in the edges of the corresponding graph while the edges in :class:`~dgl.data.QM9Dataset` are purely distance-based.
+        2. It provides edge features, and node features in addition to the atoms' coordinates and atomic numbers.
+        3. It provides another 7 regression tasks(from 12 to 19).
+
+    This class is built based on a preprocessed version of the dataset, and we provide the preprocessing datails `here <https://gist.github.com/hengruizhang98/a2da30213b2356fff18b25385c9d3cd2>`_.
+
+    Reference:
+
+    - `"MoleculeNet: A Benchmark for Molecular Machine Learning" <https://arxiv.org/abs/1703.00564>`_
+    - `"Neural Message Passing for Quantum Chemistry" <https://arxiv.org/abs/1704.01212>`_
+
+    For
+    Statistics:
+
+    - Number of graphs: 130,831.
+    - Number of regression targets: 19.
+
+    Node attributes:
+
+    - pos: the 3D coordinates of each atom.
+    - attr: the 11D atom features.
+
+    Edge attributes:
+
+    - edge_attr: the 4D bond features.
+
+    Regression targets:
+
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | Keys   | Property                         | Description                                                                       | Unit                                        |
+    +========+==================================+===================================================================================+=============================================+
+    | mu     | :math:`\mu`                      | Dipole moment                                                                     | :math:`\textrm{D}`                          |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | alpha  | :math:`\alpha`                   | Isotropic polarizability                                                          | :math:`{a_0}^3`                             |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | homo   | :math:`\epsilon_{\textrm{HOMO}}` | Highest occupied molecular orbital energy                                         | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | lumo   | :math:`\epsilon_{\textrm{LUMO}}` | Lowest unoccupied molecular orbital energy                                        | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | gap    | :math:`\Delta \epsilon`          | Gap between :math:`\epsilon_{\textrm{HOMO}}` and :math:`\epsilon_{\textrm{LUMO}}` | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | r2     | :math:`\langle R^2 \rangle`      | Electronic spatial extent                                                         | :math:`{a_0}^2`                             |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | zpve   | :math:`\textrm{ZPVE}`            | Zero point vibrational energy                                                     | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U0     | :math:`U_0`                      | Internal energy at 0K                                                             | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U      | :math:`U`                        | Internal energy at 298.15K                                                        | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | H      | :math:`H`                        | Enthalpy at 298.15K                                                               | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | G      | :math:`G`                        | Free energy at 298.15K                                                            | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | Cv     | :math:`c_{\textrm{v}}`           | Heat capavity at 298.15K                                                          | :math:`\frac{\textrm{cal}}{\textrm{mol K}}` |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U0_atom| :math:`U_0^{\textrm{ATOM}}`      | Atomization energy at 0K                                                          | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | U_atom | :math:`U^{\textrm{ATOM}}`        | Atomization energy at 298.15K                                                     | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | H_atom | :math:`H^{\textrm{ATOM}}`        | Atomization enthalpy at 298.15K                                                   | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | G_atom | :math:`G^{\textrm{ATOM}}`        | Atomization free energy at 298.15K                                                | :math:`\textrm{eV}`                         |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | A      | :math:`A`                        | Rotational constant                                                               | :math:`\textrm{GHz}`                        |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | B      | :math:`B`                        | Rotational constant                                                               | :math:`\textrm{GHz}`                        |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+    | C      | :math:`C`                        | Rotational constant                                                               | :math:`\textrm{GHz}`                        |
+    +--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
+
+    Parameters
+    ----------
+    label_keys : list
+        Names of the regression property, which should be a subset of the keys in the table above.
+        If not provided, it will load all the labels.
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False.
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_tasks : int
+        Number of prediction tasks
+    num_labels : int
+        (DEPRECATED, use num_tasks instead) Number of prediction tasks
+
+    Raises
+    ------
+    UserWarning
+        If the raw data is changed in the remote server by the author.
+
+    Examples
+    --------
+    >>> data = QM9EdgeDataset(label_keys=['mu', 'alpha'])
+    >>> data.num_tasks
+    2
+
+    >>> # iterate over the dataset
+    >>> for graph, labels in data:
+    ...     print(graph) # get information of each graph
+    ...     print(labels) # get labels of the corresponding graph
+    ...     # your code here...
+    >>>
+    """
+
+    keys = [
+        "mu",
+        "alpha",
+        "homo",
+        "lumo",
+        "gap",
+        "r2",
+        "zpve",
+        "U0",
+        "U",
+        "H",
+        "G",
+        "Cv",
+        "U0_atom",
+        "U_atom",
+        "H_atom",
+        "G_atom",
+        "A",
+        "B",
+        "C",
+    ]
+    map_dict = {}
+
+    for i, key in enumerate(keys):
+        map_dict[key] = i
+
+    def __init__(
+        self,
+        label_keys=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        if label_keys is None:
+            self.label_keys = None
+            self.num_labels = 19
+        else:
+            self.label_keys = [self.map_dict[i] for i in label_keys]
+            self.num_labels = len(label_keys)
+
+        self._url = _get_dgl_url("dataset/qm9_edge.npz")
+
+        super(QM9EdgeDataset, self).__init__(
+            name="qm9Edge",
+            raw_dir=raw_dir,
+            url=self._url,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        if not os.path.exists(self.npz_path):
+            download(self._url, path=self.npz_path)
+
+    def process(self):
+        self.load()
+
+    @property
+    def npz_path(self):
+        return f"{self.raw_dir}/qm9_edge.npz"
+
+    def has_cache(self):
+        return os.path.exists(self.npz_path)
+
+    def save(self):
+        np.savez_compressed(
+            self.npz_path,
+            n_node=self.n_node,
+            n_edge=self.n_edge,
+            node_attr=self.node_attr,
+            node_pos=self.node_pos,
+            edge_attr=self.edge_attr,
+            src=self.src,
+            dst=self.dst,
+            targets=self.targets,
+        )
+
+    def load(self):
+        data_dict = np.load(self.npz_path, allow_pickle=True)
+
+        self.n_node = data_dict["n_node"]
+        self.n_edge = data_dict["n_edge"]
+        self.node_attr = data_dict["node_attr"]
+        self.node_pos = data_dict["node_pos"]
+        self.edge_attr = data_dict["edge_attr"]
+        self.targets = data_dict["targets"]
+
+        self.src = data_dict["src"]
+        self.dst = data_dict["dst"]
+
+        self.n_cumsum = np.concatenate([[0], np.cumsum(self.n_node)])
+        self.ne_cumsum = np.concatenate([[0], np.cumsum(self.n_edge)])
+
+    def __getitem__(self, idx):
+        r"""Get graph and label by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        dgl.DGLGraph
+           The graph contains:
+
+           - ``ndata['pos']``: the coordinates of each atom
+           - ``ndata['attr']``: the features of each atom
+           - ``edata['edge_attr']``: the features of each bond
+
+        Tensor
+            Property values of molecular graphs
+        """
+
+        pos = self.node_pos[self.n_cumsum[idx] : self.n_cumsum[idx + 1]]
+        src = self.src[self.ne_cumsum[idx] : self.ne_cumsum[idx + 1]]
+        dst = self.dst[self.ne_cumsum[idx] : self.ne_cumsum[idx + 1]]
+
+        g = dgl_graph((src, dst))
+
+        g.ndata["pos"] = F.tensor(pos, dtype=F.data_type_dict["float32"])
+        g.ndata["attr"] = F.tensor(
+            self.node_attr[self.n_cumsum[idx] : self.n_cumsum[idx + 1]],
+            dtype=F.data_type_dict["float32"],
+        )
+        g.edata["edge_attr"] = F.tensor(
+            self.edge_attr[self.ne_cumsum[idx] : self.ne_cumsum[idx + 1]],
+            dtype=F.data_type_dict["float32"],
+        )
+
+        label = F.tensor(
+            self.targets[idx][self.label_keys],
+            dtype=F.data_type_dict["float32"],
+        )
+
+        if self._transform is not None:
+            g = self._transform(g)
+
+        return g, label
+
+    def __len__(self):
+        r"""Number of graphs in the dataset.
+
+        Returns
+        -------
+        int
+        """
+        return self.n_node.shape[0]
+
+    @property
+    def num_tasks(self):
+        r"""
+        Returns
+        -------
+        int
+            Number of prediction tasks
+        """
+        return self.num_labels
+
+
+QM9Edge = QM9EdgeDataset

python/dgl/data/reddit.py

@@ -0,0 +1,223 @@
+""" Reddit dataset for community detection """
+from __future__ import absolute_import
+
+import os
+
+import numpy as np
+
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import from_scipy
+from ..transforms import reorder_graph
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    _get_dgl_url,
+    deprecate_property,
+    generate_mask_tensor,
+    load_graphs,
+    save_graphs,
+)
+
+
+class RedditDataset(DGLBuiltinDataset):
+    r"""Reddit dataset for community detection (node classification)
+
+    This is a graph dataset from Reddit posts made in the month of September, 2014.
+    The node label in this case is the community, or “subreddit”, that a post belongs to.
+    The authors sampled 50 large communities and built a post-to-post graph, connecting
+    posts if the same user comments on both. In total this dataset contains 232,965
+    posts with an average degree of 492. We use the first 20 days for training and the
+    remaining days for testing (with 30% used for validation).
+
+    Reference: `<http://snap.stanford.edu/graphsage/>`_
+
+    Statistics
+
+    - Nodes: 232,965
+    - Edges: 114,615,892
+    - Node feature size: 602
+    - Number of training samples: 153,431
+    - Number of validation samples: 23,831
+    - Number of test samples: 55,703
+
+    Parameters
+    ----------
+    self_loop : bool
+        Whether load dataset with self loop connections. Default: False
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of classes for each node
+
+    Examples
+    --------
+    >>> data = RedditDataset()
+    >>> g = data[0]
+    >>> num_classes = data.num_classes
+    >>>
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>>
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>>
+    >>> # get labels
+    >>> label = g.ndata['label']
+    >>>
+    >>> # Train, Validation and Test
+    """
+
+    def __init__(
+        self,
+        self_loop=False,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self_loop_str = ""
+        if self_loop:
+            self_loop_str = "_self_loop"
+        _url = _get_dgl_url("dataset/reddit{}.zip".format(self_loop_str))
+        self._self_loop_str = self_loop_str
+        super(RedditDataset, self).__init__(
+            name="reddit{}".format(self_loop_str),
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        # graph
+        coo_adj = sp.load_npz(
+            os.path.join(
+                self.raw_path, "reddit{}_graph.npz".format(self._self_loop_str)
+            )
+        )
+        self._graph = from_scipy(coo_adj)
+        # features and labels
+        reddit_data = np.load(os.path.join(self.raw_path, "reddit_data.npz"))
+        features = reddit_data["feature"]
+        labels = reddit_data["label"]
+        # tarin/val/test indices
+        node_types = reddit_data["node_types"]
+        train_mask = node_types == 1
+        val_mask = node_types == 2
+        test_mask = node_types == 3
+        self._graph.ndata["train_mask"] = generate_mask_tensor(train_mask)
+        self._graph.ndata["val_mask"] = generate_mask_tensor(val_mask)
+        self._graph.ndata["test_mask"] = generate_mask_tensor(test_mask)
+        self._graph.ndata["feat"] = F.tensor(
+            features, dtype=F.data_type_dict["float32"]
+        )
+        self._graph.ndata["label"] = F.tensor(
+            labels, dtype=F.data_type_dict["int64"]
+        )
+        self._graph = reorder_graph(
+            self._graph,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+        self._print_info()
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        if os.path.exists(graph_path):
+            return True
+        return False
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        graphs, _ = load_graphs(graph_path)
+        self._graph = graphs[0]
+        self._graph.ndata["train_mask"] = generate_mask_tensor(
+            self._graph.ndata["train_mask"].numpy()
+        )
+        self._graph.ndata["val_mask"] = generate_mask_tensor(
+            self._graph.ndata["val_mask"].numpy()
+        )
+        self._graph.ndata["test_mask"] = generate_mask_tensor(
+            self._graph.ndata["test_mask"].numpy()
+        )
+        self._print_info()
+
+    def _print_info(self):
+        if self.verbose:
+            print("Finished data loading.")
+            print("  NumNodes: {}".format(self._graph.num_nodes()))
+            print("  NumEdges: {}".format(self._graph.num_edges()))
+            print("  NumFeats: {}".format(self._graph.ndata["feat"].shape[1]))
+            print("  NumClasses: {}".format(self.num_classes))
+            print(
+                "  NumTrainingSamples: {}".format(
+                    F.nonzero_1d(self._graph.ndata["train_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumValidationSamples: {}".format(
+                    F.nonzero_1d(self._graph.ndata["val_mask"]).shape[0]
+                )
+            )
+            print(
+                "  NumTestSamples: {}".format(
+                    F.nonzero_1d(self._graph.ndata["test_mask"]).shape[0]
+                )
+            )
+
+    @property
+    def num_classes(self):
+        r"""Number of classes for each node."""
+        return 41
+
+    def __getitem__(self, idx):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+            graph structure, node labels, node features and splitting masks:
+
+            - ``ndata['label']``: node label
+            - ``ndata['feat']``: node feature
+            - ``ndata['train_mask']``： mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']:`` mask for test node set
+        """
+        assert idx == 0, "Reddit Dataset only has one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        r"""Number of graphs in the dataset"""
+        return 1

python/dgl/data/sbm.py

@@ -0,0 +1,276 @@
+"""Dataset for stochastic block model."""
+import math
+import os
+import random
+
+import numpy as np
+import numpy.random as npr
+import scipy as sp
+
+from .. import batch
+from ..convert import from_scipy
+from .dgl_dataset import DGLDataset
+from .utils import load_graphs, load_info, save_graphs, save_info
+
+
+def sbm(n_blocks, block_size, p, q, rng=None):
+    """(Symmetric) Stochastic Block Model
+
+    Parameters
+    ----------
+    n_blocks : int
+        Number of blocks.
+    block_size : int
+        Block size.
+    p : float
+        Probability for intra-community edge.
+    q : float
+        Probability for inter-community edge.
+    rng : numpy.random.RandomState, optional
+        Random number generator.
+
+    Returns
+    -------
+    scipy sparse matrix
+        The adjacency matrix of generated graph.
+    """
+    n = n_blocks * block_size
+    p /= n
+    q /= n
+    rng = np.random.RandomState() if rng is None else rng
+
+    rows = []
+    cols = []
+    for i in range(n_blocks):
+        for j in range(i, n_blocks):
+            density = p if i == j else q
+            block = sp.sparse.random(
+                block_size,
+                block_size,
+                density,
+                random_state=rng,
+                data_rvs=lambda n: np.ones(n),
+            )
+            rows.append(block.row + i * block_size)
+            cols.append(block.col + j * block_size)
+
+    rows = np.hstack(rows)
+    cols = np.hstack(cols)
+    a = sp.sparse.coo_matrix(
+        (np.ones(rows.shape[0]), (rows, cols)), shape=(n, n)
+    )
+    adj = sp.sparse.triu(a) + sp.sparse.triu(a, 1).transpose()
+    return adj
+
+
+class SBMMixtureDataset(DGLDataset):
+    r"""Symmetric Stochastic Block Model Mixture
+
+    Reference: Appendix C of `Supervised Community Detection with Hierarchical Graph Neural Networks <https://arxiv.org/abs/1705.08415>`_
+
+    Parameters
+    ----------
+    n_graphs : int
+        Number of graphs.
+    n_nodes : int
+        Number of nodes.
+    n_communities : int
+        Number of communities.
+    k : int, optional
+        Multiplier. Default: 2
+    avg_deg : int, optional
+        Average degree. Default: 3
+    pq : list of pair of nonnegative float or str, optional
+        Random densities. This parameter is for future extension,
+        for now it's always using the default value.
+        Default: Appendix_C
+    rng : numpy.random.RandomState, optional
+        Random number generator. If not given, it's numpy.random.RandomState() with `seed=None`,
+        which read data from /dev/urandom (or the Windows analogue) if available or seed from
+        the clock otherwise.
+        Default: None
+
+    Raises
+    ------
+    RuntimeError is raised if pq is not a list or string.
+
+    Examples
+    --------
+    >>> data = SBMMixtureDataset(n_graphs=16, n_nodes=10000, n_communities=2)
+    >>> from torch.utils.data import DataLoader
+    >>> dataloader = DataLoader(data, batch_size=1, collate_fn=data.collate_fn)
+    >>> for graph, line_graph, graph_degrees, line_graph_degrees, pm_pd in dataloader:
+    ...     # your code here
+    """
+
+    def __init__(
+        self,
+        n_graphs,
+        n_nodes,
+        n_communities,
+        k=2,
+        avg_deg=3,
+        pq="Appendix_C",
+        rng=None,
+    ):
+        self._n_graphs = n_graphs
+        self._n_nodes = n_nodes
+        self._n_communities = n_communities
+        assert n_nodes % n_communities == 0
+        self._block_size = n_nodes // n_communities
+        self._k = k
+        self._avg_deg = avg_deg
+        self._pq = pq
+        self._rng = rng
+        super(SBMMixtureDataset, self).__init__(
+            name="sbmmixture",
+            hash_key=(n_graphs, n_nodes, n_communities, k, avg_deg, pq, rng),
+        )
+
+    def process(self):
+        pq = self._pq
+        if type(pq) is list:
+            assert len(pq) == self._n_graphs
+        elif type(pq) is str:
+            generator = {"Appendix_C": self._appendix_c}[pq]
+            pq = [generator() for _ in range(self._n_graphs)]
+        else:
+            raise RuntimeError()
+        self._graphs = [
+            from_scipy(sbm(self._n_communities, self._block_size, *x))
+            for x in pq
+        ]
+        self._line_graphs = [
+            g.line_graph(backtracking=False) for g in self._graphs
+        ]
+        in_degrees = lambda g: g.in_degrees().float()
+        self._graph_degrees = [in_degrees(g) for g in self._graphs]
+        self._line_graph_degrees = [in_degrees(lg) for lg in self._line_graphs]
+        self._pm_pds = list(zip(*[g.edges() for g in self._graphs]))[0]
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, "graphs_{}.bin".format(self.hash))
+
+    @property
+    def line_graph_path(self):
+        return os.path.join(
+            self.save_path, "line_graphs_{}.bin".format(self.hash)
+        )
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, "info_{}.pkl".format(self.hash))
+
+    def has_cache(self):
+        return (
+            os.path.exists(self.graph_path)
+            and os.path.exists(self.line_graph_path)
+            and os.path.exists(self.info_path)
+        )
+
+    def save(self):
+        save_graphs(self.graph_path, self._graphs)
+        save_graphs(self.line_graph_path, self._line_graphs)
+        save_info(
+            self.info_path,
+            {
+                "graph_degree": self._graph_degrees,
+                "line_graph_degree": self._line_graph_degrees,
+                "pm_pds": self._pm_pds,
+            },
+        )
+
+    def load(self):
+        self._graphs, _ = load_graphs(self.graph_path)
+        self._line_graphs, _ = load_graphs(self.line_graph_path)
+        info = load_info(self.info_path)
+        self._graph_degrees = info["graph_degree"]
+        self._line_graph_degrees = info["line_graph_degree"]
+        self._pm_pds = info["pm_pds"]
+
+    def __len__(self):
+        r"""Number of graphs in the dataset."""
+        return len(self._graphs)
+
+    def __getitem__(self, idx):
+        r"""Get one example by index
+
+        Parameters
+        ----------
+        idx : int
+            Item index
+
+        Returns
+        -------
+        graph: :class:`dgl.DGLGraph`
+            The original graph
+        line_graph: :class:`dgl.DGLGraph`
+            The line graph of `graph`
+        graph_degree: numpy.ndarray
+            In degrees for each node in `graph`
+        line_graph_degree: numpy.ndarray
+            In degrees for each node in `line_graph`
+        pm_pd: numpy.ndarray
+            Edge indicator matrices Pm and Pd
+        """
+        return (
+            self._graphs[idx],
+            self._line_graphs[idx],
+            self._graph_degrees[idx],
+            self._line_graph_degrees[idx],
+            self._pm_pds[idx],
+        )
+
+    def _appendix_c(self):
+        q = npr.uniform(0, self._avg_deg - math.sqrt(self._avg_deg))
+        p = self._k * self._avg_deg - q
+        if random.random() < 0.5:
+            return p, q
+        else:
+            return q, p
+
+    def collate_fn(self, x):
+        r"""The `collate` function for dataloader
+
+        Parameters
+        ----------
+        x : tuple
+            a batch of data that contains:
+
+            - graph: :class:`dgl.DGLGraph`
+                The original graph
+            - line_graph: :class:`dgl.DGLGraph`
+                The line graph of `graph`
+            - graph_degree: numpy.ndarray
+                In degrees for each node in `graph`
+            - line_graph_degree: numpy.ndarray
+                In degrees for each node in `line_graph`
+            - pm_pd: numpy.ndarray
+                Edge indicator matrices Pm and Pd
+
+        Returns
+        -------
+        g_batch: :class:`dgl.DGLGraph`
+            Batched graphs
+        lg_batch: :class:`dgl.DGLGraph`
+            Batched line graphs
+        degg_batch: numpy.ndarray
+            A batch of in degrees for each node in `g_batch`
+        deglg_batch: numpy.ndarray
+            A batch of in degrees for each node in `lg_batch`
+        pm_pd_batch: numpy.ndarray
+            A batch of edge indicator matrices Pm and Pd
+        """
+        g, lg, deg_g, deg_lg, pm_pd = zip(*x)
+        g_batch = batch.batch(g)
+        lg_batch = batch.batch(lg)
+        degg_batch = np.concatenate(deg_g, axis=0)
+        deglg_batch = np.concatenate(deg_lg, axis=0)
+        pm_pd_batch = np.concatenate(
+            [x + i * self._n_nodes for i, x in enumerate(pm_pd)], axis=0
+        )
+        return g_batch, lg_batch, degg_batch, deglg_batch, pm_pd_batch
+
+
+SBMMixture = SBMMixtureDataset

python/dgl/data/superpixel.py

@@ -0,0 +1,435 @@
+import os
+import pickle
+
+import numpy as np
+from scipy.spatial.distance import cdist
+from tqdm.auto import tqdm
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLDataset
+from .utils import download, extract_archive, load_graphs, save_graphs, Subset
+
+
+def sigma(dists, kth=8):
+    num_nodes = dists.shape[0]
+
+    # Compute sigma and reshape.
+    if kth > num_nodes:
+        # Handling for graphs with num_nodes less than kth.
+        sigma = np.array([1] * num_nodes).reshape(num_nodes, 1)
+    else:
+        # Get k-nearest neighbors for each node.
+        knns = np.partition(dists, kth, axis=-1)[:, : kth + 1]
+        sigma = knns.sum(axis=1).reshape((knns.shape[0], 1)) / kth
+
+    return sigma + 1e-8
+
+
+def compute_adjacency_matrix_images(coord, feat, use_feat=True):
+    coord = coord.reshape(-1, 2)
+    # Compute coordinate distance.
+    c_dist = cdist(coord, coord)
+
+    if use_feat:
+        # Compute feature distance.
+        f_dist = cdist(feat, feat)
+        # Compute adjacency.
+        A = np.exp(
+            -((c_dist / sigma(c_dist)) ** 2) - (f_dist / sigma(f_dist)) ** 2
+        )
+    else:
+        A = np.exp(-((c_dist / sigma(c_dist)) ** 2))
+
+    # Convert to symmetric matrix.
+    A = 0.5 * (A + A.T)
+    A[np.diag_indices_from(A)] = 0
+    return A
+
+
+def compute_edges_list(A, kth=9):
+    # Get k-similar neighbor indices for each node.
+    num_nodes = A.shape[0]
+    new_kth = num_nodes - kth
+
+    if num_nodes > kth:
+        knns = np.argpartition(A, new_kth - 1, axis=-1)[:, new_kth:-1]
+        knn_values = np.partition(A, new_kth - 1, axis=-1)[:, new_kth:-1]
+    else:
+        # Handling for graphs with less than kth nodes.
+        # In such cases, the resulting graph will be fully connected.
+        knns = np.tile(np.arange(num_nodes), num_nodes).reshape(
+            num_nodes, num_nodes
+        )
+        knn_values = A
+
+        # Removing self loop.
+        if num_nodes != 1:
+            knn_values = A[knns != np.arange(num_nodes)[:, None]].reshape(
+                num_nodes, -1
+            )
+            knns = knns[knns != np.arange(num_nodes)[:, None]].reshape(
+                num_nodes, -1
+            )
+    return knns, knn_values
+
+
+class SuperPixelDataset(DGLDataset):
+    def __init__(
+        self,
+        raw_dir=None,
+        name="MNIST",
+        split="train",
+        use_feature=False,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        assert split in ["train", "test"], "split not valid."
+        assert name in ["MNIST", "CIFAR10"], "name not valid."
+
+        self.use_feature = use_feature
+        self.split = split
+        self._dataset_name = name
+        self.graphs = []
+        self.labels = []
+
+        super().__init__(
+            name="Superpixel",
+            raw_dir=raw_dir,
+            url="""
+            https://www.dropbox.com/s/y2qwa77a0fxem47/superpixels.zip?dl=1
+            """,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    @property
+    def img_size(self):
+        r"""Size of dataset image."""
+        if self._dataset_name == "MNIST":
+            return 28
+        return 32
+
+    @property
+    def save_path(self):
+        r"""Directory to save the processed dataset."""
+        return os.path.join(self.raw_path, "processed")
+
+    @property
+    def raw_data_path(self):
+        r"""Path to save the raw dataset file."""
+        return os.path.join(self.raw_path, "superpixels.zip")
+
+    @property
+    def graph_path(self):
+        r"""Path to save the processed dataset file."""
+        if self.use_feature:
+            return os.path.join(
+                self.save_path,
+                f"use_feat_{self._dataset_name}_{self.split}.pkl",
+            )
+        return os.path.join(
+            self.save_path, f"{self._dataset_name}_{self.split}.pkl"
+        )
+
+    def download(self):
+        path = download(self.url, path=self.raw_data_path)
+        extract_archive(path, target_dir=self.raw_path, overwrite=True)
+
+    def process(self):
+        if self._dataset_name == "MNIST":
+            plk_file = "mnist_75sp"
+        elif self._dataset_name == "CIFAR10":
+            plk_file = "cifar10_150sp"
+
+        with open(
+            os.path.join(
+                self.raw_path, "superpixels", f"{plk_file}_{self.split}.pkl"
+            ),
+            "rb",
+        ) as f:
+            self.labels, self.sp_data = pickle.load(f)
+            self.labels = F.tensor(self.labels)
+
+        self.Adj_matrices = []
+        self.node_features = []
+        self.edges_lists = []
+        self.edge_features = []
+
+        for index, sample in enumerate(
+            tqdm(self.sp_data, desc=f"Processing {self.split} dataset")
+        ):
+            mean_px, coord = sample[:2]
+            coord = coord / self.img_size
+
+            if self.use_feature:
+                A = compute_adjacency_matrix_images(
+                    coord, mean_px
+                )  # using super-pixel locations + features
+            else:
+                A = compute_adjacency_matrix_images(
+                    coord, mean_px, False
+                )  # using only super-pixel locations
+            edges_list, edge_values_list = compute_edges_list(A)
+
+            N_nodes = A.shape[0]
+
+            mean_px = mean_px.reshape(N_nodes, -1)
+            coord = coord.reshape(N_nodes, 2)
+            x = np.concatenate((mean_px, coord), axis=1)
+
+            edge_values_list = edge_values_list.reshape(-1)
+
+            self.node_features.append(x)
+            self.edge_features.append(edge_values_list)
+            self.Adj_matrices.append(A)
+            self.edges_lists.append(edges_list)
+
+        for index in tqdm(
+            range(len(self.sp_data)), desc=f"Dump {self.split} dataset"
+        ):
+            N = self.node_features[index].shape[0]
+
+            src_nodes = []
+            dst_nodes = []
+            for src, dsts in enumerate(self.edges_lists[index]):
+                # handling for 1 node where the self loop would be the only edge
+                if N == 1:
+                    src_nodes.append(src)
+                    dst_nodes.append(dsts)
+                else:
+                    dsts = dsts[dsts != src]
+                    srcs = [src] * len(dsts)
+                    src_nodes.extend(srcs)
+                    dst_nodes.extend(dsts)
+
+            src_nodes = F.tensor(src_nodes)
+            dst_nodes = F.tensor(dst_nodes)
+
+            g = dgl_graph((src_nodes, dst_nodes), num_nodes=N)
+            g.ndata["feat"] = F.zerocopy_from_numpy(
+                self.node_features[index]
+            ).to(F.float32)
+            g.edata["feat"] = (
+                F.zerocopy_from_numpy(self.edge_features[index])
+                .to(F.float32)
+                .unsqueeze(1)
+            )
+
+            self.graphs.append(g)
+
+    def load(self):
+        self.graphs, label_dict = load_graphs(self.graph_path)
+        self.labels = label_dict["labels"]
+
+    def save(self):
+        save_graphs(
+            self.graph_path, self.graphs, labels={"labels": self.labels}
+        )
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def __len__(self):
+        return len(self.graphs)
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int or tensor
+            The sample index.
+            1-D tensor as `idx` is allowed when transform is None.
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+            Graph with node feature stored in ``feat`` field and its label.
+        or
+        :class:`dgl.data.utils.Subset`
+            Subset of the dataset at specified indices
+        """
+        if F.is_tensor(idx) and idx.dim() == 1:
+            if self._transform is None:
+                return Subset(self, idx.cpu())
+
+            raise ValueError(
+                "Tensor idx not supported when transform is not None."
+            )
+
+        if self._transform is None:
+            return self.graphs[idx], self.labels[idx]
+
+        return self._transform(self.graphs[idx]), self.labels[idx]
+
+
+class MNISTSuperPixelDataset(SuperPixelDataset):
+    r"""MNIST superpixel dataset for the graph classification task.
+
+    DGL dataset of MNIST and CIFAR10 in the benchmark-gnn which contains graphs
+    converted fromt the original MINST and CIFAR10 images.
+
+    Reference `<http://arxiv.org/abs/2003.00982>`_
+
+    Statistics:
+
+        - Train examples: 60,000
+        - Test examples: 10,000
+        - Size of dataset images: 28
+
+    Parameters
+    ----------
+    raw_dir : str
+        Directory to store all the downloaded raw datasets.
+        Default: "~/.dgl/".
+    split : str
+        Should be chosen from ["train", "test"]
+        Default: "train".
+    use_feature: bool
+
+        - True: Adj matrix defined from super-pixel locations + features
+        - False: Adj matrix defined from super-pixel locations (only)
+
+        Default: False.
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False.
+    verbose : bool
+        Whether to print out progress information.
+        Default: False.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Examples
+    ---------
+    >>> from dgl.data import MNISTSuperPixelDataset
+
+    >>> # MNIST dataset
+    >>> train_dataset = MNISTSuperPixelDataset(split="train")
+    >>> len(train_dataset)
+    60000
+    >>> graph, label = train_dataset[0]
+    >>> graph
+    Graph(num_nodes=71, num_edges=568,
+        ndata_schemes={'feat': Scheme(shape=(3,), dtype=torch.float32)}
+        edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)})
+
+    >>> # support tensor to be index when transform is None
+    >>> # see details in __getitem__ function
+    >>> import torch
+    >>> idx = torch.tensor([0, 1, 2])
+    >>> train_dataset_subset = train_dataset[idx]
+    >>> train_dataset_subset[0]
+    Graph(num_nodes=71, num_edges=568,
+        ndata_schemes={'feat': Scheme(shape=(3,), dtype=torch.float32)}
+        edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)})
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        split="train",
+        use_feature=False,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        super().__init__(
+            raw_dir=raw_dir,
+            name="MNIST",
+            split=split,
+            use_feature=use_feature,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+
+class CIFAR10SuperPixelDataset(SuperPixelDataset):
+    r"""CIFAR10 superpixel dataset for the graph classification task.
+
+    DGL dataset of CIFAR10 in the benchmark-gnn which contains graphs
+    converted fromt the original CIFAR10 images.
+
+    Reference `<http://arxiv.org/abs/2003.00982>`_
+
+    Statistics:
+
+        - Train examples: 50,000
+        - Test examples: 10,000
+        - Size of dataset images: 32
+
+    Parameters
+    ----------
+    raw_dir : str
+        Directory to store all the downloaded raw datasets.
+        Default: "~/.dgl/".
+    split : str
+        Should be chosen from ["train", "test"]
+        Default: "train".
+    use_feature: bool
+
+        - True: Adj matrix defined from super-pixel locations + features
+        - False: Adj matrix defined from super-pixel locations (only)
+
+        Default: False.
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False.
+    verbose : bool
+        Whether to print out progress information.
+        Default: False.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Examples
+    ---------
+    >>> from dgl.data import CIFAR10SuperPixelDataset
+
+    >>> # CIFAR10 dataset
+    >>> train_dataset = CIFAR10SuperPixelDataset(split="train")
+    >>> len(train_dataset)
+    50000
+    >>> graph, label = train_dataset[0]
+    >>> graph
+    Graph(num_nodes=123, num_edges=984,
+        ndata_schemes={'feat': Scheme(shape=(5,), dtype=torch.float32)}
+        edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)}),
+
+    >>> # support tensor to be index when transform is None
+    >>> # see details in __getitem__ function
+    >>> import torch
+    >>> idx = torch.tensor([0, 1, 2])
+    >>> train_dataset_subset = train_dataset[idx]
+    >>> train_dataset_subset[0]
+    Graph(num_nodes=123, num_edges=984,
+        ndata_schemes={'feat': Scheme(shape=(5,), dtype=torch.float32)}
+        edata_schemes={'feat': Scheme(shape=(1,), dtype=torch.float32)}),
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        split="train",
+        use_feature=False,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        super().__init__(
+            raw_dir=raw_dir,
+            name="CIFAR10",
+            split=split,
+            use_feature=use_feature,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )

python/dgl/data/synthetic.py

@@ -0,0 +1,834 @@
+"""Synthetic graph datasets."""
+import math
+import os
+import pickle
+import random
+
+import networkx as nx
+import numpy as np
+
+from .. import backend as F
+from ..batch import batch
+from ..convert import graph
+from ..transforms import reorder_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, download, load_graphs, save_graphs
+
+
+class BAShapeDataset(DGLBuiltinDataset):
+    r"""BA-SHAPES dataset from `GNNExplainer: Generating Explanations for Graph Neural Networks
+    <https://arxiv.org/abs/1903.03894>`__
+
+    This is a synthetic dataset for node classification. It is generated by performing the
+    following steps in order.
+
+    - Construct a base Barabási–Albert (BA) graph.
+    - Construct a set of five-node house-structured network motifs.
+    - Attach the motifs to randomly selected nodes of the base graph.
+    - Perturb the graph by adding random edges.
+    - Nodes are assigned to 4 classes. Nodes of label 0 belong to the base BA graph. Nodes of
+      label 1, 2, 3 are separately at the middle, bottom, or top of houses.
+    - Generate constant feature for all nodes, which is 1.
+
+    Parameters
+    ----------
+    num_base_nodes : int, optional
+        Number of nodes in the base BA graph. Default: 300
+    num_base_edges_per_node : int, optional
+        Number of edges to attach from a new node to existing nodes in constructing the base BA
+        graph. Default: 5
+    num_motifs : int, optional
+        Number of house-structured network motifs to use. Default: 80
+    perturb_ratio : float, optional
+        Number of random edges to add in perturbation divided by the number of edges in the
+        original graph. Default: 0.01
+    seed : integer, random_state, or None, optional
+        Indicator of random number generation state. Default: None
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to always generate the data from scratch rather than load a cached version.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import BAShapeDataset
+    >>> dataset = BAShapeDataset()
+    >>> dataset.num_classes
+    4
+    >>> g = dataset[0]
+    >>> label = g.ndata['label']
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self,
+        num_base_nodes=300,
+        num_base_edges_per_node=5,
+        num_motifs=80,
+        perturb_ratio=0.01,
+        seed=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self.num_base_nodes = num_base_nodes
+        self.num_base_edges_per_node = num_base_edges_per_node
+        self.num_motifs = num_motifs
+        self.perturb_ratio = perturb_ratio
+        self.seed = seed
+        super(BAShapeDataset, self).__init__(
+            name="BA-SHAPES",
+            url=None,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        g = nx.barabasi_albert_graph(
+            self.num_base_nodes, self.num_base_edges_per_node, self.seed
+        )
+        edges = list(g.edges())
+        src, dst = map(list, zip(*edges))
+        n = self.num_base_nodes
+
+        # Nodes in the base BA graph belong to class 0
+        node_labels = [0] * n
+        # The motifs will be evenly attached to the nodes in the base graph.
+        spacing = math.floor(n / self.num_motifs)
+
+        for motif_id in range(self.num_motifs):
+            # Construct a five-node house-structured network motif
+            motif_edges = [
+                (n, n + 1),
+                (n + 1, n + 2),
+                (n + 2, n + 3),
+                (n + 3, n),
+                (n + 4, n),
+                (n + 4, n + 1),
+            ]
+            motif_src, motif_dst = map(list, zip(*motif_edges))
+            src.extend(motif_src)
+            dst.extend(motif_dst)
+
+            # Nodes at the middle of a house belong to class 1
+            # Nodes at the bottom of a house belong to class 2
+            # Nodes at the top of a house belong to class 3
+            node_labels.extend([1, 1, 2, 2, 3])
+
+            # Attach the motif to the base BA graph
+            src.append(n)
+            dst.append(int(motif_id * spacing))
+            n += 5
+
+        g = graph((src, dst), num_nodes=n)
+
+        # Perturb the graph by adding non-self-loop random edges
+        num_real_edges = g.num_edges()
+        max_ratio = (n * (n - 1) - num_real_edges) / num_real_edges
+        assert (
+            self.perturb_ratio <= max_ratio
+        ), "perturb_ratio cannot exceed {:.4f}".format(max_ratio)
+        num_random_edges = int(num_real_edges * self.perturb_ratio)
+
+        if self.seed is not None:
+            np.random.seed(self.seed)
+        for _ in range(num_random_edges):
+            while True:
+                u = np.random.randint(0, n)
+                v = np.random.randint(0, n)
+                if (not g.has_edges_between(u, v)) and (u != v):
+                    break
+            g.add_edges(u, v)
+
+        g.ndata["label"] = F.tensor(node_labels, F.int64)
+        g.ndata["feat"] = F.ones((n, 1), F.float32, F.cpu())
+        self._graph = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        save_graphs(str(self.graph_path), self._graph)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+        self._graph = graphs[0]
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return 4
+
+
+class BACommunityDataset(DGLBuiltinDataset):
+    r"""BA-COMMUNITY dataset from `GNNExplainer: Generating Explanations for Graph Neural Networks
+    <https://arxiv.org/abs/1903.03894>`__
+
+    This is a synthetic dataset for node classification. It is generated by performing the
+    following steps in order.
+
+    - Construct a base Barabási–Albert (BA) graph.
+    - Construct a set of five-node house-structured network motifs.
+    - Attach the motifs to randomly selected nodes of the base graph.
+    - Perturb the graph by adding random edges.
+    - Nodes are assigned to 4 classes. Nodes of label 0 belong to the base BA graph. Nodes of
+      label 1, 2, 3 are separately at the middle, bottom, or top of houses.
+    - Generate normally distributed features of length 10
+    - Repeat the above steps to generate another graph. Its nodes are assigned to class
+      4, 5, 6, 7. Its node features are generated with a distinct normal distribution.
+    - Join the two graphs by randomly adding edges between them.
+
+    Parameters
+    ----------
+    num_base_nodes : int, optional
+        Number of nodes in each base BA graph. Default: 300
+    num_base_edges_per_node : int, optional
+        Number of edges to attach from a new node to existing nodes in constructing a base BA
+        graph. Default: 4
+    num_motifs : int, optional
+        Number of house-structured network motifs to use in constructing each graph. Default: 80
+    perturb_ratio : float, optional
+        Number of random edges to add to a graph in perturbation divided by the number of original
+        edges in it. Default: 0.01
+    num_inter_edges : int, optional
+        Number of random edges to add between the two graphs. Default: 350
+    seed : integer, random_state, or None, optional
+        Indicator of random number generation state. Default: None
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to always generate the data from scratch rather than load a cached version.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import BACommunityDataset
+    >>> dataset = BACommunityDataset()
+    >>> dataset.num_classes
+    8
+    >>> g = dataset[0]
+    >>> label = g.ndata['label']
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self,
+        num_base_nodes=300,
+        num_base_edges_per_node=4,
+        num_motifs=80,
+        perturb_ratio=0.01,
+        num_inter_edges=350,
+        seed=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self.num_base_nodes = num_base_nodes
+        self.num_base_edges_per_node = num_base_edges_per_node
+        self.num_motifs = num_motifs
+        self.perturb_ratio = perturb_ratio
+        self.num_inter_edges = num_inter_edges
+        self.seed = seed
+        super(BACommunityDataset, self).__init__(
+            name="BA-COMMUNITY",
+            url=None,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        if self.seed is not None:
+            random.seed(self.seed)
+            np.random.seed(self.seed)
+
+        # Construct two BA-SHAPES graphs
+        g1 = BAShapeDataset(
+            self.num_base_nodes,
+            self.num_base_edges_per_node,
+            self.num_motifs,
+            self.perturb_ratio,
+            force_reload=True,
+            verbose=False,
+        )[0]
+        g2 = BAShapeDataset(
+            self.num_base_nodes,
+            self.num_base_edges_per_node,
+            self.num_motifs,
+            self.perturb_ratio,
+            force_reload=True,
+            verbose=False,
+        )[0]
+
+        # Join them and randomly add edges between them
+        g = batch([g1, g2])
+        num_nodes = g.num_nodes() // 2
+        src = np.random.randint(0, num_nodes, (self.num_inter_edges,))
+        dst = np.random.randint(
+            num_nodes, 2 * num_nodes, (self.num_inter_edges,)
+        )
+        src = F.astype(F.zerocopy_from_numpy(src), g.idtype)
+        dst = F.astype(F.zerocopy_from_numpy(dst), g.idtype)
+        g.add_edges(src, dst)
+        g.ndata["label"] = F.cat(
+            [g1.ndata["label"], g2.ndata["label"] + 4], dim=0
+        )
+
+        # feature generation
+        random_mu = [0.0] * 8
+        random_sigma = [1.0] * 8
+
+        mu_1, sigma_1 = np.array([-1.0] * 2 + random_mu), np.array(
+            [0.5] * 2 + random_sigma
+        )
+        feat1 = np.random.multivariate_normal(mu_1, np.diag(sigma_1), num_nodes)
+
+        mu_2, sigma_2 = np.array([1.0] * 2 + random_mu), np.array(
+            [0.5] * 2 + random_sigma
+        )
+        feat2 = np.random.multivariate_normal(mu_2, np.diag(sigma_2), num_nodes)
+
+        feat = np.concatenate([feat1, feat2])
+        g.ndata["feat"] = F.zerocopy_from_numpy(feat)
+        self._graph = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        save_graphs(str(self.graph_path), self._graph)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+        self._graph = graphs[0]
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return 8
+
+
+class TreeCycleDataset(DGLBuiltinDataset):
+    r"""TREE-CYCLES dataset from `GNNExplainer: Generating Explanations for Graph Neural Networks
+    <https://arxiv.org/abs/1903.03894>`__
+
+    This is a synthetic dataset for node classification. It is generated by performing the
+    following steps in order.
+
+    - Construct a balanced binary tree as the base graph.
+    - Construct a set of cycle motifs.
+    - Attach the motifs to randomly selected nodes of the base graph.
+    - Perturb the graph by adding random edges.
+    - Generate constant feature for all nodes, which is 1.
+    - Nodes in the tree belong to class 0 and nodes in cycles belong to class 1.
+
+    Parameters
+    ----------
+    tree_height : int, optional
+        Height of the balanced binary tree. Default: 8
+    num_motifs : int, optional
+        Number of cycle motifs to use. Default: 60
+    cycle_size : int, optional
+        Number of nodes in a cycle motif. Default: 6
+    perturb_ratio : float, optional
+        Number of random edges to add in perturbation divided by the
+        number of original edges in the graph. Default: 0.01
+    seed : integer, random_state, or None, optional
+        Indicator of random number generation state. Default: None
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to always generate the data from scratch rather than load a cached version.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import TreeCycleDataset
+    >>> dataset = TreeCycleDataset()
+    >>> dataset.num_classes
+    2
+    >>> g = dataset[0]
+    >>> label = g.ndata['label']
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self,
+        tree_height=8,
+        num_motifs=60,
+        cycle_size=6,
+        perturb_ratio=0.01,
+        seed=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self.tree_height = tree_height
+        self.num_motifs = num_motifs
+        self.cycle_size = cycle_size
+        self.perturb_ratio = perturb_ratio
+        self.seed = seed
+        super(TreeCycleDataset, self).__init__(
+            name="TREE-CYCLES",
+            url=None,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        if self.seed is not None:
+            np.random.seed(self.seed)
+
+        g = nx.balanced_tree(r=2, h=self.tree_height)
+        edges = list(g.edges())
+        src, dst = map(list, zip(*edges))
+        n = nx.number_of_nodes(g)
+
+        # Nodes in the base tree graph belong to class 0
+        node_labels = [0] * n
+        # The motifs will be evenly attached to the nodes in the base graph.
+        spacing = math.floor(n / self.num_motifs)
+
+        for motif_id in range(self.num_motifs):
+            # Construct a six-node cycle
+            motif_edges = [(n + i, n + i + 1) for i in range(5)]
+            motif_edges.append((n + 5, n))
+            motif_src, motif_dst = map(list, zip(*motif_edges))
+            src.extend(motif_src)
+            dst.extend(motif_dst)
+
+            # Nodes in cycles belong to class 1
+            node_labels.extend([1] * self.cycle_size)
+
+            # Attach the motif to the base tree graph
+            anchor = int(motif_id * spacing)
+            src.append(n)
+            dst.append(anchor)
+
+            if np.random.random() > 0.5:
+                a = np.random.randint(1, 4)
+                b = np.random.randint(1, 4)
+                src.append(n + a)
+                dst.append(anchor + b)
+
+            n += self.cycle_size
+
+        g = graph((src, dst), num_nodes=n)
+
+        # Perturb the graph by adding non-self-loop random edges
+        num_real_edges = g.num_edges()
+        max_ratio = (n * (n - 1) - num_real_edges) / num_real_edges
+        assert (
+            self.perturb_ratio <= max_ratio
+        ), "perturb_ratio cannot exceed {:.4f}".format(max_ratio)
+        num_random_edges = int(num_real_edges * self.perturb_ratio)
+
+        for _ in range(num_random_edges):
+            while True:
+                u = np.random.randint(0, n)
+                v = np.random.randint(0, n)
+                if (not g.has_edges_between(u, v)) and (u != v):
+                    break
+            g.add_edges(u, v)
+
+        g.ndata["label"] = F.tensor(node_labels, F.int64)
+        g.ndata["feat"] = F.ones((n, 1), F.float32, F.cpu())
+        self._graph = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        save_graphs(str(self.graph_path), self._graph)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+        self._graph = graphs[0]
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return 2
+
+
+class TreeGridDataset(DGLBuiltinDataset):
+    r"""TREE-GRIDS dataset from `GNNExplainer: Generating Explanations for Graph Neural Networks
+    <https://arxiv.org/abs/1903.03894>`__
+
+    This is a synthetic dataset for node classification. It is generated by performing the
+    following steps in order.
+
+    - Construct a balanced binary tree as the base graph.
+    - Construct a set of n-by-n grid motifs.
+    - Attach the motifs to randomly selected nodes of the base graph.
+    - Perturb the graph by adding random edges.
+    - Generate constant feature for all nodes, which is 1.
+    - Nodes in the tree belong to class 0 and nodes in grids belong to class 1.
+
+    Parameters
+    ----------
+    tree_height : int, optional
+        Height of the balanced binary tree. Default: 8
+    num_motifs : int, optional
+        Number of grid motifs to use. Default: 80
+    grid_size : int, optional
+        The number of nodes in a grid motif will be grid_size ^ 2. Default: 3
+    perturb_ratio : float, optional
+        Number of random edges to add in perturbation divided by the
+        number of original edges in the graph. Default: 0.1
+    seed : integer, random_state, or None, optional
+        Indicator of random number generation state. Default: None
+    raw_dir : str, optional
+        Raw file directory to store the processed data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to always generate the data from scratch rather than load a cached version.
+        Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import TreeGridDataset
+    >>> dataset = TreeGridDataset()
+    >>> dataset.num_classes
+    2
+    >>> g = dataset[0]
+    >>> label = g.ndata['label']
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self,
+        tree_height=8,
+        num_motifs=80,
+        grid_size=3,
+        perturb_ratio=0.1,
+        seed=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=True,
+        transform=None,
+    ):
+        self.tree_height = tree_height
+        self.num_motifs = num_motifs
+        self.grid_size = grid_size
+        self.perturb_ratio = perturb_ratio
+        self.seed = seed
+        super(TreeGridDataset, self).__init__(
+            name="TREE-GRIDS",
+            url=None,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        if self.seed is not None:
+            np.random.seed(self.seed)
+
+        g = nx.balanced_tree(r=2, h=self.tree_height)
+        edges = list(g.edges())
+        src, dst = map(list, zip(*edges))
+        n = nx.number_of_nodes(g)
+
+        # Nodes in the base tree graph belong to class 0
+        node_labels = [0] * n
+        # The motifs will be evenly attached to the nodes in the base graph.
+        spacing = math.floor(n / self.num_motifs)
+
+        # Construct an n-by-n grid
+        motif_g = nx.grid_graph([self.grid_size, self.grid_size])
+        grid_size = nx.number_of_nodes(motif_g)
+        motif_g = nx.convert_node_labels_to_integers(motif_g, first_label=0)
+        motif_edges = list(motif_g.edges())
+        motif_src, motif_dst = map(list, zip(*motif_edges))
+        motif_src, motif_dst = np.array(motif_src), np.array(motif_dst)
+
+        for motif_id in range(self.num_motifs):
+            src.extend((motif_src + n).tolist())
+            dst.extend((motif_dst + n).tolist())
+
+            # Nodes in grids belong to class 1
+            node_labels.extend([1] * grid_size)
+
+            # Attach the motif to the base tree graph
+            src.append(n)
+            dst.append(int(motif_id * spacing))
+
+            n += grid_size
+
+        g = graph((src, dst), num_nodes=n)
+
+        # Perturb the graph by adding non-self-loop random edges
+        num_real_edges = g.num_edges()
+        max_ratio = (n * (n - 1) - num_real_edges) / num_real_edges
+        assert (
+            self.perturb_ratio <= max_ratio
+        ), "perturb_ratio cannot exceed {:.4f}".format(max_ratio)
+        num_random_edges = int(num_real_edges * self.perturb_ratio)
+
+        for _ in range(num_random_edges):
+            while True:
+                u = np.random.randint(0, n)
+                v = np.random.randint(0, n)
+                if (not g.has_edges_between(u, v)) and (u != v):
+                    break
+            g.add_edges(u, v)
+
+        g.ndata["label"] = F.tensor(node_labels, F.int64)
+        g.ndata["feat"] = F.ones((n, 1), F.float32, F.cpu())
+        self._graph = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        save_graphs(str(self.graph_path), self._graph)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        graphs, _ = load_graphs(str(self.graph_path))
+        self._graph = graphs[0]
+
+    def __getitem__(self, idx):
+        assert idx == 0, "This dataset has only one graph."
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)
+
+    def __len__(self):
+        return 1
+
+    @property
+    def num_classes(self):
+        return 2
+
+
+class BA2MotifDataset(DGLBuiltinDataset):
+    r"""BA-2motifs dataset from `Parameterized Explainer for Graph Neural Network
+    <https://arxiv.org/abs/2011.04573>`__
+
+    This is a synthetic dataset for graph classification. It was generated by
+    performing the following steps in order.
+
+    - Construct 1000 base Barabási–Albert (BA) graphs.
+    - Attach house-structured network motifs to half of the base BA graphs.
+    - Attach five-node cycle motifs to the rest base BA graphs.
+    - Assign each graph to one of two classes according to the type of the attached motif.
+
+    Parameters
+    ----------
+    raw_dir : str, optional
+        Raw file directory to download and store the data. Default: ~/.dgl/
+    force_reload : bool, optional
+        Whether to reload the dataset. Default: False
+    verbose : bool, optional
+        Whether to print progress information. Default: True
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access. Default: None
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of graph classes
+
+    Examples
+    --------
+
+    >>> from dgl.data import BA2MotifDataset
+    >>> dataset = BA2MotifDataset()
+    >>> dataset.num_classes
+    2
+    >>> # Get the first graph and its label
+    >>> g, label = dataset[0]
+    >>> feat = g.ndata['feat']
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=True, transform=None
+    ):
+        super(BA2MotifDataset, self).__init__(
+            name="BA-2motifs",
+            url=_get_dgl_url("dataset/BA-2motif.pkl"),
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def download(self):
+        r"""Automatically download data."""
+        file_path = os.path.join(self.raw_dir, self.name + ".pkl")
+        download(self.url, path=file_path)
+
+    def process(self):
+        file_path = os.path.join(self.raw_dir, self.name + ".pkl")
+        with open(file_path, "rb") as f:
+            adjs, features, labels = pickle.load(f)
+
+        self.graphs = []
+        self.labels = F.tensor(labels, F.int64)
+
+        for i in range(len(adjs)):
+            g = graph(adjs[i].nonzero())
+            g.ndata["feat"] = F.zerocopy_from_numpy(features[i])
+            self.graphs.append(g)
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "{}_dgl_graph.bin".format(self.name)
+        )
+
+    def save(self):
+        label_dict = {"labels": self.labels}
+        save_graphs(str(self.graph_path), self.graphs, label_dict)
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        self.graphs, label_dict = load_graphs(str(self.graph_path))
+        self.labels = label_dict["labels"]
+
+    def __getitem__(self, idx):
+        g = self.graphs[idx]
+        if self._transform is not None:
+            g = self._transform(g)
+        return g, self.labels[idx]
+
+    def __len__(self):
+        return len(self.graphs)
+
+    @property
+    def num_classes(self):
+        return 2

python/dgl/data/tensor_serialize.py

@@ -0,0 +1,69 @@
+"""For Tensor Serialization"""
+from __future__ import absolute_import
+
+from .. import backend as F
+from .._ffi.function import _init_api
+from ..ndarray import NDArray
+
+__all__ = ["save_tensors", "load_tensors"]
+
+_init_api("dgl.data.tensor_serialize")
+
+
+def save_tensors(filename, tensor_dict):
+    """
+    Save dict of tensors to file
+
+    Parameters
+    ----------
+    filename : str
+        File name to store dict of tensors.
+    tensor_dict: dict of dgl NDArray or backend tensor
+        Python dict using string as key and tensor as value
+
+    Returns
+    ----------
+    status : bool
+        Return whether save operation succeeds
+    """
+    nd_dict = {}
+    is_empty_dict = len(tensor_dict) == 0
+    for key, value in tensor_dict.items():
+        if not isinstance(key, str):
+            raise Exception("Dict key has to be str")
+        if F.is_tensor(value):
+            nd_dict[key] = F.zerocopy_to_dgl_ndarray(value)
+        elif isinstance(value, NDArray):
+            nd_dict[key] = value
+        else:
+            raise Exception(
+                "Dict value has to be backend tensor or dgl ndarray"
+            )
+
+    return _CAPI_SaveNDArrayDict(filename, nd_dict, is_empty_dict)
+
+
+def load_tensors(filename, return_dgl_ndarray=False):
+    """
+    load dict of tensors from file
+
+    Parameters
+    ----------
+    filename : str
+        File name to load dict of tensors.
+    return_dgl_ndarray: bool
+        Whether return dict of dgl NDArrays or backend tensors
+
+    Returns
+    ---------
+    tensor_dict : dict
+        dict of tensor or ndarray based on return_dgl_ndarray flag
+    """
+    nd_dict = _CAPI_LoadNDArrayDict(filename)
+    tensor_dict = {}
+    for key, value in nd_dict.items():
+        if return_dgl_ndarray:
+            tensor_dict[key] = value
+        else:
+            tensor_dict[key] = F.zerocopy_from_dgl_ndarray(value)
+    return tensor_dict

python/dgl/data/tree.py

@@ -0,0 +1,305 @@
+"""Tree-structured data.
+Including:
+    - Stanford Sentiment Treebank
+"""
+from __future__ import absolute_import
+
+import os
+
+from collections import OrderedDict
+
+import networkx as nx
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import from_networkx
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import (
+    _get_dgl_url,
+    deprecate_property,
+    load_graphs,
+    load_info,
+    save_graphs,
+    save_info,
+)
+
+__all__ = ["SST", "SSTDataset"]
+
+
+class SSTDataset(DGLBuiltinDataset):
+    r"""Stanford Sentiment Treebank dataset.
+
+    Each sample is the constituency tree of a sentence. The leaf nodes
+    represent words. The word is a int value stored in the ``x`` feature field.
+    The non-leaf node has a special value ``PAD_WORD`` in the ``x`` field.
+    Each node also has a sentiment annotation: 5 classes (very negative,
+    negative, neutral, positive and very positive). The sentiment label is a
+    int value stored in the ``y`` feature field.
+    Official site: `<http://nlp.stanford.edu/sentiment/index.html>`_
+
+    Statistics:
+
+    - Train examples: 8,544
+    - Dev examples: 1,101
+    - Test examples: 2,210
+    - Number of classes for each node: 5
+
+    Parameters
+    ----------
+    mode : str, optional
+        Should be one of ['train', 'dev', 'test', 'tiny']
+        Default: train
+    glove_embed_file : str, optional
+        The path to pretrained glove embedding file.
+        Default: None
+    vocab_file : str, optional
+        Optional vocabulary file. If not given, the default vacabulary file is used.
+        Default: None
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset. Default: False
+    verbose : bool
+        Whether to print out progress information. Default: True.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    vocab : OrderedDict
+        Vocabulary of the dataset
+    num_classes : int
+        Number of classes for each node
+    pretrained_emb: Tensor
+        Pretrained glove embedding with respect the vocabulary.
+    vocab_size : int
+        The size of the vocabulary
+
+    Notes
+    -----
+    All the samples will be loaded and preprocessed in the memory first.
+
+    Examples
+    --------
+    >>> # get dataset
+    >>> train_data = SSTDataset()
+    >>> dev_data = SSTDataset(mode='dev')
+    >>> test_data = SSTDataset(mode='test')
+    >>> tiny_data = SSTDataset(mode='tiny')
+    >>>
+    >>> len(train_data)
+    8544
+    >>> train_data.num_classes
+    5
+    >>> glove_embed = train_data.pretrained_emb
+    >>> train_data.vocab_size
+    19536
+    >>> train_data[0]
+    Graph(num_nodes=71, num_edges=70,
+      ndata_schemes={'x': Scheme(shape=(), dtype=torch.int64), 'y': Scheme(shape=(), dtype=torch.int64), 'mask': Scheme(shape=(), dtype=torch.int64)}
+      edata_schemes={})
+    >>> for tree in train_data:
+    ...     input_ids = tree.ndata['x']
+    ...     labels = tree.ndata['y']
+    ...     mask = tree.ndata['mask']
+    ...     # your code here
+    """
+
+    PAD_WORD = -1  # special pad word id
+    UNK_WORD = -1  # out-of-vocabulary word id
+
+    def __init__(
+        self,
+        mode="train",
+        glove_embed_file=None,
+        vocab_file=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        assert mode in ["train", "dev", "test", "tiny"]
+        _url = _get_dgl_url("dataset/sst.zip")
+        self._glove_embed_file = glove_embed_file if mode == "train" else None
+        self.mode = mode
+        self._vocab_file = vocab_file
+        super(SSTDataset, self).__init__(
+            name="sst",
+            url=_url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        from nltk.corpus.reader import BracketParseCorpusReader
+
+        # load vocab file
+        self._vocab = OrderedDict()
+        vocab_file = (
+            self._vocab_file
+            if self._vocab_file is not None
+            else os.path.join(self.raw_path, "vocab.txt")
+        )
+        with open(vocab_file, encoding="utf-8") as vf:
+            for line in vf.readlines():
+                line = line.strip()
+                self._vocab[line] = len(self._vocab)
+
+        # filter glove
+        if self._glove_embed_file is not None and os.path.exists(
+            self._glove_embed_file
+        ):
+            glove_emb = {}
+            with open(self._glove_embed_file, "r", encoding="utf-8") as pf:
+                for line in pf.readlines():
+                    sp = line.split(" ")
+                    if sp[0].lower() in self._vocab:
+                        glove_emb[sp[0].lower()] = np.asarray(
+                            [float(x) for x in sp[1:]]
+                        )
+        files = ["{}.txt".format(self.mode)]
+        corpus = BracketParseCorpusReader(self.raw_path, files)
+        sents = corpus.parsed_sents(files[0])
+
+        # initialize with glove
+        pretrained_emb = []
+        fail_cnt = 0
+        for line in self._vocab.keys():
+            if self._glove_embed_file is not None and os.path.exists(
+                self._glove_embed_file
+            ):
+                if not line.lower() in glove_emb:
+                    fail_cnt += 1
+                pretrained_emb.append(
+                    glove_emb.get(
+                        line.lower(), np.random.uniform(-0.05, 0.05, 300)
+                    )
+                )
+
+        self._pretrained_emb = None
+        if self._glove_embed_file is not None and os.path.exists(
+            self._glove_embed_file
+        ):
+            self._pretrained_emb = F.tensor(np.stack(pretrained_emb, 0))
+            print(
+                "Miss word in GloVe {0:.4f}".format(
+                    1.0 * fail_cnt / len(self._pretrained_emb)
+                )
+            )
+        # build trees
+        self._trees = []
+        for sent in sents:
+            self._trees.append(self._build_tree(sent))
+
+    def _build_tree(self, root):
+        g = nx.DiGraph()
+
+        def _rec_build(nid, node):
+            for child in node:
+                cid = g.number_of_nodes()
+                if isinstance(child[0], str) or isinstance(child[0], bytes):
+                    # leaf node
+                    word = self.vocab.get(child[0].lower(), self.UNK_WORD)
+                    g.add_node(cid, x=word, y=int(child.label()), mask=1)
+                else:
+                    g.add_node(
+                        cid, x=SSTDataset.PAD_WORD, y=int(child.label()), mask=0
+                    )
+                    _rec_build(cid, child)
+                g.add_edge(cid, nid)
+
+        # add root
+        g.add_node(0, x=SSTDataset.PAD_WORD, y=int(root.label()), mask=0)
+        _rec_build(0, root)
+        ret = from_networkx(g, node_attrs=["x", "y", "mask"])
+        return ret
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, self.mode + "_dgl_graph.bin")
+
+    @property
+    def vocab_path(self):
+        return os.path.join(self.save_path, "vocab.pkl")
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path) and os.path.exists(
+            self.vocab_path
+        )
+
+    def save(self):
+        save_graphs(self.graph_path, self._trees)
+        save_info(self.vocab_path, {"vocab": self.vocab})
+        if self.pretrained_emb:
+            emb_path = os.path.join(self.save_path, "emb.pkl")
+            save_info(emb_path, {"embed": self.pretrained_emb})
+
+    def load(self):
+        emb_path = os.path.join(self.save_path, "emb.pkl")
+
+        self._trees = load_graphs(self.graph_path)[0]
+        self._vocab = load_info(self.vocab_path)["vocab"]
+        self._pretrained_emb = None
+        if os.path.exists(emb_path):
+            self._pretrained_emb = load_info(emb_path)["embed"]
+
+    @property
+    def vocab(self):
+        r"""Vocabulary
+
+        Returns
+        -------
+        OrderedDict
+        """
+        return self._vocab
+
+    @property
+    def pretrained_emb(self):
+        r"""Pre-trained word embedding, if given."""
+        return self._pretrained_emb
+
+    def __getitem__(self, idx):
+        r"""Get graph by index
+
+        Parameters
+        ----------
+        idx : int
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            graph structure, word id for each node, node labels and masks.
+
+            - ``ndata['x']``: word id of the node
+            - ``ndata['y']:`` label of the node
+            - ``ndata['mask']``: 1 if the node is a leaf, otherwise 0
+        """
+        if self._transform is None:
+            return self._trees[idx]
+        else:
+            return self._transform(self._trees[idx])
+
+    def __len__(self):
+        r"""Number of graphs in the dataset."""
+        return len(self._trees)
+
+    @property
+    def vocab_size(self):
+        r"""Vocabulary size."""
+        return len(self._vocab)
+
+    @property
+    def num_classes(self):
+        r"""Number of classes for each node."""
+        return 5
+
+
+SST = SSTDataset

python/dgl/data/tu.py

@@ -0,0 +1,532 @@
+from __future__ import absolute_import
+
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph as dgl_graph
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import load_graphs, load_info, loadtxt, save_graphs, save_info
+
+
+class LegacyTUDataset(DGLBuiltinDataset):
+    r"""LegacyTUDataset contains lots of graph kernel datasets for graph classification.
+
+    Parameters
+    ----------
+    name : str
+        Dataset Name, such as ``ENZYMES``, ``DD``, ``COLLAB``, ``MUTAG``, can be the
+        datasets name on `<https://chrsmrrs.github.io/datasets/docs/datasets/>`_.
+    use_pandas : bool
+        Numpy's file read function has performance issue when file is large,
+        using pandas can be faster.
+        Default: False
+    hidden_size : int
+        Some dataset doesn't contain features.
+        Use constant node features initialization instead, with hidden size as ``hidden_size``.
+        Default : 10
+    max_allow_node : int
+        Remove graphs that contains more nodes than ``max_allow_node``.
+        Default : None
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    max_num_node : int
+        Maximum number of nodes
+    num_classes : int
+        Number of classes
+    num_labels : numpy.int64
+        (DEPRECATED, use num_classes instead) Number of classes
+
+    Notes
+    -----
+    LegacyTUDataset uses provided node feature by default. If no feature provided, it uses one-hot node label instead.
+    If neither labels provided, it uses constant for node feature.
+
+    The dataset sorts graphs by their labels.
+    Shuffle is preferred before manual train/val split.
+
+    Examples
+    --------
+    >>> data = LegacyTUDataset('DD')
+
+    The dataset instance is an iterable
+
+    >>> len(data)
+    1178
+    >>> g, label = data[1024]
+    >>> g
+    Graph(num_nodes=88, num_edges=410,
+          ndata_schemes={'feat': Scheme(shape=(89,), dtype=torch.float32), '_ID': Scheme(shape=(), dtype=torch.int64)}
+          edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
+    >>> label
+    tensor(1)
+
+    Batch the graphs and labels for mini-batch training
+
+    >>> graphs, labels = zip(*[data[i] for i in range(16)])
+    >>> batched_graphs = dgl.batch(graphs)
+    >>> batched_labels = torch.tensor(labels)
+    >>> batched_graphs
+    Graph(num_nodes=9539, num_edges=47382,
+          ndata_schemes={'feat': Scheme(shape=(89,), dtype=torch.float32), '_ID': Scheme(shape=(), dtype=torch.int64)}
+          edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
+    """
+
+    _url = r"https://www.chrsmrrs.com/graphkerneldatasets/{}.zip"
+
+    def __init__(
+        self,
+        name,
+        use_pandas=False,
+        hidden_size=10,
+        max_allow_node=None,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        url = self._url.format(name)
+        self.hidden_size = hidden_size
+        self.max_allow_node = max_allow_node
+        self.use_pandas = use_pandas
+        super(LegacyTUDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            hash_key=(name, use_pandas, hidden_size, max_allow_node),
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.data_mode = None
+
+        if self.use_pandas:
+            import pandas as pd
+
+            DS_edge_list = self._idx_from_zero(
+                pd.read_csv(
+                    self._file_path("A"), delimiter=",", dtype=int, header=None
+                ).values
+            )
+        else:
+            DS_edge_list = self._idx_from_zero(
+                np.genfromtxt(self._file_path("A"), delimiter=",", dtype=int)
+            )
+
+        DS_indicator = self._idx_from_zero(
+            np.genfromtxt(self._file_path("graph_indicator"), dtype=int)
+        )
+        if os.path.exists(self._file_path("graph_labels")):
+            DS_graph_labels = self._idx_from_zero(
+                np.genfromtxt(self._file_path("graph_labels"), dtype=int)
+            )
+            self.num_labels = max(DS_graph_labels) + 1
+            self.graph_labels = DS_graph_labels
+        elif os.path.exists(self._file_path("graph_attributes")):
+            DS_graph_labels = np.genfromtxt(
+                self._file_path("graph_attributes"), dtype=float
+            )
+            self.num_labels = None
+            self.graph_labels = DS_graph_labels
+        else:
+            raise Exception("Unknown graph label or graph attributes")
+
+        g = dgl_graph(([], []))
+        g.add_nodes(int(DS_edge_list.max()) + 1)
+        g.add_edges(DS_edge_list[:, 0], DS_edge_list[:, 1])
+
+        node_idx_list = []
+        self.max_num_node = 0
+        for idx in range(np.max(DS_indicator) + 1):
+            node_idx = np.where(DS_indicator == idx)
+            node_idx_list.append(node_idx[0])
+            if len(node_idx[0]) > self.max_num_node:
+                self.max_num_node = len(node_idx[0])
+
+        self.graph_lists = [g.subgraph(node_idx) for node_idx in node_idx_list]
+
+        try:
+            DS_node_labels = self._idx_from_zero(
+                np.loadtxt(self._file_path("node_labels"), dtype=int)
+            )
+            g.ndata["node_label"] = F.tensor(DS_node_labels)
+            one_hot_node_labels = self._to_onehot(DS_node_labels)
+            for idxs, g in zip(node_idx_list, self.graph_lists):
+                g.ndata["feat"] = F.tensor(
+                    one_hot_node_labels[idxs, :], F.float32
+                )
+            self.data_mode = "node_label"
+        except IOError:
+            print("No Node Label Data")
+
+        try:
+            DS_node_attr = np.loadtxt(
+                self._file_path("node_attributes"), delimiter=","
+            )
+            if DS_node_attr.ndim == 1:
+                DS_node_attr = np.expand_dims(DS_node_attr, -1)
+            for idxs, g in zip(node_idx_list, self.graph_lists):
+                g.ndata["feat"] = F.tensor(DS_node_attr[idxs, :], F.float32)
+            self.data_mode = "node_attr"
+        except IOError:
+            print("No Node Attribute Data")
+
+        if "feat" not in g.ndata.keys():
+            for idxs, g in zip(node_idx_list, self.graph_lists):
+                g.ndata["feat"] = F.ones(
+                    (g.num_nodes(), self.hidden_size), F.float32, F.cpu()
+                )
+            self.data_mode = "constant"
+            if self.verbose:
+                print(
+                    "Use Constant one as Feature with hidden size {}".format(
+                        self.hidden_size
+                    )
+                )
+
+        # remove graphs that are too large by user given standard
+        # optional pre-processing steop in conformity with Rex Ying's original
+        # DiffPool implementation
+        if self.max_allow_node:
+            preserve_idx = []
+            if self.verbose:
+                print("original dataset length : ", len(self.graph_lists))
+            for i, g in enumerate(self.graph_lists):
+                if g.num_nodes() <= self.max_allow_node:
+                    preserve_idx.append(i)
+            self.graph_lists = [self.graph_lists[i] for i in preserve_idx]
+            if self.verbose:
+                print(
+                    "after pruning graphs that are too big : ",
+                    len(self.graph_lists),
+                )
+            self.graph_labels = [self.graph_labels[i] for i in preserve_idx]
+            self.max_num_node = self.max_allow_node
+        self.graph_labels = F.tensor(self.graph_labels)
+
+    def save(self):
+        label_dict = {"labels": self.graph_labels}
+        info_dict = {
+            "max_num_node": self.max_num_node,
+            "num_labels": self.num_labels,
+        }
+        save_graphs(str(self.graph_path), self.graph_lists, label_dict)
+        save_info(str(self.info_path), info_dict)
+
+    def load(self):
+        graphs, label_dict = load_graphs(str(self.graph_path))
+        info_dict = load_info(str(self.info_path))
+
+        self.graph_lists = graphs
+        self.graph_labels = label_dict["labels"]
+        self.max_num_node = info_dict["max_num_node"]
+        self.num_labels = info_dict["num_labels"]
+
+    @property
+    def graph_path(self):
+        return os.path.join(
+            self.save_path, "legacy_tu_{}_{}.bin".format(self.name, self.hash)
+        )
+
+    @property
+    def info_path(self):
+        return os.path.join(
+            self.save_path, "legacy_tu_{}_{}.pkl".format(self.name, self.hash)
+        )
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+        return False
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+            Graph with node feature stored in ``feat`` field and node label in ``node_label`` if available.
+            And its label.
+        """
+        g = self.graph_lists[idx]
+        if self._transform is not None:
+            g = self._transform(g)
+        return g, self.graph_labels[idx]
+
+    def __len__(self):
+        """Return the number of graphs in the dataset."""
+        return len(self.graph_lists)
+
+    def _file_path(self, category):
+        return os.path.join(
+            self.raw_path, self.name, "{}_{}.txt".format(self.name, category)
+        )
+
+    @staticmethod
+    def _idx_from_zero(idx_tensor):
+        return idx_tensor - np.min(idx_tensor)
+
+    @staticmethod
+    def _to_onehot(label_tensor):
+        label_num = label_tensor.shape[0]
+        assert np.min(label_tensor) == 0
+        one_hot_tensor = np.zeros((label_num, np.max(label_tensor) + 1))
+        one_hot_tensor[np.arange(label_num), label_tensor] = 1
+        return one_hot_tensor
+
+    def statistics(self):
+        return (
+            self.graph_lists[0].ndata["feat"].shape[1],
+            self.num_labels,
+            self.max_num_node,
+        )
+
+    @property
+    def num_classes(self):
+        return int(self.num_labels)
+
+
+class TUDataset(DGLBuiltinDataset):
+    r"""
+    TUDataset contains lots of graph kernel datasets for graph classification.
+
+    Parameters
+    ----------
+    name : str
+        Dataset Name, such as ``ENZYMES``, ``DD``, ``COLLAB``, ``MUTAG``, can be the
+        datasets name on `<https://chrsmrrs.github.io/datasets/docs/datasets/>`_.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    max_num_node : int
+        Maximum number of nodes
+    num_classes : int
+        Number of classes
+    num_labels : int
+        (DEPRECATED, use num_classes instead) Number of classes
+
+    Notes
+    -----
+    **IMPORTANT:** Some of the datasets have duplicate edges exist in the graphs, e.g.
+    the edges in ``IMDB-BINARY`` are all duplicated.  DGL faithfully keeps the duplicates
+    as per the original data.  Other frameworks such as PyTorch Geometric removes the
+    duplicates by default.  You can remove the duplicate edges with :func:`dgl.to_simple`.
+
+    Graphs may have node labels, node attributes, edge labels, and edge attributes,
+    varing from different dataset.
+
+    Labels are mapped to :math:`\lbrace 0,\cdots,n-1 \rbrace` where :math:`n` is the
+    number of labels (some datasets have raw labels :math:`\lbrace -1, 1 \rbrace` which
+    will be mapped to :math:`\lbrace 0, 1 \rbrace`). In previous versions, the minimum
+    label was added so that :math:`\lbrace -1, 1 \rbrace` was mapped to
+    :math:`\lbrace 0, 2 \rbrace`.
+
+    The dataset sorts graphs by their labels.
+    Shuffle is preferred before manual train/val split.
+
+    Examples
+    --------
+    >>> data = TUDataset('DD')
+
+    The dataset instance is an iterable
+
+    >>> len(data)
+    1178
+    >>> g, label = data[1024]
+    >>> g
+    Graph(num_nodes=88, num_edges=410,
+          ndata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64), 'node_labels': Scheme(shape=(1,), dtype=torch.int64)}
+          edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
+    >>> label
+    tensor([1])
+
+    Batch the graphs and labels for mini-batch training
+
+    >>> graphs, labels = zip(*[data[i] for i in range(16)])
+    >>> batched_graphs = dgl.batch(graphs)
+    >>> batched_labels = torch.tensor(labels)
+    >>> batched_graphs
+    Graph(num_nodes=9539, num_edges=47382,
+          ndata_schemes={'node_labels': Scheme(shape=(1,), dtype=torch.int64), '_ID': Scheme(shape=(), dtype=torch.int64)}
+          edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
+
+    """
+
+    _url = r"https://www.chrsmrrs.com/graphkerneldatasets/{}.zip"
+
+    def __init__(
+        self,
+        name,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        url = self._url.format(name)
+        super(TUDataset, self).__init__(
+            name=name,
+            url=url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        DS_edge_list = self._idx_from_zero(
+            loadtxt(self._file_path("A"), delimiter=",").astype(int)
+        )
+        DS_indicator = self._idx_from_zero(
+            loadtxt(self._file_path("graph_indicator"), delimiter=",").astype(
+                int
+            )
+        )
+
+        if os.path.exists(self._file_path("graph_labels")):
+            DS_graph_labels = self._idx_reset(
+                loadtxt(self._file_path("graph_labels"), delimiter=",").astype(
+                    int
+                )
+            )
+            self.num_labels = int(max(DS_graph_labels) + 1)
+            self.graph_labels = F.tensor(DS_graph_labels)
+        elif os.path.exists(self._file_path("graph_attributes")):
+            DS_graph_labels = loadtxt(
+                self._file_path("graph_attributes"), delimiter=","
+            ).astype(float)
+            self.num_labels = None
+            self.graph_labels = F.tensor(DS_graph_labels)
+        else:
+            raise Exception("Unknown graph label or graph attributes")
+
+        g = dgl_graph(([], []))
+        g.add_nodes(int(DS_edge_list.max()) + 1)
+        g.add_edges(DS_edge_list[:, 0], DS_edge_list[:, 1])
+
+        node_idx_list = []
+        self.max_num_node = 0
+        for idx in range(np.max(DS_indicator) + 1):
+            node_idx = np.where(DS_indicator == idx)
+            node_idx_list.append(node_idx[0])
+            if len(node_idx[0]) > self.max_num_node:
+                self.max_num_node = len(node_idx[0])
+
+        self.attr_dict = {
+            "node_labels": ("ndata", "node_labels"),
+            "node_attributes": ("ndata", "node_attr"),
+            "edge_labels": ("edata", "edge_labels"),
+            "edge_attributes": ("edata", "node_labels"),
+        }
+
+        for filename, field_name in self.attr_dict.items():
+            try:
+                data = loadtxt(self._file_path(filename), delimiter=",")
+                if "label" in filename:
+                    data = F.tensor(self._idx_from_zero(data))
+                else:
+                    data = F.tensor(data)
+                getattr(g, field_name[0])[field_name[1]] = data
+            except IOError:
+                pass
+
+        self.graph_lists = [g.subgraph(node_idx) for node_idx in node_idx_list]
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, "tu_{}.bin".format(self.name))
+
+    @property
+    def info_path(self):
+        return os.path.join(self.save_path, "tu_{}.pkl".format(self.name))
+
+    def save(self):
+        label_dict = {"labels": self.graph_labels}
+        info_dict = {
+            "max_num_node": self.max_num_node,
+            "num_labels": self.num_labels,
+        }
+        save_graphs(str(self.graph_path), self.graph_lists, label_dict)
+        save_info(str(self.info_path), info_dict)
+
+    def load(self):
+        graphs, label_dict = load_graphs(str(self.graph_path))
+        info_dict = load_info(str(self.info_path))
+
+        self.graph_lists = graphs
+        self.graph_labels = label_dict["labels"]
+        self.max_num_node = info_dict["max_num_node"]
+        self.num_labels = info_dict["num_labels"]
+
+    def has_cache(self):
+        if os.path.exists(self.graph_path) and os.path.exists(self.info_path):
+            return True
+        return False
+
+    def __getitem__(self, idx):
+        """Get the idx-th sample.
+
+        Parameters
+        ---------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        (:class:`dgl.DGLGraph`, Tensor)
+            Graph with node feature stored in ``feat`` field and node label in ``node_labels`` if available.
+            And its label.
+        """
+        g = self.graph_lists[idx]
+        if self._transform is not None:
+            g = self._transform(g)
+        return g, self.graph_labels[idx]
+
+    def __len__(self):
+        """Return the number of graphs in the dataset."""
+        return len(self.graph_lists)
+
+    def _file_path(self, category):
+        return os.path.join(
+            self.raw_path, self.name, "{}_{}.txt".format(self.name, category)
+        )
+
+    @staticmethod
+    def _idx_from_zero(idx_tensor):
+        return idx_tensor - np.min(idx_tensor)
+
+    @staticmethod
+    def _idx_reset(idx_tensor):
+        """Maps n unique labels to {0, ..., n-1} in an ordered fashion."""
+        labels = np.unique(idx_tensor)
+        relabel_map = {x: i for i, x in enumerate(labels)}
+        new_idx_tensor = np.vectorize(relabel_map.get)(idx_tensor)
+        return new_idx_tensor
+
+    def statistics(self):
+        return (
+            self.graph_lists[0].ndata["feat"].shape[1],
+            self.num_labels,
+            self.max_num_node,
+        )
+
+    @property
+    def num_classes(self):
+        return self.num_labels

python/dgl/data/utils.py

@@ -0,0 +1,683 @@
+"""Dataset utilities."""
+from __future__ import absolute_import
+
+import errno
+import hashlib
+import os
+import pickle
+import sys
+import warnings
+
+import networkx.algorithms as A
+
+import numpy as np
+import requests
+from tqdm.auto import tqdm
+
+from .. import backend as F
+from .graph_serialize import load_graphs, load_labels, save_graphs
+from .tensor_serialize import load_tensors, save_tensors
+
+__all__ = [
+    "loadtxt",
+    "download",
+    "check_sha1",
+    "extract_archive",
+    "get_download_dir",
+    "Subset",
+    "split_dataset",
+    "save_graphs",
+    "load_graphs",
+    "load_labels",
+    "save_tensors",
+    "load_tensors",
+    "add_nodepred_split",
+    "add_node_property_split",
+    "mask_nodes_by_property",
+]
+
+
+def loadtxt(path, delimiter, dtype=None):
+    try:
+        import pandas as pd
+
+        df = pd.read_csv(path, delimiter=delimiter, header=None)
+        return df.values
+    except ImportError:
+        warnings.warn(
+            "Pandas is not installed, now using numpy.loadtxt to load data, "
+            "which could be extremely slow. Accelerate by installing pandas"
+        )
+        return np.loadtxt(path, delimiter=delimiter)
+
+
+def _get_dgl_url(file_url):
+    """Get DGL online url for download."""
+    dgl_repo_url = "https://data.dgl.ai/"
+    repo_url = os.environ.get("DGL_REPO", dgl_repo_url)
+    if repo_url[-1] != "/":
+        repo_url = repo_url + "/"
+    return repo_url + file_url
+
+
+def split_dataset(dataset, frac_list=None, shuffle=False, random_state=None):
+    """Split dataset into training, validation and test set.
+
+    Parameters
+    ----------
+    dataset
+        We assume ``len(dataset)`` gives the number of datapoints and ``dataset[i]``
+        gives the ith datapoint.
+    frac_list : list or None, optional
+        A list of length 3 containing the fraction to use for training,
+        validation and test. If None, we will use [0.8, 0.1, 0.1].
+    shuffle : bool, optional
+        By default we perform a consecutive split of the dataset. If True,
+        we will first randomly shuffle the dataset.
+    random_state : None, int or array_like, optional
+        Random seed used to initialize the pseudo-random number generator.
+        Can be any integer between 0 and 2**32 - 1 inclusive, an array
+        (or other sequence) of such integers, or None (the default).
+        If seed is None, then RandomState will try to read data from /dev/urandom
+        (or the Windows analogue) if available or seed from the clock otherwise.
+
+    Returns
+    -------
+    list of length 3
+        Subsets for training, validation and test.
+    """
+    from itertools import accumulate
+
+    if frac_list is None:
+        frac_list = [0.8, 0.1, 0.1]
+    frac_list = np.asarray(frac_list)
+    assert np.allclose(
+        np.sum(frac_list), 1.0
+    ), "Expect frac_list sum to 1, got {:.4f}".format(np.sum(frac_list))
+    num_data = len(dataset)
+    lengths = (num_data * frac_list).astype(int)
+    lengths[-1] = num_data - np.sum(lengths[:-1])
+    if shuffle:
+        indices = np.random.RandomState(seed=random_state).permutation(num_data)
+    else:
+        indices = np.arange(num_data)
+    return [
+        Subset(dataset, indices[offset - length : offset])
+        for offset, length in zip(accumulate(lengths), lengths)
+    ]
+
+
+def download(
+    url,
+    path=None,
+    overwrite=True,
+    sha1_hash=None,
+    retries=5,
+    verify_ssl=True,
+    log=True,
+):
+    """Download a given URL.
+
+    Codes borrowed from mxnet/gluon/utils.py
+
+    Parameters
+    ----------
+    url : str
+        URL to download.
+    path : str, optional
+        Destination path to store downloaded file. By default stores to the
+        current directory with the same name as in url.
+    overwrite : bool, optional
+        Whether to overwrite the destination file if it already exists.
+        By default always overwrites the downloaded file.
+    sha1_hash : str, optional
+        Expected sha1 hash in hexadecimal digits. Will ignore existing file when hash is specified
+        but doesn't match.
+    retries : integer, default 5
+        The number of times to attempt downloading in case of failure or non 200 return codes.
+    verify_ssl : bool, default True
+        Verify SSL certificates.
+    log : bool, default True
+        Whether to print the progress for download
+
+    Returns
+    -------
+    str
+        The file path of the downloaded file.
+    """
+    if path is None:
+        fname = url.split("/")[-1]
+        # Empty filenames are invalid
+        assert fname, (
+            "Can't construct file-name from this URL. "
+            "Please set the `path` option manually."
+        )
+    else:
+        path = os.path.expanduser(path)
+        if os.path.isdir(path):
+            fname = os.path.join(path, url.split("/")[-1])
+        else:
+            fname = path
+    assert retries >= 0, "Number of retries should be at least 0"
+
+    if not verify_ssl:
+        warnings.warn(
+            "Unverified HTTPS request is being made (verify_ssl=False). "
+            "Adding certificate verification is strongly advised."
+        )
+
+    if (
+        overwrite
+        or not os.path.exists(fname)
+        or (sha1_hash and not check_sha1(fname, sha1_hash))
+    ):
+        dirname = os.path.dirname(os.path.abspath(os.path.expanduser(fname)))
+        if not os.path.exists(dirname):
+            os.makedirs(dirname)
+        while retries + 1 > 0:
+            # Disable pyling too broad Exception
+            # pylint: disable=W0703
+            try:
+                if log:
+                    print("Downloading %s from %s..." % (fname, url))
+                r = requests.get(url, stream=True, verify=verify_ssl)
+                if r.status_code != 200:
+                    raise RuntimeError("Failed downloading url %s" % url)
+                # Get the total file size.
+                total_size = int(r.headers.get("content-length", 0))
+                with tqdm(
+                    total=total_size, unit="B", unit_scale=True, desc=fname
+                ) as bar:
+                    with open(fname, "wb") as f:
+                        for chunk in r.iter_content(chunk_size=1024):
+                            if chunk:  # filter out keep-alive new chunks
+                                f.write(chunk)
+                                bar.update(len(chunk))
+                if sha1_hash and not check_sha1(fname, sha1_hash):
+                    raise UserWarning(
+                        "File {} is downloaded but the content hash does not match."
+                        " The repo may be outdated or download may be incomplete. "
+                        'If the "repo_url" is overridden, consider switching to '
+                        "the default repo.".format(fname)
+                    )
+                break
+            except Exception as e:
+                retries -= 1
+                if retries <= 0:
+                    raise e
+                else:
+                    if log:
+                        print(
+                            "download failed, retrying, {} attempt{} left".format(
+                                retries, "s" if retries > 1 else ""
+                            )
+                        )
+
+    return fname
+
+
+def check_sha1(filename, sha1_hash):
+    """Check whether the sha1 hash of the file content matches the expected hash.
+
+    Codes borrowed from mxnet/gluon/utils.py
+
+    Parameters
+    ----------
+    filename : str
+        Path to the file.
+    sha1_hash : str
+        Expected sha1 hash in hexadecimal digits.
+
+    Returns
+    -------
+    bool
+        Whether the file content matches the expected hash.
+    """
+    sha1 = hashlib.sha1()
+    with open(filename, "rb") as f:
+        while True:
+            data = f.read(1048576)
+            if not data:
+                break
+            sha1.update(data)
+
+    return sha1.hexdigest() == sha1_hash
+
+
+def extract_archive(file, target_dir, overwrite=True):
+    """Extract archive file.
+
+    Parameters
+    ----------
+    file : str
+        Absolute path of the archive file.
+    target_dir : str
+        Target directory of the archive to be uncompressed.
+    overwrite : bool, default True
+        Whether to overwrite the contents inside the directory.
+        By default always overwrites.
+    """
+    if os.path.exists(target_dir) and not overwrite:
+        return
+    print("Extracting file to {}".format(target_dir))
+    if (
+        file.endswith(".tar.gz")
+        or file.endswith(".tar")
+        or file.endswith(".tgz")
+    ):
+        import tarfile
+
+        with tarfile.open(file, "r") as archive:
+
+            def is_within_directory(directory, target):
+                abs_directory = os.path.abspath(directory)
+                abs_target = os.path.abspath(target)
+                prefix = os.path.commonprefix([abs_directory, abs_target])
+                return prefix == abs_directory
+
+            def safe_extract(
+                tar, path=".", members=None, *, numeric_owner=False
+            ):
+                for member in tar.getmembers():
+                    member_path = os.path.join(path, member.name)
+                    if not is_within_directory(path, member_path):
+                        raise Exception("Attempted Path Traversal in Tar File")
+                tar.extractall(path, members, numeric_owner=numeric_owner)
+
+            safe_extract(archive, path=target_dir)
+    elif file.endswith(".gz"):
+        import gzip
+        import shutil
+
+        with gzip.open(file, "rb") as f_in:
+            target_file = os.path.join(target_dir, os.path.basename(file)[:-3])
+            with open(target_file, "wb") as f_out:
+                shutil.copyfileobj(f_in, f_out)
+    elif file.endswith(".zip"):
+        import zipfile
+
+        with zipfile.ZipFile(file, "r") as archive:
+            archive.extractall(path=target_dir)
+    else:
+        raise Exception("Unrecognized file type: " + file)
+
+
+def get_download_dir():
+    """Get the absolute path to the download directory.
+
+    Returns
+    -------
+    dirname : str
+        Path to the download directory
+    """
+    default_dir = os.path.join(os.path.expanduser("~"), ".dgl")
+    dirname = os.environ.get("DGL_DOWNLOAD_DIR", default_dir)
+    if not os.path.exists(dirname):
+        os.makedirs(dirname)
+    return dirname
+
+
+def makedirs(path):
+    try:
+        os.makedirs(os.path.expanduser(os.path.normpath(path)))
+    except OSError as e:
+        if e.errno != errno.EEXIST and os.path.isdir(path):
+            raise e
+
+
+def save_info(path, info):
+    """Save dataset related information into disk.
+
+    Parameters
+    ----------
+    path : str
+        File to save information.
+    info : dict
+        A python dict storing information to save on disk.
+    """
+    with open(path, "wb") as pf:
+        pickle.dump(info, pf)
+
+
+def load_info(path):
+    """Load dataset related information from disk.
+
+    Parameters
+    ----------
+    path : str
+        File to load information from.
+
+    Returns
+    -------
+    info : dict
+        A python dict storing information loaded from disk.
+    """
+    with open(path, "rb") as pf:
+        info = pickle.load(pf)
+    return info
+
+
+def deprecate_property(old, new):
+    warnings.warn(
+        "Property {} will be deprecated, please use {} instead.".format(
+            old, new
+        )
+    )
+
+
+def deprecate_function(old, new):
+    warnings.warn(
+        "Function {} will be deprecated, please use {} instead.".format(
+            old, new
+        )
+    )
+
+
+def deprecate_class(old, new):
+    warnings.warn(
+        "Class {} will be deprecated, please use {} instead.".format(old, new)
+    )
+
+
+def idx2mask(idx, len):
+    """Create mask."""
+    mask = np.zeros(len)
+    mask[idx] = 1
+    return mask
+
+
+def generate_mask_tensor(mask):
+    """Generate mask tensor according to different backend
+    For torch and tensorflow, it will create a bool tensor
+    For mxnet, it will create a float tensor
+    Parameters
+    ----------
+    mask: numpy ndarray
+        input mask tensor
+    """
+    assert isinstance(mask, np.ndarray), (
+        "input for generate_mask_tensor" "should be an numpy ndarray"
+    )
+    if F.backend_name == "mxnet":
+        return F.tensor(mask, dtype=F.data_type_dict["float32"])
+    else:
+        return F.tensor(mask, dtype=F.data_type_dict["bool"])
+
+
+class Subset(object):
+    """Subset of a dataset at specified indices
+
+    Code adapted from PyTorch.
+
+    Parameters
+    ----------
+    dataset
+        dataset[i] should return the ith datapoint
+    indices : list
+        List of datapoint indices to construct the subset
+    """
+
+    def __init__(self, dataset, indices):
+        self.dataset = dataset
+        self.indices = indices
+
+    def __getitem__(self, item):
+        """Get the datapoint indexed by item
+
+        Returns
+        -------
+        tuple
+            datapoint
+        """
+        return self.dataset[self.indices[item]]
+
+    def __len__(self):
+        """Get subset size
+
+        Returns
+        -------
+        int
+            Number of datapoints in the subset
+        """
+        return len(self.indices)
+
+
+def add_nodepred_split(dataset, ratio, ntype=None):
+    """Split the given dataset into training, validation and test sets for
+    transductive node predction task.
+
+    It adds three node mask arrays ``'train_mask'``, ``'val_mask'`` and ``'test_mask'``,
+    to each graph in the dataset. Each sample in the dataset thus must be a :class:`DGLGraph`.
+
+    Fix the random seed of NumPy to make the result deterministic::
+
+        numpy.random.seed(42)
+
+    Parameters
+    ----------
+    dataset : DGLDataset
+        The dataset to modify.
+    ratio : (float, float, float)
+        Split ratios for training, validation and test sets. Must sum to one.
+    ntype : str, optional
+        The node type to add mask for.
+
+    Examples
+    --------
+    >>> dataset = dgl.data.AmazonCoBuyComputerDataset()
+    >>> print('train_mask' in dataset[0].ndata)
+    False
+    >>> dgl.data.utils.add_nodepred_split(dataset, [0.8, 0.1, 0.1])
+    >>> print('train_mask' in dataset[0].ndata)
+    True
+    """
+    if len(ratio) != 3:
+        raise ValueError(
+            f"Split ratio must be a float triplet but got {ratio}."
+        )
+    for i in range(len(dataset)):
+        g = dataset[i]
+        n = g.num_nodes(ntype)
+        idx = np.arange(0, n)
+        np.random.shuffle(idx)
+        n_train, n_val, n_test = (
+            int(n * ratio[0]),
+            int(n * ratio[1]),
+            int(n * ratio[2]),
+        )
+        train_mask = generate_mask_tensor(idx2mask(idx[:n_train], n))
+        val_mask = generate_mask_tensor(
+            idx2mask(idx[n_train : n_train + n_val], n)
+        )
+        test_mask = generate_mask_tensor(idx2mask(idx[n_train + n_val :], n))
+        g.nodes[ntype].data["train_mask"] = train_mask
+        g.nodes[ntype].data["val_mask"] = val_mask
+        g.nodes[ntype].data["test_mask"] = test_mask
+
+
+def mask_nodes_by_property(property_values, part_ratios, random_seed=None):
+    """Provide the split masks for a node split with distributional shift based on a given
+    node property, as proposed in `Evaluating Robustness and Uncertainty of Graph Models
+    Under Structural Distributional Shifts <https://arxiv.org/abs/2302.13875>`__
+
+    It considers the in-distribution (ID) and out-of-distribution (OOD) subsets of nodes.
+    The ID subset includes training, validation and testing parts, while the OOD subset
+    includes validation and testing parts. It sorts the nodes in the ascending order of
+    their property values, splits them into 5 non-intersecting parts, and creates 5
+    associated node mask arrays:
+        - 3 for the ID nodes: ``'in_train_mask'``, ``'in_valid_mask'``, ``'in_test_mask'``,
+        - and 2 for the OOD nodes: ``'out_valid_mask'``, ``'out_test_mask'``.
+
+    Parameters
+    ----------
+    property_values : numpy ndarray
+        The node property (float) values by which the dataset will be split.
+        The length of the array must be equal to the number of nodes in graph.
+    part_ratios : list
+        A list of 5 ratios for training, ID validation, ID test,
+        OOD validation, OOD testing parts. The values in the list must sum to one.
+    random_seed : int, optional
+        Random seed to fix for the initial permutation of nodes. It is
+        used to create a random order for the nodes that have the same
+        property values or belong to the ID subset. (default: None)
+
+    Returns
+    ----------
+    split_masks : dict
+        A python dict storing the mask names as keys and the corresponding
+        node mask arrays as values.
+
+    Examples
+    --------
+    >>> num_nodes = 1000
+    >>> property_values = np.random.uniform(size=num_nodes)
+    >>> part_ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
+    >>> split_masks = dgl.data.utils.mask_nodes_by_property(property_values, part_ratios)
+    >>> print('in_valid_mask' in split_masks)
+    True
+    """
+
+    num_nodes = len(property_values)
+    part_sizes = np.round(num_nodes * np.array(part_ratios)).astype(int)
+    part_sizes[-1] -= np.sum(part_sizes) - num_nodes
+
+    generator = np.random.RandomState(random_seed)
+    permutation = generator.permutation(num_nodes)
+
+    node_indices = np.arange(num_nodes)[permutation]
+    property_values = property_values[permutation]
+    in_distribution_size = np.sum(part_sizes[:3])
+
+    node_indices_ordered = node_indices[np.argsort(property_values)]
+    node_indices_ordered[:in_distribution_size] = generator.permutation(
+        node_indices_ordered[:in_distribution_size]
+    )
+
+    sections = np.cumsum(part_sizes)
+    node_split = np.split(node_indices_ordered, sections)[:-1]
+    mask_names = [
+        "in_train_mask",
+        "in_valid_mask",
+        "in_test_mask",
+        "out_valid_mask",
+        "out_test_mask",
+    ]
+    split_masks = {}
+
+    for mask_name, node_indices in zip(mask_names, node_split):
+        split_mask = idx2mask(node_indices, num_nodes)
+        split_masks[mask_name] = generate_mask_tensor(split_mask)
+
+    return split_masks
+
+
+def add_node_property_split(
+    dataset, part_ratios, property_name, ascending=True, random_seed=None
+):
+    """Create a node split with distributional shift based on a given node property,
+    as proposed in `Evaluating Robustness and Uncertainty of Graph Models Under
+    Structural Distributional Shifts <https://arxiv.org/abs/2302.13875>`__
+
+    It splits the nodes of each graph in the given dataset into 5 non-intersecting
+    parts based on their structural properties. This can be used for transductive node
+    prediction task with distributional shifts.
+
+    It considers the in-distribution (ID) and out-of-distribution (OOD) subsets of nodes.
+    The ID subset includes training, validation and testing parts, while the OOD subset
+    includes validation and testing parts. As a result, it creates 5 associated node mask
+    arrays for each graph:
+        - 3 for the ID nodes: ``'in_train_mask'``, ``'in_valid_mask'``, ``'in_test_mask'``,
+        - and 2 for the OOD nodes: ``'out_valid_mask'``, ``'out_test_mask'``.
+
+    This function implements 3 particular strategies for inducing distributional shifts
+    in graph — based on **popularity**, **locality** or **density**.
+
+    Parameters
+    ----------
+    dataset : :class:`~DGLDataset` or list of :class:`~dgl.DGLGraph`
+        The dataset to induce structural distributional shift.
+    part_ratios : list
+        A list of 5 ratio values for training, ID validation, ID test,
+        OOD validation and OOD test parts. The values must sum to 1.0.
+    property_name : str
+        The name of the node property to be used, which must be
+        ``'popularity'``, ``'locality'`` or ``'density'``.
+    ascending : bool, optional
+        Whether to sort nodes in the ascending order of the node property,
+        so that nodes with greater values of the property are considered
+        to be OOD (default: True)
+    random_seed : int, optional
+        Random seed to fix for the initial permutation of nodes. It is
+        used to create a random order for the nodes that have the same
+        property values or belong to the ID subset. (default: None)
+
+    Examples
+    --------
+    >>> dataset = dgl.data.AmazonCoBuyComputerDataset()
+    >>> print('in_valid_mask' in dataset[0].ndata)
+    False
+    >>> part_ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
+    >>> property_name = 'popularity'
+    >>> dgl.data.utils.add_node_property_split(dataset, part_ratios, property_name)
+    >>> print('in_valid_mask' in dataset[0].ndata)
+    True
+    """
+
+    assert property_name in [
+        "popularity",
+        "locality",
+        "density",
+    ], "The name of property has to be 'popularity', 'locality', or 'density'"
+
+    assert len(part_ratios) == 5, "part_ratios must contain 5 values"
+
+    import networkx as nx
+
+    for idx in range(len(dataset)):
+        graph_dgl = dataset[idx]
+        graph_nx = nx.Graph(graph_dgl.to_networkx())
+
+        compute_property_fn = _property_name_to_compute_fn[property_name]
+        property_values = compute_property_fn(graph_nx, ascending)
+
+        node_masks = mask_nodes_by_property(
+            property_values, part_ratios, random_seed
+        )
+
+        for mask_name, node_mask in node_masks.items():
+            graph_dgl.ndata[mask_name] = node_mask
+
+
+def _compute_popularity_property(graph_nx, ascending=True):
+    direction = -1 if ascending else 1
+    property_values = direction * np.array(list(A.pagerank(graph_nx).values()))
+    return property_values
+
+
+def _compute_locality_property(graph_nx, ascending=True):
+    num_nodes = graph_nx.number_of_nodes()
+    pagerank_values = np.array(list(A.pagerank(graph_nx).values()))
+
+    personalization = dict(zip(range(num_nodes), [0.0] * num_nodes))
+    personalization[np.argmax(pagerank_values)] = 1.0
+
+    direction = -1 if ascending else 1
+    property_values = direction * np.array(
+        list(A.pagerank(graph_nx, personalization=personalization).values())
+    )
+    return property_values
+
+
+def _compute_density_property(graph_nx, ascending=True):
+    direction = -1 if ascending else 1
+    property_values = direction * np.array(
+        list(A.clustering(graph_nx).values())
+    )
+    return property_values
+
+
+_property_name_to_compute_fn = {
+    "popularity": _compute_popularity_property,
+    "locality": _compute_locality_property,
+    "density": _compute_density_property,
+}

python/dgl/data/wikics.py

@@ -0,0 +1,173 @@
+"""Wiki-CS Dataset"""
+import itertools
+import json
+import os
+
+import numpy as np
+
+from .. import backend as F
+from ..convert import graph
+from ..transforms import reorder_graph, to_bidirected
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, generate_mask_tensor, load_graphs, save_graphs
+
+
+class WikiCSDataset(DGLBuiltinDataset):
+    r"""Wiki-CS is a Wikipedia-based dataset for node classification from `Wiki-CS: A Wikipedia-Based
+    Benchmark for Graph Neural Networks <https://arxiv.org/abs/2007.02901v2>`_
+
+    The dataset consists of nodes corresponding to Computer Science articles, with edges based on
+    hyperlinks and 10 classes representing different branches of the field.
+
+    WikiCS dataset statistics:
+
+    - Nodes: 11,701
+    - Edges: 431,726 (note that the original dataset has 216,123 edges but DGL adds
+      the reverse edges and removes the duplicate edges, hence with a different number)
+    - Number of classes: 10
+    - Node feature size: 300
+    - Number of different train, validation, stopping splits: 20
+    - Number of test split: 1
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+    >>> from dgl.data import WikiCSDataset
+    >>> dataset = WikiCSDataset()
+    >>> dataset.num_classes
+    10
+    >>> g = dataset[0]
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>> # get node labels
+    >>> labels = g.ndata['label']
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> stopping_mask = g.ndata['stopping_mask']
+    >>> test_mask = g.ndata['test_mask']
+    >>> # The shape of train, val and stopping masks are (num_nodes, num_splits).
+    >>> # The num_splits is the number of different train, validation, stopping splits.
+    >>> # Due to the number of test spilt is 1, the shape of test mask is (num_nodes,).
+    >>> print(train_mask.shape, val_mask.shape, stopping_mask.shape)
+    (11701, 20) (11701, 20) (11701, 20)
+    >>> print(test_mask.shape)
+    (11701,)
+    """
+
+    def __init__(
+        self, raw_dir=None, force_reload=False, verbose=False, transform=None
+    ):
+        _url = _get_dgl_url("dataset/wiki_cs.zip")
+        super(WikiCSDataset, self).__init__(
+            name="wiki_cs",
+            raw_dir=raw_dir,
+            url=_url,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """process raw data to graph, labels and masks"""
+        with open(os.path.join(self.raw_path, "data.json")) as f:
+            data = json.load(f)
+        features = F.tensor(np.array(data["features"]), dtype=F.float32)
+        labels = F.tensor(np.array(data["labels"]), dtype=F.int64)
+
+        train_masks = np.array(data["train_masks"], dtype=bool).T
+        val_masks = np.array(data["val_masks"], dtype=bool).T
+        stopping_masks = np.array(data["stopping_masks"], dtype=bool).T
+        test_mask = np.array(data["test_mask"], dtype=bool)
+
+        edges = [[(i, j) for j in js] for i, js in enumerate(data["links"])]
+        edges = np.array(list(itertools.chain(*edges)))
+        src, dst = edges[:, 0], edges[:, 1]
+
+        g = graph((src, dst))
+        g = to_bidirected(g)
+
+        g.ndata["feat"] = features
+        g.ndata["label"] = labels
+        g.ndata["train_mask"] = generate_mask_tensor(train_masks)
+        g.ndata["val_mask"] = generate_mask_tensor(val_masks)
+        g.ndata["stopping_mask"] = generate_mask_tensor(stopping_masks)
+        g.ndata["test_mask"] = generate_mask_tensor(test_mask)
+
+        g = reorder_graph(
+            g,
+            node_permute_algo="rcmk",
+            edge_permute_algo="dst",
+            store_ids=False,
+        )
+
+        self._graph = g
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        return os.path.exists(graph_path)
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        g, _ = load_graphs(graph_path)
+        self._graph = g[0]
+
+    @property
+    def num_classes(self):
+        return 10
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return 1
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index, WikiCSDataset has only one graph object
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``ndata['feat']``: node features
+            - ``ndata['label']``: node labels
+            - ``ndata['train_mask']``: train mask is for retrieving the nodes for training.
+            - ``ndata['val_mask']``: val mask is for retrieving the nodes for hyperparameter tuning.
+            - ``ndata['stopping_mask']``: stopping mask is for retrieving the nodes for early stopping criterion.
+            - ``ndata['test_mask']``: test mask is for retrieving the nodes for testing.
+
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)

python/dgl/data/yelp.py

@@ -0,0 +1,177 @@
+"""Yelp Dataset"""
+import json
+import os
+
+import numpy as np
+import scipy.sparse as sp
+
+from .. import backend as F
+from ..convert import from_scipy
+from ..transforms import reorder_graph
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, generate_mask_tensor, load_graphs, save_graphs
+
+
+class YelpDataset(DGLBuiltinDataset):
+    r"""Yelp dataset for node classification from `GraphSAINT: Graph Sampling Based Inductive
+    Learning Method <https://arxiv.org/abs/1907.04931>`_
+
+    The task of this dataset is categorizing types of businesses based on customer reviewers and
+    friendship.
+
+    Yelp dataset statistics:
+
+    - Nodes: 716,847
+    - Edges: 13,954,819
+    - Number of classes: 100 (Multi-class)
+    - Node feature size: 300
+
+    Parameters
+    ----------
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: ~/.dgl/
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False
+    verbose : bool
+        Whether to print out progress information.
+        Default: False
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+    reorder : bool
+        Whether to reorder the graph using :func:`~dgl.reorder_graph`.
+        Default: False.
+
+    Attributes
+    ----------
+    num_classes : int
+        Number of node classes
+
+    Examples
+    --------
+    >>> dataset = YelpDataset()
+    >>> dataset.num_classes
+    100
+    >>> g = dataset[0]
+    >>> # get node feature
+    >>> feat = g.ndata['feat']
+    >>> # get node labels
+    >>> labels = g.ndata['label']
+    >>> # get data split
+    >>> train_mask = g.ndata['train_mask']
+    >>> val_mask = g.ndata['val_mask']
+    >>> test_mask = g.ndata['test_mask']
+    """
+
+    def __init__(
+        self,
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+        reorder=False,
+    ):
+        _url = _get_dgl_url("dataset/yelp.zip")
+        self._reorder = reorder
+        super(YelpDataset, self).__init__(
+            name="yelp",
+            raw_dir=raw_dir,
+            url=_url,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        """process raw data to graph, labels and masks"""
+        coo_adj = sp.load_npz(os.path.join(self.raw_path, "adj_full.npz"))
+        g = from_scipy(coo_adj)
+
+        features = np.load(os.path.join(self.raw_path, "feats.npy"))
+        features = F.tensor(features, dtype=F.float32)
+
+        y = [-1] * features.shape[0]
+        with open(os.path.join(self.raw_path, "class_map.json")) as f:
+            class_map = json.load(f)
+            for key, item in class_map.items():
+                y[int(key)] = item
+        labels = F.tensor(np.array(y), dtype=F.int64)
+
+        with open(os.path.join(self.raw_path, "role.json")) as f:
+            role = json.load(f)
+
+        train_mask = np.zeros(features.shape[0], dtype=bool)
+        train_mask[role["tr"]] = True
+
+        val_mask = np.zeros(features.shape[0], dtype=bool)
+        val_mask[role["va"]] = True
+
+        test_mask = np.zeros(features.shape[0], dtype=bool)
+        test_mask[role["te"]] = True
+
+        g.ndata["feat"] = features
+        g.ndata["label"] = labels
+        g.ndata["train_mask"] = generate_mask_tensor(train_mask)
+        g.ndata["val_mask"] = generate_mask_tensor(val_mask)
+        g.ndata["test_mask"] = generate_mask_tensor(test_mask)
+
+        if self._reorder:
+            self._graph = reorder_graph(
+                g,
+                node_permute_algo="rcmk",
+                edge_permute_algo="dst",
+                store_ids=False,
+            )
+        else:
+            self._graph = g
+
+    def has_cache(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        return os.path.exists(graph_path)
+
+    def save(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        save_graphs(graph_path, self._graph)
+
+    def load(self):
+        graph_path = os.path.join(self.save_path, "dgl_graph.bin")
+        g, _ = load_graphs(graph_path)
+        self._graph = g[0]
+
+    @property
+    def num_classes(self):
+        return 100
+
+    def __len__(self):
+        r"""The number of graphs in the dataset."""
+        return 1
+
+    def __getitem__(self, idx):
+        r"""Get graph object
+
+        Parameters
+        ----------
+        idx : int
+            Item index, FlickrDataset has only one graph object
+
+        Returns
+        -------
+        :class:`dgl.DGLGraph`
+
+            The graph contains:
+
+            - ``ndata['label']``: node label
+            - ``ndata['feat']``: node feature
+            - ``ndata['train_mask']``: mask for training node set
+            - ``ndata['val_mask']``: mask for validation node set
+            - ``ndata['test_mask']``: mask for test node set
+
+        """
+        assert idx == 0, "This dataset has only one graph"
+        if self._transform is None:
+            return self._graph
+        else:
+            return self._transform(self._graph)

python/dgl/data/zinc.py

@@ -0,0 +1,137 @@
+import os
+
+from .dgl_dataset import DGLBuiltinDataset
+from .utils import _get_dgl_url, load_graphs
+
+
+class ZINCDataset(DGLBuiltinDataset):
+    r"""ZINC dataset for the graph regression task.
+
+    A subset (12K) of ZINC molecular graphs (250K) dataset is used to
+    regress a molecular property known as the constrained solubility.
+    For each molecular graph, the node features are the types of heavy
+    atoms, between which the edge features are the types of bonds.
+    Each graph contains 9-37 nodes and 16-84 edges.
+
+    Reference `<https://arxiv.org/pdf/2003.00982.pdf>`_
+
+    Statistics:
+
+    Train examples: 10,000
+    Valid examples: 1,000
+    Test examples: 1,000
+    Average number of nodes: 23.16
+    Average number of edges: 39.83
+    Number of atom types: 28
+    Number of bond types: 4
+
+    Parameters
+    ----------
+    mode : str, optional
+        Should be chosen from ["train", "valid", "test"]
+        Default: "train".
+    raw_dir : str
+        Raw file directory to download/contains the input data directory.
+        Default: "~/.dgl/".
+    force_reload : bool
+        Whether to reload the dataset.
+        Default: False.
+    verbose : bool
+        Whether to print out progress information.
+        Default: False.
+    transform : callable, optional
+        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
+        a transformed version. The :class:`~dgl.DGLGraph` object will be
+        transformed before every access.
+
+    Attributes
+    ----------
+    num_atom_types : int
+        Number of atom types.
+    num_bond_types : int
+        Number of bond types.
+
+    Examples
+    ---------
+    >>> from dgl.data import ZINCDataset
+
+    >>> training_set = ZINCDataset(mode="train")
+    >>> training_set.num_atom_types
+    28
+    >>> len(training_set)
+    10000
+    >>> graph, label = training_set[0]
+    >>> graph
+    Graph(num_nodes=29, num_edges=64,
+        ndata_schemes={'feat': Scheme(shape=(), dtype=torch.int64)}
+        edata_schemes={'feat': Scheme(shape=(), dtype=torch.int64)})
+    """
+
+    def __init__(
+        self,
+        mode="train",
+        raw_dir=None,
+        force_reload=False,
+        verbose=False,
+        transform=None,
+    ):
+        self._url = _get_dgl_url("dataset/ZINC12k.zip")
+        self.mode = mode
+
+        super(ZINCDataset, self).__init__(
+            name="zinc",
+            url=self._url,
+            raw_dir=raw_dir,
+            force_reload=force_reload,
+            verbose=verbose,
+            transform=transform,
+        )
+
+    def process(self):
+        self.load()
+
+    @property
+    def graph_path(self):
+        return os.path.join(self.save_path, "ZincDGL_{}.bin".format(self.mode))
+
+    def has_cache(self):
+        return os.path.exists(self.graph_path)
+
+    def load(self):
+        self._graphs, self._labels = load_graphs(self.graph_path)
+
+    @property
+    def num_atom_types(self):
+        return 28
+
+    @property
+    def num_bond_types(self):
+        return 4
+
+    def __len__(self):
+        return len(self._graphs)
+
+    def __getitem__(self, idx):
+        r"""Get one example by index.
+
+        Parameters
+        ----------
+        idx : int
+            The sample index.
+
+        Returns
+        -------
+        dgl.DGLGraph
+            Each graph contains:
+
+            - ``ndata['feat']``: Types of heavy atoms as node features
+            - ``edata['feat']``: Types of bonds as edge features
+
+        Tensor
+            Constrained solubility as graph label
+        """
+        labels = self._labels["g_label"]
+        if self._transform is None:
+            return self._graphs[idx], labels[idx]
+        else:
+            return self._transform(self._graphs[idx]), labels[idx]

python/dgl/dataloading/__init__.py

@@ -0,0 +1,14 @@
+"""Package for dataloaders and samplers."""
+
+from .. import backend as F
+from . import negative_sampler
+from .base import *
+from .cluster_gcn import *
+from .graphsaint import *
+from .labor_sampler import *
+from .neighbor_sampler import *
+from .shadow import *
+
+if F.get_preferred_backend() == "pytorch":
+    from .spot_target import *
+    from .dataloader import *

python/dgl/dataloading/base.py

@@ -0,0 +1,658 @@
+"""Base classes and functionalities for dataloaders"""
+import inspect
+from collections.abc import Mapping
+
+from .. import backend as F
+from ..base import EID, NID
+from ..convert import heterograph
+from ..frame import LazyFeature
+from ..transforms import compact_graphs
+from ..utils import context_of, recursive_apply
+
+
+def _set_lazy_features(x, xdata, feature_names):
+    if feature_names is None:
+        return
+    if not isinstance(feature_names, Mapping):
+        xdata.update({k: LazyFeature(k) for k in feature_names})
+    else:
+        for type_, names in feature_names.items():
+            x[type_].data.update({k: LazyFeature(k) for k in names})
+
+
+def set_node_lazy_features(g, feature_names):
+    """Assign lazy features to the ``ndata`` of the input graph for prefetching optimization.
+
+    When used in a :class:`~dgl.dataloading.Sampler`, lazy features mark which data
+    should be fetched before computation in model. See :ref:`guide-minibatch-prefetching`
+    for a detailed explanation.
+
+    If the graph is homogeneous, this is equivalent to:
+
+    .. code:: python
+
+       g.ndata.update({k: LazyFeature(k, g.ndata[dgl.NID]) for k in feature_names})
+
+    If the graph is heterogeneous, this is equivalent to:
+
+    .. code:: python
+
+        for type_, names in feature_names.items():
+            g.nodes[type_].data.update(
+                {k: LazyFeature(k, g.nodes[type_].data[dgl.NID]) for k in names})
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The graph.
+    feature_names : list[str] or dict[str, list[str]]
+        The feature names to prefetch.
+
+    See also
+    --------
+    dgl.LazyFeature
+    """
+    return _set_lazy_features(g.nodes, g.ndata, feature_names)
+
+
+def set_edge_lazy_features(g, feature_names):
+    """Assign lazy features to the ``edata`` of the input graph for prefetching optimization.
+
+    When used in a :class:`~dgl.dataloading.Sampler`, lazy features mark which data
+    should be fetched before computation in model. See :ref:`guide-minibatch-prefetching`
+    for a detailed explanation.
+
+    If the graph is homogeneous, this is equivalent to:
+
+    .. code:: python
+
+       g.edata.update({k: LazyFeature(k, g.edata[dgl.EID]) for k in feature_names})
+
+    If the graph is heterogeneous, this is equivalent to:
+
+    .. code:: python
+
+        for type_, names in feature_names.items():
+            g.edges[type_].data.update(
+                {k: LazyFeature(k, g.edges[type_].data[dgl.EID]) for k in names})
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The graph.
+    feature_names : list[str] or dict[etype, list[str]]
+        The feature names to prefetch. The ``etype`` key is either a string
+        or a triplet.
+
+    See also
+    --------
+    dgl.LazyFeature
+    """
+    return _set_lazy_features(g.edges, g.edata, feature_names)
+
+
+def set_src_lazy_features(g, feature_names):
+    """Assign lazy features to the ``srcdata`` of the input graph for prefetching optimization.
+
+    When used in a :class:`~dgl.dataloading.Sampler`, lazy features mark which data
+    should be fetched before computation in model. See :ref:`guide-minibatch-prefetching`
+    for a detailed explanation.
+
+    If the graph is homogeneous, this is equivalent to:
+
+    .. code:: python
+
+       g.srcdata.update({k: LazyFeature(k, g.srcdata[dgl.NID]) for k in feature_names})
+
+    If the graph is heterogeneous, this is equivalent to:
+
+    .. code:: python
+
+        for type_, names in feature_names.items():
+            g.srcnodes[type_].data.update(
+                {k: LazyFeature(k, g.srcnodes[type_].data[dgl.NID]) for k in names})
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The graph.
+    feature_names : list[str] or dict[str, list[str]]
+        The feature names to prefetch.
+
+    See also
+    --------
+    dgl.LazyFeature
+    """
+    return _set_lazy_features(g.srcnodes, g.srcdata, feature_names)
+
+
+def set_dst_lazy_features(g, feature_names):
+    """Assign lazy features to the ``dstdata`` of the input graph for prefetching optimization.
+
+    When used in a :class:`~dgl.dataloading.Sampler`, lazy features mark which data
+    should be fetched before computation in model. See :ref:`guide-minibatch-prefetching`
+    for a detailed explanation.
+
+    If the graph is homogeneous, this is equivalent to:
+
+    .. code:: python
+
+       g.dstdata.update({k: LazyFeature(k, g.dstdata[dgl.NID]) for k in feature_names})
+
+    If the graph is heterogeneous, this is equivalent to:
+
+    .. code:: python
+
+        for type_, names in feature_names.items():
+            g.dstnodes[type_].data.update(
+                {k: LazyFeature(k, g.dstnodes[type_].data[dgl.NID]) for k in names})
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The graph.
+    feature_names : list[str] or dict[str, list[str]]
+        The feature names to prefetch.
+
+    See also
+    --------
+    dgl.LazyFeature
+    """
+    return _set_lazy_features(g.dstnodes, g.dstdata, feature_names)
+
+
+class Sampler(object):
+    """Base class for graph samplers.
+
+    All graph samplers must subclass this class and override the ``sample``
+    method.
+
+    .. code:: python
+
+        from dgl.dataloading import Sampler
+
+        class SubgraphSampler(Sampler):
+            def __init__(self):
+                super().__init__()
+
+            def sample(self, g, indices):
+                return g.subgraph(indices)
+    """
+
+    def sample(self, g, indices):
+        """Abstract sample method.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        indices : object
+            Any object representing the indices selected in the current minibatch.
+        """
+        raise NotImplementedError
+
+
+class BlockSampler(Sampler):
+    """Base class for sampling mini-batches in the form of Message-passing
+    Flow Graphs (MFGs).
+
+    It provides prefetching options to fetch the node features for the first MFG's ``srcdata``,
+    the node labels for the last MFG's ``dstdata`` and the edge features of all MFG's ``edata``.
+
+    Parameters
+    ----------
+    prefetch_node_feats : list[str] or dict[str, list[str]], optional
+        The node data to prefetch for the first MFG.
+
+        DGL will populate the first layer's MFG's ``srcnodes`` and ``srcdata`` with
+        the node data of the given names from the original graph.
+    prefetch_labels : list[str] or dict[str, list[str]], optional
+        The node data to prefetch for the last MFG.
+
+        DGL will populate the last layer's MFG's ``dstnodes`` and ``dstdata`` with
+        the node data of the given names from the original graph.
+    prefetch_edge_feats : list[str] or dict[etype, list[str]], optional
+        The edge data names to prefetch for all the MFGs.
+
+        DGL will populate every MFG's ``edges`` and ``edata`` with the edge data
+        of the given names from the original graph.
+    output_device : device, optional
+        The device of the output subgraphs or MFGs.  Default is the same as the
+        minibatch of seed nodes.
+    """
+
+    def __init__(
+        self,
+        prefetch_node_feats=None,
+        prefetch_labels=None,
+        prefetch_edge_feats=None,
+        output_device=None,
+    ):
+        super().__init__()
+        self.prefetch_node_feats = prefetch_node_feats or []
+        self.prefetch_labels = prefetch_labels or []
+        self.prefetch_edge_feats = prefetch_edge_feats or []
+        self.output_device = output_device
+
+    def sample_blocks(self, g, seed_nodes, exclude_eids=None):
+        """Generates a list of blocks from the given seed nodes.
+
+        This function must return a triplet where the first element is the input node IDs
+        for the first GNN layer (a tensor or a dict of tensors for heterogeneous graphs),
+        the second element is the output node IDs for the last GNN layer, and the third
+        element is the said list of blocks.
+        """
+        raise NotImplementedError
+
+    def assign_lazy_features(self, result):
+        """Assign lazy features for prefetching."""
+        input_nodes, output_nodes, blocks = result
+        set_src_lazy_features(blocks[0], self.prefetch_node_feats)
+        set_dst_lazy_features(blocks[-1], self.prefetch_labels)
+        for block in blocks:
+            set_edge_lazy_features(block, self.prefetch_edge_feats)
+        return input_nodes, output_nodes, blocks
+
+    def sample(
+        self, g, seed_nodes, exclude_eids=None
+    ):  # pylint: disable=arguments-differ
+        """Sample a list of blocks from the given seed nodes."""
+        result = self.sample_blocks(g, seed_nodes, exclude_eids=exclude_eids)
+        return self.assign_lazy_features(result)
+
+
+def _find_exclude_eids_with_reverse_id(g, eids, reverse_eid_map):
+    if isinstance(eids, Mapping):
+        eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
+        exclude_eids = {
+            k: F.cat([v, F.gather_row(reverse_eid_map[k], v)], 0)
+            for k, v in eids.items()
+        }
+    else:
+        exclude_eids = F.cat([eids, F.gather_row(reverse_eid_map, eids)], 0)
+    return exclude_eids
+
+
+def _find_exclude_eids_with_reverse_types(g, eids, reverse_etype_map):
+    exclude_eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
+    reverse_etype_map = {
+        g.to_canonical_etype(k): g.to_canonical_etype(v)
+        for k, v in reverse_etype_map.items()
+    }
+    for k, v in reverse_etype_map.items():
+        if k in exclude_eids:
+            if v in exclude_eids:
+                exclude_eids[v] = F.unique(
+                    F.cat((exclude_eids[k], exclude_eids[v]), dim=0)
+                )
+            else:
+                exclude_eids[v] = exclude_eids[k]
+    return exclude_eids
+
+
+def _find_exclude_eids(g, exclude_mode, eids, **kwargs):
+    if exclude_mode is None:
+        return None
+    elif callable(exclude_mode):
+        return exclude_mode(eids)
+    elif F.is_tensor(exclude_mode) or (
+        isinstance(exclude_mode, Mapping)
+        and all(F.is_tensor(v) for v in exclude_mode.values())
+    ):
+        return exclude_mode
+    elif exclude_mode == "self":
+        return eids
+    elif exclude_mode == "reverse_id":
+        return _find_exclude_eids_with_reverse_id(
+            g, eids, kwargs["reverse_eid_map"]
+        )
+    elif exclude_mode == "reverse_types":
+        return _find_exclude_eids_with_reverse_types(
+            g, eids, kwargs["reverse_etype_map"]
+        )
+    else:
+        raise ValueError("unsupported mode {}".format(exclude_mode))
+
+
+def find_exclude_eids(
+    g,
+    seed_edges,
+    exclude,
+    reverse_eids=None,
+    reverse_etypes=None,
+    output_device=None,
+):
+    """Find all edge IDs to exclude according to :attr:`exclude_mode`.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The graph.
+    exclude :
+        Can be either of the following,
+
+        None (default)
+            Does not exclude any edge.
+
+        'self'
+            Exclude the given edges themselves but nothing else.
+
+        'reverse_id'
+            Exclude all edges specified in ``eids``, as well as their reverse edges
+            of the same edge type.
+
+            The mapping from each edge ID to its reverse edge ID is specified in
+            the keyword argument ``reverse_eid_map``.
+
+            This mode assumes that the reverse of an edge with ID ``e`` and type
+            ``etype`` will have ID ``reverse_eid_map[e]`` and type ``etype``.
+
+        'reverse_types'
+            Exclude all edges specified in ``eids``, as well as their reverse
+            edges of the corresponding edge types.
+
+            The mapping from each edge type to its reverse edge type is specified
+            in the keyword argument ``reverse_etype_map``.
+
+            This mode assumes that the reverse of an edge with ID ``e`` and type ``etype``
+            will have ID ``e`` and type ``reverse_etype_map[etype]``.
+
+        callable
+            Any function that takes in a single argument :attr:`seed_edges` and returns
+            a tensor or dict of tensors.
+    eids : Tensor or dict[etype, Tensor]
+        The edge IDs.
+    reverse_eids : Tensor or dict[etype, Tensor]
+        The mapping from edge ID to its reverse edge ID.
+    reverse_etypes : dict[etype, etype]
+        The mapping from edge etype to its reverse edge type.
+    output_device : device
+        The device of the output edge IDs.
+    """
+    exclude_eids = _find_exclude_eids(
+        g,
+        exclude,
+        seed_edges,
+        reverse_eid_map=reverse_eids,
+        reverse_etype_map=reverse_etypes,
+    )
+    if exclude_eids is not None and output_device is not None:
+        exclude_eids = recursive_apply(
+            exclude_eids, lambda x: F.copy_to(x, output_device)
+        )
+    return exclude_eids
+
+
+class EdgePredictionSampler(Sampler):
+    """Sampler class that wraps an existing sampler for node classification into another
+    one for edge classification or link prediction.
+
+    See also
+    --------
+    as_edge_prediction_sampler
+    """
+
+    def __init__(
+        self,
+        sampler,
+        exclude=None,
+        reverse_eids=None,
+        reverse_etypes=None,
+        negative_sampler=None,
+        prefetch_labels=None,
+    ):
+        super().__init__()
+        # Check if the sampler's sample method has an optional third argument.
+        argspec = inspect.getfullargspec(sampler.sample)
+        if len(argspec.args) < 4:  # ['self', 'g', 'indices', 'exclude_eids']
+            raise TypeError(
+                "This sampler does not support edge or link prediction; please add an"
+                "optional third argument for edge IDs to exclude in its sample() method."
+            )
+        self.reverse_eids = reverse_eids
+        self.reverse_etypes = reverse_etypes
+        self.exclude = exclude
+        self.sampler = sampler
+        self.negative_sampler = negative_sampler
+        self.prefetch_labels = prefetch_labels or []
+        self.output_device = sampler.output_device
+
+    def _build_neg_graph(self, g, seed_edges):
+        neg_srcdst = self.negative_sampler(g, seed_edges)
+        if not isinstance(neg_srcdst, Mapping):
+            assert len(g.canonical_etypes) == 1, (
+                "graph has multiple or no edge types; "
+                "please return a dict in negative sampler."
+            )
+            neg_srcdst = {g.canonical_etypes[0]: neg_srcdst}
+
+        dtype = F.dtype(list(neg_srcdst.values())[0][0])
+        ctx = context_of(seed_edges) if seed_edges is not None else g.device
+        neg_edges = {
+            etype: neg_srcdst.get(
+                etype,
+                (
+                    F.copy_to(F.tensor([], dtype), ctx=ctx),
+                    F.copy_to(F.tensor([], dtype), ctx=ctx),
+                ),
+            )
+            for etype in g.canonical_etypes
+        }
+        neg_pair_graph = heterograph(
+            neg_edges, {ntype: g.num_nodes(ntype) for ntype in g.ntypes}
+        )
+        return neg_pair_graph
+
+    def assign_lazy_features(self, result):
+        """Assign lazy features for prefetching."""
+        pair_graph = result[1]
+        set_edge_lazy_features(pair_graph, self.prefetch_labels)
+        # In-place updates
+        return result
+
+    def sample(self, g, seed_edges):  # pylint: disable=arguments-differ
+        """Samples a list of blocks, as well as a subgraph containing the sampled
+        edges from the original graph.
+
+        If :attr:`negative_sampler` is given, also returns another graph containing the
+        negative pairs as edges.
+        """
+        if isinstance(seed_edges, Mapping):
+            seed_edges = {
+                g.to_canonical_etype(k): v for k, v in seed_edges.items()
+            }
+        exclude = self.exclude
+        pair_graph = g.edge_subgraph(
+            seed_edges, relabel_nodes=False, output_device=self.output_device
+        )
+        eids = pair_graph.edata[EID]
+
+        if self.negative_sampler is not None:
+            neg_graph = self._build_neg_graph(g, seed_edges)
+            pair_graph, neg_graph = compact_graphs([pair_graph, neg_graph])
+        else:
+            pair_graph = compact_graphs(pair_graph)
+
+        pair_graph.edata[EID] = eids
+        seed_nodes = pair_graph.ndata[NID]
+
+        exclude_eids = find_exclude_eids(
+            g,
+            seed_edges,
+            exclude,
+            self.reverse_eids,
+            self.reverse_etypes,
+            self.output_device,
+        )
+
+        input_nodes, _, blocks = self.sampler.sample(
+            g, seed_nodes, exclude_eids
+        )
+
+        if self.negative_sampler is None:
+            return self.assign_lazy_features((input_nodes, pair_graph, blocks))
+        else:
+            return self.assign_lazy_features(
+                (input_nodes, pair_graph, neg_graph, blocks)
+            )
+
+
+def as_edge_prediction_sampler(
+    sampler,
+    exclude=None,
+    reverse_eids=None,
+    reverse_etypes=None,
+    negative_sampler=None,
+    prefetch_labels=None,
+):
+    """Create an edge-wise sampler from a node-wise sampler.
+
+    For each batch of edges, the sampler applies the provided node-wise sampler to
+    their source and destination nodes to extract subgraphs. It also generates negative
+    edges if a negative sampler is provided, and extract subgraphs for their incident
+    nodes as well.
+
+    For each iteration, the sampler will yield
+
+    * A tensor of input nodes necessary for computing the representation on edges, or
+      a dictionary of node type names and such tensors.
+
+    * A subgraph that contains only the edges in the minibatch and their incident nodes.
+      Note that the graph has an identical metagraph with the original graph.
+
+    * If a negative sampler is given, another graph that contains the "negative edges",
+      connecting the source and destination nodes yielded from the given negative sampler.
+
+    * The subgraphs or MFGs returned by the provided node-wise sampler, generated
+      from the incident nodes of the edges in the minibatch (as well as those of the
+      negative edges if applicable).
+
+    Parameters
+    ----------
+    sampler : Sampler
+        The node-wise sampler object.  It additionally requires that the :attr:`sample`
+        method must have an optional third argument :attr:`exclude_eids` representing the
+        edge IDs to exclude from neighborhood.  The argument will be either a tensor
+        for homogeneous graphs or a dict of edge types and tensors for heterogeneous
+        graphs.
+    exclude : Union[str, callable], optional
+        Whether and how to exclude dependencies related to the sampled edges in the
+        minibatch.  Possible values are
+
+        * None, for not excluding any edges.
+
+        * ``self``, for excluding the edges in the current minibatch.
+
+        * ``reverse_id``, for excluding not only the edges in the current minibatch but
+          also their reverse edges according to the ID mapping in the argument
+          :attr:`reverse_eids`.
+
+        * ``reverse_types``, for excluding not only the edges in the current minibatch
+          but also their reverse edges stored in another type according to
+          the argument :attr:`reverse_etypes`.
+
+        * User-defined exclusion rule. It is a callable with edges in the current
+          minibatch as a single argument and should return the edges to be excluded.
+    reverse_eids : Tensor or dict[etype, Tensor], optional
+        A tensor of reverse edge ID mapping.  The i-th element indicates the ID of
+        the i-th edge's reverse edge.
+
+        If the graph is heterogeneous, this argument requires a dictionary of edge
+        types and the reverse edge ID mapping tensors.
+    reverse_etypes : dict[etype, etype], optional
+        The mapping from the original edge types to their reverse edge types.
+    negative_sampler : callable, optional
+        The negative sampler.
+    prefetch_labels : list[str] or dict[etype, list[str]], optional
+        The edge labels to prefetch for the returned positive pair graph.
+
+        See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
+
+    Examples
+    --------
+    The following example shows how to train a 3-layer GNN for edge classification on a
+    set of edges ``train_eid`` on a homogeneous undirected graph. Each node takes
+    messages from all neighbors.
+
+    Given an array of source node IDs ``src`` and another array of destination
+    node IDs ``dst``, the following code creates a bidirectional graph:
+
+    >>> g = dgl.graph((torch.cat([src, dst]), torch.cat([dst, src])))
+
+    Edge :math:`i`'s reverse edge in the graph above is edge :math:`i + |E|`. Therefore, we can
+    create a reverse edge mapping ``reverse_eids`` by:
+
+    >>> E = len(src)
+    >>> reverse_eids = torch.cat([torch.arange(E, 2 * E), torch.arange(0, E)])
+
+    By passing ``reverse_eids`` to the edge sampler, the edges in the current mini-batch and their
+    reversed edges will be excluded from the extracted subgraphs to avoid information leakage.
+
+    >>> sampler = dgl.dataloading.as_edge_prediction_sampler(
+    ...     dgl.dataloading.NeighborSampler([15, 10, 5]),
+    ...     exclude='reverse_id', reverse_eids=reverse_eids)
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, train_eid, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, pair_graph, blocks in dataloader:
+    ...     train_on(input_nodes, pair_graph, blocks)
+
+    For link prediction, one can provide a negative sampler to sample negative edges.
+    The code below uses DGL's :class:`~dgl.dataloading.negative_sampler.Uniform`
+    to generate 5 negative samples per edge:
+
+    >>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
+    >>> sampler = dgl.dataloading.as_edge_prediction_sampler(
+    ...     dgl.dataloading.NeighborSampler([15, 10, 5]),
+    ...     sampler, exclude='reverse_id', reverse_eids=reverse_eids,
+    ...     negative_sampler=neg_sampler)
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, train_eid, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
+    ...     train_on(input_nodes, pair_graph, neg_pair_graph, blocks)
+
+    For heterogeneous graphs, reverse edges may belong to a different relation. For example,
+    the relations "user-click-item" and "item-click-by-user" in the graph below are
+    mutual reverse.
+
+    >>> g = dgl.heterograph({
+    ...     ('user', 'click', 'item'): (user, item),
+    ...     ('item', 'clicked-by', 'user'): (item, user)})
+
+    To correctly exclude edges from each mini-batch, set ``exclude='reverse_types'`` and
+    pass a dictionary ``{'click': 'clicked-by', 'clicked-by': 'click'}`` to the
+    ``reverse_etypes`` argument.
+
+    >>> sampler = dgl.dataloading.as_edge_prediction_sampler(
+    ...     dgl.dataloading.NeighborSampler([15, 10, 5]),
+    ...     exclude='reverse_types',
+    ...     reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'})
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, {'click': train_eid}, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, pair_graph, blocks in dataloader:
+    ...     train_on(input_nodes, pair_graph, blocks)
+
+    For link prediction, provide a negative sampler to generate negative samples:
+
+    >>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
+    >>> sampler = dgl.dataloading.as_edge_prediction_sampler(
+    ...     dgl.dataloading.NeighborSampler([15, 10, 5]),
+    ...     exclude='reverse_types',
+    ...     reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'},
+    ...     negative_sampler=neg_sampler)
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, train_eid, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
+    ...     train_on(input_nodes, pair_graph, neg_pair_graph, blocks)
+    """
+    return EdgePredictionSampler(
+        sampler,
+        exclude=exclude,
+        reverse_eids=reverse_eids,
+        reverse_etypes=reverse_etypes,
+        negative_sampler=negative_sampler,
+        prefetch_labels=prefetch_labels,
+    )

python/dgl/dataloading/capped_neighbor_sampler.py

@@ -0,0 +1,190 @@
+"""Capped neighbor sampler."""
+from collections import defaultdict
+
+import numpy as np
+import torch
+
+from ..sampling.utils import EidExcluder
+from .base import Sampler, set_edge_lazy_features, set_node_lazy_features
+
+
+class CappedNeighborSampler(Sampler):
+    """Subgraph sampler that sets an upper bound on the number of nodes included in
+    each layer of the sampled subgraph. At each layer, the frontier is randomly
+    subsampled. Rare node types can also be upsampled by taking the scaled square
+    root of the sampling probabilities. The sampler returns the subgraph induced by
+    all the sampled nodes.
+
+    This code was contributed by a community member
+    ([@ayushnoori](https://github.com/ayushnoori)). There aren't currently any unit
+    tests in place to verify its functionality, so please be cautious if you need
+    to make any changes to the code's logic.
+
+    Parameters
+    ----------
+    fanouts : list[int] or dict[etype, int]
+        List of neighbors to sample per edge type for each GNN layer, with the i-th
+        element being the fanout for the i-th GNN layer.
+        - If only a single integer is provided, DGL assumes that every edge type
+            will have the same fanout.
+        - If -1 is provided for one edge type on one layer, then all inbound edges
+            of that edge type will be included.
+    fixed_k : int
+            The number of nodes to sample for each GNN layer.
+    upsample_rare_types : bool
+        Whether or not to upsample rare node types.
+    replace : bool, default True
+        Whether to sample with replacement.
+    prob : str, optional
+        If given, the probability of each neighbor being sampled is proportional
+        to the edge feature value with the given name in ``g.edata``. The feature must be
+        a scalar on each edge.
+    """
+
+    def __init__(
+        self,
+        fanouts,
+        fixed_k,
+        upsample_rare_types,
+        replace=False,
+        prob=None,
+        prefetch_node_feats=None,
+        prefetch_edge_feats=None,
+        output_device=None,
+    ):
+        super().__init__()
+        self.fanouts = fanouts
+        self.replace = replace
+        self.fixed_k = fixed_k
+        self.upsample_rare_types = upsample_rare_types
+        self.prob = prob
+        self.prefetch_node_feats = prefetch_node_feats
+        self.prefetch_edge_feats = prefetch_edge_feats
+        self.output_device = output_device
+
+    def sample(
+        self, g, indices, exclude_eids=None
+    ):  # pylint: disable=arguments-differ
+        """Sampling function.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph to sample from.
+        indices : Tensor or dict[str, Tensor]
+            Nodes which induce the subgraph.
+        exclude_eids : Tensor or dict[etype, Tensor], optional
+            The edges to exclude from the sampled subgraph.
+
+        Returns
+        -------
+        input_nodes : Tensor or dict[str, Tensor]
+            The node IDs inducing the subgraph.
+        output_nodes : Tensor or dict[str, Tensor]
+            The node IDs that are sampled in this minibatch.
+        subg : DGLGraph
+            The subgraph itself.
+        """
+
+        # Define empty dictionary to store reached nodes.
+        output_nodes = indices
+        all_reached_nodes = [indices]
+
+        # Iterate over fanout.
+        for fanout in reversed(self.fanouts):
+
+            # Sample frontier.
+            frontier = g.sample_neighbors(
+                indices,
+                fanout,
+                output_device=self.output_device,
+                replace=self.replace,
+                prob=self.prob,
+                exclude_edges=exclude_eids,
+            )
+
+            # Get reached nodes.
+            curr_reached = defaultdict(list)
+            for c_etype in frontier.canonical_etypes:
+                (src_type, _, _) = c_etype
+                src, _ = frontier.edges(etype=c_etype)
+                curr_reached[src_type].append(src)
+
+            # De-duplication.
+            curr_reached = {
+                ntype: torch.unique(torch.cat(srcs))
+                for ntype, srcs in curr_reached.items()
+            }
+
+            # Generate type sampling probabilties.
+            type_count = {
+                node_type: indices.shape[0]
+                for node_type, indices in curr_reached.items()
+            }
+            total_count = sum(type_count.values())
+            probs = {
+                node_type: count / total_count
+                for node_type, count in type_count.items()
+            }
+
+            # Upsample rare node types.
+            if self.upsample_rare_types:
+
+                # Take scaled square root of probabilities.
+                prob_dist = list(probs.values())
+                prob_dist = np.sqrt(prob_dist)
+                prob_dist = prob_dist / prob_dist.sum()
+
+                # Update probabilities.
+                probs = {
+                    node_type: prob_dist[i]
+                    for i, node_type in enumerate(probs.keys())
+                }
+
+            # Generate node counts per type.
+            n_per_type = {
+                node_type: int(self.fixed_k * prob)
+                for node_type, prob in probs.items()
+            }
+            remainder = self.fixed_k - sum(n_per_type.values())
+            for _ in range(remainder):
+                node_type = np.random.choice(
+                    list(probs.keys()), p=list(probs.values())
+                )
+                n_per_type[node_type] += 1
+
+            # Downsample nodes.
+            curr_reached_k = {}
+            for node_type, node_ids in curr_reached.items():
+
+                # Get number of total nodes and number to sample.
+                num_nodes = node_ids.shape[0]
+                n_to_sample = min(num_nodes, n_per_type[node_type])
+
+                # Downsample nodes of current type.
+                random_indices = torch.randperm(num_nodes)[:n_to_sample]
+                curr_reached_k[node_type] = node_ids[random_indices]
+
+            # Update seed nodes.
+            indices = curr_reached_k
+            all_reached_nodes.append(curr_reached_k)
+
+        # Merge all reached nodes before sending to `DGLGraph.subgraph`.
+        merged_nodes = {}
+        for ntype in g.ntypes:
+            merged_nodes[ntype] = torch.unique(
+                torch.cat(
+                    [reached.get(ntype, []) for reached in all_reached_nodes]
+                )
+            )
+        subg = g.subgraph(
+            merged_nodes, relabel_nodes=True, output_device=self.output_device
+        )
+
+        if exclude_eids is not None:
+            subg = EidExcluder(exclude_eids)(subg)
+
+        set_node_lazy_features(subg, self.prefetch_node_feats)
+        set_edge_lazy_features(subg, self.prefetch_edge_feats)
+
+        return indices, output_nodes, subg

python/dgl/dataloading/cluster_gcn.py

@@ -0,0 +1,155 @@
+"""Cluster-GCN samplers."""
+import os
+import pickle
+
+import numpy as np
+
+from .. import backend as F
+from ..base import DGLError
+from ..partition import metis_partition_assignment
+from .base import Sampler, set_edge_lazy_features, set_node_lazy_features
+
+
+class ClusterGCNSampler(Sampler):
+    """Cluster sampler from `Cluster-GCN: An Efficient Algorithm for Training
+    Deep and Large Graph Convolutional Networks
+    <https://arxiv.org/abs/1905.07953>`__
+
+    This sampler first partitions the graph with METIS partitioning, then it caches the nodes of
+    each partition to a file within the given cache directory.
+
+    The sampler then selects the graph partitions according to the provided
+    partition IDs, take the union of all nodes in those partitions, and return an
+    induced subgraph in its :attr:`sample` method.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The original graph.  Must be homogeneous and on CPU.
+    k : int
+        The number of partitions.
+    cache_path : str
+        The path to the cache directory for storing the partition result.
+    balance_ntypes, balkance_edges, mode :
+        Passed to :func:`dgl.metis_partition_assignment`.
+    prefetch_ndata : list[str], optional
+        The node data to prefetch for the subgraph.
+
+        See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
+    prefetch_edata : list[str], optional
+        The edge data to prefetch for the subgraph.
+
+        See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
+    output_device : device, optional
+        The device of the output subgraphs or MFGs.  Default is the same as the
+        minibatch of partition indices.
+
+    Examples
+    --------
+    **Node classification**
+
+    With this sampler, the data loader will accept the list of partition IDs as
+    indices to iterate over.  For instance, the following code first splits the
+    graph into 1000 partitions using METIS, and at each iteration it gets a subgraph
+    induced by the nodes covered by 20 randomly selected partitions.
+
+    >>> num_parts = 1000
+    >>> sampler = dgl.dataloading.ClusterGCNSampler(g, num_parts)
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, torch.arange(num_parts), sampler,
+    ...     batch_size=20, shuffle=True, drop_last=False, num_workers=4)
+    >>> for subg in dataloader:
+    ...     train_on(subg)
+    """
+
+    def __init__(
+        self,
+        g,
+        k,
+        cache_path="cluster_gcn.pkl",
+        balance_ntypes=None,
+        balance_edges=False,
+        mode="k-way",
+        prefetch_ndata=None,
+        prefetch_edata=None,
+        output_device=None,
+    ):
+        super().__init__()
+        if os.path.exists(cache_path):
+            try:
+                with open(cache_path, "rb") as f:
+                    (
+                        self.partition_offset,
+                        self.partition_node_ids,
+                    ) = pickle.load(f)
+            except (EOFError, TypeError, ValueError):
+                raise DGLError(
+                    f"The contents in the cache file {cache_path} is invalid. "
+                    f"Please remove the cache file {cache_path} or specify another path."
+                )
+            if len(self.partition_offset) != k + 1:
+                raise DGLError(
+                    f"Number of partitions in the cache does not match the value of k. "
+                    f"Please remove the cache file {cache_path} or specify another path."
+                )
+            if len(self.partition_node_ids) != g.num_nodes():
+                raise DGLError(
+                    f"Number of nodes in the cache does not match the given graph. "
+                    f"Please remove the cache file {cache_path} or specify another path."
+                )
+        else:
+            partition_ids = metis_partition_assignment(
+                g,
+                k,
+                balance_ntypes=balance_ntypes,
+                balance_edges=balance_edges,
+                mode=mode,
+            )
+            partition_ids = F.asnumpy(partition_ids)
+            partition_node_ids = np.argsort(partition_ids)
+            partition_size = F.zerocopy_from_numpy(
+                np.bincount(partition_ids, minlength=k)
+            )
+            partition_offset = F.zerocopy_from_numpy(
+                np.insert(np.cumsum(partition_size), 0, 0)
+            )
+            partition_node_ids = F.zerocopy_from_numpy(partition_node_ids)
+            with open(cache_path, "wb") as f:
+                pickle.dump((partition_offset, partition_node_ids), f)
+            self.partition_offset = partition_offset
+            self.partition_node_ids = partition_node_ids
+
+        self.prefetch_ndata = prefetch_ndata or []
+        self.prefetch_edata = prefetch_edata or []
+        self.output_device = output_device
+
+    def sample(self, g, partition_ids):  # pylint: disable=arguments-differ
+        """Sampling function.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph to sample from.
+        partition_ids : Tensor
+            A 1-D integer tensor of partition IDs.
+
+        Returns
+        -------
+        DGLGraph
+            The sampled subgraph.
+        """
+        node_ids = F.cat(
+            [
+                self.partition_node_ids[
+                    self.partition_offset[i] : self.partition_offset[i + 1]
+                ]
+                for i in F.asnumpy(partition_ids)
+            ],
+            0,
+        )
+        sg = g.subgraph(
+            node_ids, relabel_nodes=True, output_device=self.output_device
+        )
+        set_node_lazy_features(sg, self.prefetch_ndata)
+        set_edge_lazy_features(sg, self.prefetch_edata)
+        return sg

python/dgl/dataloading/graphsaint.py

@@ -0,0 +1,163 @@
+"""GraphSAINT samplers."""
+from ..base import DGLError
+from ..random import choice
+from ..sampling import pack_traces, random_walk
+from .base import Sampler, set_edge_lazy_features, set_node_lazy_features
+
+try:
+    import torch
+except ImportError:
+    pass
+
+
+class SAINTSampler(Sampler):
+    """Random node/edge/walk sampler from
+    `GraphSAINT: Graph Sampling Based Inductive Learning Method
+    <https://arxiv.org/abs/1907.04931>`__
+
+    For each call, the sampler samples a node subset and then returns a node induced subgraph.
+    There are three options for sampling node subsets:
+
+    - For :attr:`'node'` sampler, the probability to sample a node is in proportion
+      to its out-degree.
+    - The :attr:`'edge'` sampler first samples an edge subset and then use the
+      end nodes of the edges.
+    - The :attr:`'walk'` sampler uses the nodes visited by random walks. It uniformly selects
+      a number of root nodes and then performs a fixed-length random walk from each root node.
+
+    Parameters
+    ----------
+    mode : str
+        The sampler to use, which can be :attr:`'node'`, :attr:`'edge'`, or :attr:`'walk'`.
+    budget : int or tuple[int]
+        Sampler configuration.
+
+        - For :attr:`'node'` sampler, budget specifies the number of nodes
+          in each sampled subgraph.
+        - For :attr:`'edge'` sampler, budget specifies the number of edges
+          to sample for inducing a subgraph.
+        - For :attr:`'walk'` sampler, budget is a tuple. budget[0] specifies
+          the number of root nodes to generate random walks. budget[1] specifies
+          the length of a random walk.
+
+    cache : bool, optional
+        If False, it will not cache the probability arrays for sampling. Setting
+        it to False is required if you want to use the sampler across different graphs.
+    prefetch_ndata : list[str], optional
+        The node data to prefetch for the subgraph.
+
+        See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
+    prefetch_edata : list[str], optional
+        The edge data to prefetch for the subgraph.
+
+        See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
+    output_device : device, optional
+        The device of the output subgraphs.
+
+    Examples
+    --------
+
+    >>> import torch
+    >>> from dgl.dataloading import SAINTSampler, DataLoader
+    >>> num_iters = 1000
+    >>> sampler = SAINTSampler(mode='node', budget=6000)
+    >>> # Assume g.ndata['feat'] and g.ndata['label'] hold node features and labels
+    >>> dataloader = DataLoader(g, torch.arange(num_iters), sampler, num_workers=4)
+    >>> for subg in dataloader:
+    ...     train_on(subg)
+    """
+
+    def __init__(
+        self,
+        mode,
+        budget,
+        cache=True,
+        prefetch_ndata=None,
+        prefetch_edata=None,
+        output_device="cpu",
+    ):
+        super().__init__()
+        self.budget = budget
+        if mode == "node":
+            self.sampler = self.node_sampler
+        elif mode == "edge":
+            self.sampler = self.edge_sampler
+        elif mode == "walk":
+            self.sampler = self.walk_sampler
+        else:
+            raise DGLError(
+                f"Expect mode to be 'node', 'edge' or 'walk', got {mode}."
+            )
+
+        self.cache = cache
+        self.prob = None
+        self.prefetch_ndata = prefetch_ndata or []
+        self.prefetch_edata = prefetch_edata or []
+        self.output_device = output_device
+
+    def node_sampler(self, g):
+        """Node ID sampler for random node sampler"""
+        # Alternatively, this can be realized by uniformly sampling an edge subset,
+        # and then take the src node of the sampled edges. However, the number of edges
+        # is typically much larger than the number of nodes.
+        if self.cache and self.prob is not None:
+            prob = self.prob
+        else:
+            prob = g.out_degrees().float().clamp(min=1)
+            if self.cache:
+                self.prob = prob
+        return (
+            torch.multinomial(prob, num_samples=self.budget, replacement=True)
+            .unique()
+            .type(g.idtype)
+        )
+
+    def edge_sampler(self, g):
+        """Node ID sampler for random edge sampler"""
+        src, dst = g.edges()
+        if self.cache and self.prob is not None:
+            prob = self.prob
+        else:
+            in_deg = g.in_degrees().float().clamp(min=1)
+            out_deg = g.out_degrees().float().clamp(min=1)
+            # We can reduce the sample space by half if graphs are always symmetric.
+            prob = 1.0 / in_deg[dst.long()] + 1.0 / out_deg[src.long()]
+            prob /= prob.sum()
+            if self.cache:
+                self.prob = prob
+        sampled_edges = torch.unique(
+            choice(len(prob), size=self.budget, prob=prob)
+        )
+        sampled_nodes = torch.cat([src[sampled_edges], dst[sampled_edges]])
+        return sampled_nodes.unique().type(g.idtype)
+
+    def walk_sampler(self, g):
+        """Node ID sampler for random walk sampler"""
+        num_roots, walk_length = self.budget
+        sampled_roots = torch.randint(0, g.num_nodes(), (num_roots,))
+        traces, types = random_walk(g, nodes=sampled_roots, length=walk_length)
+        sampled_nodes, _, _, _ = pack_traces(traces, types)
+        return sampled_nodes.unique().type(g.idtype)
+
+    def sample(self, g, indices):
+        """Sampling function
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph to sample from.
+        indices : Tensor
+            Placeholder not used.
+
+        Returns
+        -------
+        DGLGraph
+            The sampled subgraph.
+        """
+        node_ids = self.sampler(g)
+        sg = g.subgraph(
+            node_ids, relabel_nodes=True, output_device=self.output_device
+        )
+        set_node_lazy_features(sg, self.prefetch_ndata)
+        set_edge_lazy_features(sg, self.prefetch_edata)
+        return sg

python/dgl/dataloading/labor_sampler.py

@@ -0,0 +1,255 @@
+#
+#   Copyright (c) 2022 by Contributors
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+#   Based off of neighbor_sampler.py
+#
+
+"""Data loading components for labor sampling"""
+from numpy.random import default_rng
+
+from .. import backend as F
+from ..base import EID, NID
+from ..random import choice
+from ..transforms import to_block
+from .base import BlockSampler
+
+
+class LaborSampler(BlockSampler):
+    """Sampler that builds computational dependency of node representations via
+    labor sampling for multilayer GNN from the NeurIPS 2023 paper
+    `Layer-Neighbor Sampling -- Defusing Neighborhood Explosion in GNNs
+    <https://arxiv.org/abs/2210.13339>`__
+
+    This sampler will make every node gather messages from a fixed number of
+    neighbors per edge type. The neighbors are picked uniformly with default
+    parameters. For every vertex t that will be considered to be sampled, there
+    will be a single random variate r_t.
+
+    Parameters
+    ----------
+    fanouts : list[int] or list[dict[etype, int]]
+        List of neighbors to sample per edge type for each GNN layer, with the
+        i-th element being the fanout for the i-th GNN layer.
+
+        If only a single integer is provided, DGL assumes that every edge type
+        will have the same fanout.
+
+        If -1 is provided for one edge type on one layer, then all inbound edges
+        of that edge type will be included.
+    edge_dir : str, default ``'in'``
+        Can be either ``'in'`` where the neighbors will be sampled according to
+        incoming edges, or ``'out'`` otherwise, same as
+        :func:`dgl.sampling.sample_neighbors`.
+    prob : str, optional
+        If given, the probability of each neighbor being sampled is proportional
+        to the edge feature value with the given name in ``g.edata``.
+        The feature must be a scalar on each edge. In this case, the returned
+        blocks edata include ``'edge_weights'`` that needs to be used in the
+        message passing operation.
+    importance_sampling : int, default ``0``
+        Whether to use importance sampling or uniform sampling, use of negative
+        values optimizes importance sampling probabilities until convergence
+        while use of positive values runs optimization steps that many times.
+        If the value is i, then LABOR-i variant is used. When used with a
+        nonzero parameter, the returned blocks edata include ``'edge_weights'``
+        that needs to be used in the message passing operation.
+    layer_dependency : bool, default ``False``
+        Specifies whether different layers should use same random variates.
+        Results into a reduction in the number of vertices sampled, but may
+        degrade the quality slightly.
+    batch_dependency : int, default ``1``
+        Specifies whether different minibatches should use similar random
+        variates. Results in a higher temporal access locality of sampled
+        vertices, but may degrade the quality slightly.
+    prefetch_node_feats : list[str] or dict[ntype, list[str]], optional
+        The source node data to prefetch for the first MFG, corresponding to the
+        input node features necessary for the first GNN layer.
+    prefetch_labels : list[str] or dict[ntype, list[str]], optional
+        The destination node data to prefetch for the last MFG, corresponding to
+        the node labels of the minibatch.
+    prefetch_edge_feats : list[str] or dict[etype, list[str]], optional
+        The edge data names to prefetch for all the MFGs, corresponding to the
+        edge features necessary for all GNN layers.
+    output_device : device, optional
+        The device of the output subgraphs or MFGs.  Default is the same as the
+        minibatch of seed nodes.
+
+    Examples
+    --------
+    **Node classification**
+
+    To train a 3-layer GNN for node classification on a set of nodes
+    ``train_nid`` on a homogeneous graph where each node takes messages from
+    5, 10, 15 neighbors for the first, second, and third layer respectively
+    (assuming the backend is PyTorch):
+
+    >>> sampler = dgl.dataloading.LaborSampler([5, 10, 15])
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, train_nid, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     train_on(blocks)
+
+    If training on a heterogeneous graph and you want different number of
+    neighbors for each edge type, one should instead provide a list of dicts.
+    Each dict would specify the number of neighbors to pick per edge type.
+
+    >>> sampler = dgl.dataloading.LaborSampler([
+    ...     {('user', 'follows', 'user'): 5,
+    ...      ('user', 'plays', 'game'): 4,
+    ...      ('game', 'played-by', 'user'): 3}] * 3)
+
+    If you would like non-uniform labor sampling:
+
+    >>> # any non-negative 1D vector works
+    >>> g.edata['p'] = torch.rand(g.num_edges())
+    >>> sampler = dgl.dataloading.LaborSampler([5, 10, 15], prob='p')
+
+    **Edge classification and link prediction**
+
+    This class can also work for edge classification and link prediction
+    together with :func:`as_edge_prediction_sampler`.
+
+    >>> sampler = dgl.dataloading.LaborSampler([5, 10, 15])
+    >>> sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, train_eid, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+
+    See the documentation :func:`as_edge_prediction_sampler` for more details.
+
+    Notes
+    -----
+    For the concept of MFGs, please refer to
+    :ref:`User Guide Section 6 <guide-minibatch>` and
+    :doc:`Minibatch Training Tutorials
+    <tutorials/large/L0_neighbor_sampling_overview>`.
+    """
+
+    def __init__(
+        self,
+        fanouts,
+        edge_dir="in",
+        prob=None,
+        importance_sampling=0,
+        layer_dependency=False,
+        batch_dependency=1,
+        prefetch_node_feats=None,
+        prefetch_labels=None,
+        prefetch_edge_feats=None,
+        output_device=None,
+    ):
+        super().__init__(
+            prefetch_node_feats=prefetch_node_feats,
+            prefetch_labels=prefetch_labels,
+            prefetch_edge_feats=prefetch_edge_feats,
+            output_device=output_device,
+        )
+        self.fanouts = fanouts
+        self.edge_dir = edge_dir
+        self.prob = prob
+        self.importance_sampling = importance_sampling
+        self.layer_dependency = layer_dependency
+        self.cnt = F.zeros(2, F.int64, F.cpu())
+        self.cnt[0] = -1
+        self.cnt[1] = batch_dependency
+        self.random_seed = F.zeros(
+            2 if self.cnt[1] > 1 else 1, F.int64, F.cpu()
+        )
+        self.set_seed(None if batch_dependency > 0 else choice(1e18, 1).item())
+
+    def set_seed(self, random_seed=None):
+        """Updates the underlying seed for the sampler
+
+        Calling this function enforces the sampling algorithm to use the same
+        seed on every edge type. This can reduce the number of nodes being
+        sampled because the passed random_seed makes it so that for any seed
+        vertex ``s`` and its neighbor ``t``, the rolled random variate ``r_t``
+        is the same for any instance of this class with the same random seed.
+        When sampling as part of the same batch, one would want identical seeds
+        so that LABOR can globally sample. One example is that for heterogenous
+        graphs, there is a single random seed passed for each edge type. This
+        will sample much fewer vertices compared to having unique random seeds
+        for each edge type. If one called this function individually for each
+        edge type for a heterogenous graph with different random seeds, then it
+        would run LABOR locally for each edge type, resulting into a larger
+        number of vertices being sampled.
+
+        If this function is called without any parameters, we get the random
+        seed by getting a random number from DGL. Call this function if multiple
+        instances of LaborSampler are used to sample as part of a single batch.
+
+        Parameters
+        ----------
+        random_seed : int, default ``None``
+            The random seed to be used for next sampling call.
+        """
+        if random_seed is None:
+            self.cnt[0] += 1
+            if self.cnt[1] > 0 and self.cnt[0] % self.cnt[1] == 0:
+                if self.cnt[0] <= 0 or self.cnt[1] <= 1:
+                    if not hasattr(self, "rng"):
+                        self.rng = default_rng(choice(1e18, 1).item())
+                    self.random_seed[0] = self.rng.integers(1e18)
+                    if self.cnt[1] > 1:
+                        self.random_seed[1] = self.rng.integers(1e18)
+                else:
+                    self.random_seed[0] = self.random_seed[1]
+                    self.random_seed[1] = self.rng.integers(1e18)
+        else:
+            self.rng = default_rng(random_seed)
+            self.random_seed[0] = self.rng.integers(1e18)
+            if self.cnt[1] > 1:
+                self.random_seed[1] = self.rng.integers(1e18)
+            self.cnt[0] = 0
+
+    def sample_blocks(self, g, seed_nodes, exclude_eids=None):
+        output_nodes = seed_nodes
+        blocks = []
+        for i, fanout in enumerate(reversed(self.fanouts)):
+            random_seed_i = F.zerocopy_to_dgl_ndarray(
+                self.random_seed + (i if not self.layer_dependency else 0)
+            )
+            if self.cnt[1] <= 1:
+                seed2_contr = 0
+            else:
+                seed2_contr = ((self.cnt[0] % self.cnt[1]) / self.cnt[1]).item()
+            frontier, importances = g.sample_labors(
+                seed_nodes,
+                fanout,
+                edge_dir=self.edge_dir,
+                prob=self.prob,
+                importance_sampling=self.importance_sampling,
+                random_seed=random_seed_i,
+                seed2_contribution=seed2_contr,
+                output_device=self.output_device,
+                exclude_edges=exclude_eids,
+            )
+            eid = frontier.edata[EID]
+            block = to_block(
+                frontier, seed_nodes, include_dst_in_src=True, src_nodes=None
+            )
+            block.edata[EID] = eid
+            if len(g.canonical_etypes) > 1:
+                for etype, importance in zip(g.canonical_etypes, importances):
+                    if importance.shape[0] == block.num_edges(etype):
+                        block.edata["edge_weights"][etype] = importance
+            elif importances[0].shape[0] == block.num_edges():
+                block.edata["edge_weights"] = importances[0]
+            seed_nodes = block.srcdata[NID]
+            blocks.insert(0, block)
+
+        self.set_seed()
+        return seed_nodes, output_nodes, blocks

python/dgl/dataloading/negative_sampler.py

@@ -0,0 +1,126 @@
+"""Negative samplers"""
+from collections.abc import Mapping
+
+from .. import backend as F
+
+
+class _BaseNegativeSampler(object):
+    def _generate(self, g, eids, canonical_etype):
+        raise NotImplementedError
+
+    def __call__(self, g, eids):
+        """Returns negative samples.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        eids : Tensor or dict[etype, Tensor]
+            The sampled edges in the minibatch.
+
+        Returns
+        -------
+        tuple[Tensor, Tensor] or dict[etype, tuple[Tensor, Tensor]]
+            The returned source-destination pairs as negative samples.
+        """
+        if isinstance(eids, Mapping):
+            eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
+            neg_pair = {k: self._generate(g, v, k) for k, v in eids.items()}
+        else:
+            assert (
+                len(g.canonical_etypes) == 1
+            ), "please specify a dict of etypes and ids for graphs with multiple edge types"
+            neg_pair = self._generate(g, eids, g.canonical_etypes[0])
+
+        return neg_pair
+
+
+class PerSourceUniform(_BaseNegativeSampler):
+    """Negative sampler that randomly chooses negative destination nodes
+    for each source node according to a uniform distribution.
+
+    For each edge ``(u, v)`` of type ``(srctype, etype, dsttype)``, DGL generates
+    :attr:`k` pairs of negative edges ``(u, v')``, where ``v'`` is chosen
+    uniformly from all the nodes of type ``dsttype``.  The resulting edges will
+    also have type ``(srctype, etype, dsttype)``.
+
+    Parameters
+    ----------
+    k : int
+        The number of negative samples per edge.
+
+    Examples
+    --------
+    >>> g = dgl.graph(([0, 1, 2], [1, 2, 3]))
+    >>> neg_sampler = dgl.dataloading.negative_sampler.PerSourceUniform(2)
+    >>> neg_sampler(g, torch.tensor([0, 1]))
+    (tensor([0, 0, 1, 1]), tensor([1, 0, 2, 3]))
+    """
+
+    def __init__(self, k):
+        self.k = k
+
+    def _generate(self, g, eids, canonical_etype):
+        _, _, vtype = canonical_etype
+        shape = F.shape(eids)
+        dtype = F.dtype(eids)
+        ctx = F.context(eids)
+        shape = (shape[0] * self.k,)
+        src, _ = g.find_edges(eids, etype=canonical_etype)
+        src = F.repeat(src, self.k, 0)
+        dst = F.randint(shape, dtype, ctx, 0, g.num_nodes(vtype))
+        return src, dst
+
+
+# Alias
+Uniform = PerSourceUniform
+
+
+class GlobalUniform(_BaseNegativeSampler):
+    """Negative sampler that randomly chooses negative source-destination pairs according
+    to a uniform distribution.
+
+    For each edge ``(u, v)`` of type ``(srctype, etype, dsttype)``, DGL generates at most
+    :attr:`k` pairs of negative edges ``(u', v')``, where ``u'`` is chosen uniformly from
+    all the nodes of type ``srctype`` and ``v'`` is chosen uniformly from all the nodes
+    of type ``dsttype``.  The resulting edges will also have type
+    ``(srctype, etype, dsttype)``.  DGL guarantees that the sampled pairs will not have
+    edges in between.
+
+    Parameters
+    ----------
+    k : int
+        The desired number of negative samples to generate per edge.
+    exclude_self_loops : bool, optional
+        Whether to exclude self-loops from negative samples.  (Default: True)
+    replace : bool, optional
+        Whether to sample with replacement.  Setting it to True will make things
+        faster.  (Default: False)
+
+    Notes
+    -----
+    This negative sampler will try to generate as many negative samples as possible, but
+    it may rarely return less than :attr:`k` negative samples per edge.
+    This is more likely to happen if a graph is so small or dense that not many unique
+    negative samples exist.
+
+    Examples
+    --------
+    >>> g = dgl.graph(([0, 1, 2], [1, 2, 3]))
+    >>> neg_sampler = dgl.dataloading.negative_sampler.GlobalUniform(2, True)
+    >>> neg_sampler(g, torch.LongTensor([0, 1]))
+    (tensor([0, 1, 3, 2]), tensor([2, 0, 2, 1]))
+    """
+
+    def __init__(self, k, exclude_self_loops=True, replace=False):
+        self.k = k
+        self.exclude_self_loops = exclude_self_loops
+        self.replace = replace
+
+    def _generate(self, g, eids, canonical_etype):
+        return g.global_uniform_negative_sampling(
+            len(eids) * self.k,
+            self.exclude_self_loops,
+            self.replace,
+            canonical_etype,
+        )

python/dgl/dataloading/neighbor_sampler.py

@@ -0,0 +1,246 @@
+"""Data loading components for neighbor sampling"""
+
+from .. import backend as F
+from ..base import EID, NID
+from ..heterograph import DGLGraph
+from ..transforms import to_block
+from ..utils import get_num_threads
+from .base import BlockSampler
+
+
+class NeighborSampler(BlockSampler):
+    """Sampler that builds computational dependency of node representations via
+    neighbor sampling for multilayer GNN.
+
+    This sampler will make every node gather messages from a fixed number of neighbors
+    per edge type.  The neighbors are picked uniformly.
+
+    Parameters
+    ----------
+    fanouts : list[int] or list[dict[etype, int]]
+        List of neighbors to sample per edge type for each GNN layer, with the i-th
+        element being the fanout for the i-th GNN layer.
+
+        If only a single integer is provided, DGL assumes that every edge type
+        will have the same fanout.
+
+        If -1 is provided for one edge type on one layer, then all inbound edges
+        of that edge type will be included.
+    edge_dir : str, default ``'in'``
+        Can be either ``'in' `` where the neighbors will be sampled according to
+        incoming edges, or ``'out'`` otherwise, same as :func:`dgl.sampling.sample_neighbors`.
+    prob : str, optional
+        If given, the probability of each neighbor being sampled is proportional
+        to the edge feature value with the given name in ``g.edata``.  The feature must be
+        a scalar on each edge.
+
+        This argument is mutually exclusive with :attr:`mask`.  If you want to
+        specify both a mask and a probability, consider multiplying the probability
+        with the mask instead.
+    mask : str, optional
+        If given, a neighbor could be picked only if the edge mask with the given
+        name in ``g.edata`` is True.  The data must be boolean on each edge.
+
+        This argument is mutually exclusive with :attr:`prob`.  If you want to
+        specify both a mask and a probability, consider multiplying the probability
+        with the mask instead.
+    replace : bool, default False
+        Whether to sample with replacement
+    prefetch_node_feats : list[str] or dict[ntype, list[str]], optional
+        The source node data to prefetch for the first MFG, corresponding to the
+        input node features necessary for the first GNN layer.
+    prefetch_labels : list[str] or dict[ntype, list[str]], optional
+        The destination node data to prefetch for the last MFG, corresponding to
+        the node labels of the minibatch.
+    prefetch_edge_feats : list[str] or dict[etype, list[str]], optional
+        The edge data names to prefetch for all the MFGs, corresponding to the
+        edge features necessary for all GNN layers.
+    output_device : device, optional
+        The device of the output subgraphs or MFGs.  Default is the same as the
+        minibatch of seed nodes.
+    fused : bool, default True
+        If True and device is CPU fused sample neighbors is invoked. This version
+        requires seed_nodes to be unique
+
+    Examples
+    --------
+    **Node classification**
+
+    To train a 3-layer GNN for node classification on a set of nodes ``train_nid`` on
+    a homogeneous graph where each node takes messages from 5, 10, 15 neighbors for
+    the first, second, and third layer respectively (assuming the backend is PyTorch):
+
+    >>> sampler = dgl.dataloading.NeighborSampler([5, 10, 15])
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, train_nid, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     train_on(blocks)
+
+    If training on a heterogeneous graph and you want different number of neighbors for each
+    edge type, one should instead provide a list of dicts.  Each dict would specify the
+    number of neighbors to pick per edge type.
+
+    >>> sampler = dgl.dataloading.NeighborSampler([
+    ...     {('user', 'follows', 'user'): 5,
+    ...      ('user', 'plays', 'game'): 4,
+    ...      ('game', 'played-by', 'user'): 3}] * 3)
+
+    If you would like non-uniform neighbor sampling:
+
+    >>> g.edata['p'] = torch.rand(g.num_edges())   # any non-negative 1D vector works
+    >>> sampler = dgl.dataloading.NeighborSampler([5, 10, 15], prob='p')
+
+    Or sampling on edge masks:
+
+    >>> g.edata['mask'] = torch.rand(g.num_edges()) < 0.2   # any 1D boolean mask works
+    >>> sampler = dgl.dataloading.NeighborSampler([5, 10, 15], prob='mask')
+
+    **Edge classification and link prediction**
+
+    This class can also work for edge classification and link prediction together
+    with :func:`as_edge_prediction_sampler`.
+
+    >>> sampler = dgl.dataloading.NeighborSampler([5, 10, 15])
+    >>> sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, train_eid, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+
+    See the documentation :func:`as_edge_prediction_sampler` for more details.
+
+    Notes
+    -----
+    For the concept of MFGs, please refer to
+    :ref:`User Guide Section 6 <guide-minibatch>` and
+    :doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
+    """
+
+    def __init__(
+        self,
+        fanouts,
+        edge_dir="in",
+        prob=None,
+        mask=None,
+        replace=False,
+        prefetch_node_feats=None,
+        prefetch_labels=None,
+        prefetch_edge_feats=None,
+        output_device=None,
+        fused=True,
+    ):
+        super().__init__(
+            prefetch_node_feats=prefetch_node_feats,
+            prefetch_labels=prefetch_labels,
+            prefetch_edge_feats=prefetch_edge_feats,
+            output_device=output_device,
+        )
+        self.fanouts = fanouts
+        self.edge_dir = edge_dir
+        if mask is not None and prob is not None:
+            raise ValueError(
+                "Mask and probability arguments are mutually exclusive. "
+                "Consider multiplying the probability with the mask "
+                "to achieve the same goal."
+            )
+        self.prob = prob or mask
+        self.replace = replace
+        self.fused = fused
+        self.mapping = {}
+        self.g = None
+
+    def sample_blocks(self, g, seed_nodes, exclude_eids=None):
+        output_nodes = seed_nodes
+        blocks = []
+        # sample_neighbors_fused function requires multithreading to be more efficient
+        # than sample_neighbors
+        if self.fused and get_num_threads() > 1:
+            cpu = F.device_type(g.device) == "cpu"
+            if isinstance(seed_nodes, dict):
+                for ntype in list(seed_nodes.keys()):
+                    if not cpu:
+                        break
+                    cpu = (
+                        cpu and F.device_type(seed_nodes[ntype].device) == "cpu"
+                    )
+            else:
+                cpu = cpu and F.device_type(seed_nodes.device) == "cpu"
+            if cpu and isinstance(g, DGLGraph) and F.backend_name == "pytorch":
+                if self.g != g:
+                    self.mapping = {}
+                    self.g = g
+                for fanout in reversed(self.fanouts):
+                    block = g.sample_neighbors_fused(
+                        seed_nodes,
+                        fanout,
+                        edge_dir=self.edge_dir,
+                        prob=self.prob,
+                        replace=self.replace,
+                        exclude_edges=exclude_eids,
+                        mapping=self.mapping,
+                    )
+                    seed_nodes = block.srcdata[NID]
+                    blocks.insert(0, block)
+                return seed_nodes, output_nodes, blocks
+
+        for fanout in reversed(self.fanouts):
+            frontier = g.sample_neighbors(
+                seed_nodes,
+                fanout,
+                edge_dir=self.edge_dir,
+                prob=self.prob,
+                replace=self.replace,
+                output_device=self.output_device,
+                exclude_edges=exclude_eids,
+            )
+            block = to_block(frontier, seed_nodes)
+            # If sampled from graphbolt-backed DistGraph, `EID` may not be in
+            # the block. If not exists, we should remove it from the block.
+            if EID in frontier.edata.keys():
+                block.edata[EID] = frontier.edata[EID]
+            else:
+                del block.edata[EID]
+            seed_nodes = block.srcdata[NID]
+            blocks.insert(0, block)
+
+        return seed_nodes, output_nodes, blocks
+
+
+MultiLayerNeighborSampler = NeighborSampler
+
+
+class MultiLayerFullNeighborSampler(NeighborSampler):
+    """Sampler that builds computational dependency of node representations by taking messages
+    from all neighbors for multilayer GNN.
+
+    This sampler will make every node gather messages from every single neighbor per edge type.
+
+    Parameters
+    ----------
+    num_layers : int
+        The number of GNN layers to sample.
+    kwargs :
+        Passed to :class:`dgl.dataloading.NeighborSampler`.
+
+    Examples
+    --------
+    To train a 3-layer GNN for node classification on a set of nodes ``train_nid`` on
+    a homogeneous graph where each node takes messages from all neighbors for the first,
+    second, and third layer respectively (assuming the backend is PyTorch):
+
+    >>> sampler = dgl.dataloading.MultiLayerFullNeighborSampler(3)
+    >>> dataloader = dgl.dataloading.DataLoader(
+    ...     g, train_nid, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     train_on(blocks)
+
+    Notes
+    -----
+    For the concept of MFGs, please refer to
+    :ref:`User Guide Section 6 <guide-minibatch>` and
+    :doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
+    """
+
+    def __init__(self, num_layers, **kwargs):
+        super().__init__([-1] * num_layers, **kwargs)