Skip to content

GitLab

"examples/mxnet/vscode:/vscode.git/clone" did not exist on "1990e797e1f5ef5bc12dbdde6ee03c8066754bf1"
Unverified Commit 76bb5404 authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] Black auto fix. (#4682) 


Co-authored-by: default avatarSteve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
parent a208e886
Hide whitespace changes
Inline Side-by-side
Showing with 696 additions and 340 deletions
python/dgl/function/base.py
View file @ 76bb5404
"""Built-in function base class"""
from __future__ import absolute_import
__all__ = ['BuiltinFunction', 'TargetCode']
__all__ = ["BuiltinFunction", "TargetCode"]
class TargetCode(object):
......@@ -10,6 +10,7 @@ class TargetCode(object):
Note: must be consistent with the target code definition in C++ side:
src/kernel/binary_reduce_common.h
"""
SRC = 0
DST = 1
EDGE = 2
......@@ -23,6 +24,7 @@ class TargetCode(object):
class BuiltinFunction(object):
"""Base builtin function class."""
@property
def name(self):
"""Return the name of this builtin function."""
......
python/dgl/function/reducer.py
View file @ 76bb5404
......@@ -4,16 +4,15 @@ from __future__ import absolute_import
import sys
from .base import BuiltinFunction, TargetCode
from .._deprecate.runtime import ir
from .._deprecate.runtime.ir import var
from .base import BuiltinFunction, TargetCode
class ReduceFunction(BuiltinFunction):
"""Base builtin reduce function class."""
def _invoke(self, graph, edge_frame, out_size, edge_map=None,
out_map=None):
def _invoke(self, graph, edge_frame, out_size, edge_map=None, out_map=None):
"""Symbolic computation of this builtin function to create
runtime.executor
"""
......@@ -28,21 +27,28 @@ class ReduceFunction(BuiltinFunction):
class SimpleReduceFunction(ReduceFunction):
"""Builtin reduce function that aggregates a single field into another
single field."""
def __init__(self, name, msg_field, out_field):
self._name = name
self.msg_field = msg_field
self.out_field = out_field
def _invoke(self, graph, edge_frame, out_size, edge_map=None,
out_map=None):
def _invoke(self, graph, edge_frame, out_size, edge_map=None, out_map=None):
"""Symbolic execution of this builtin function"""
reducer = self._name
graph = var.GRAPH(graph)
edge_map = var.MAP(edge_map)
out_map = var.MAP(out_map)
edge_data = ir.READ_COL(edge_frame, var.STR(self.msg_field))
return ir.COPY_REDUCE(reducer, graph, TargetCode.EDGE, edge_data,
out_size, edge_map, out_map)
return ir.COPY_REDUCE(
reducer,
graph,
TargetCode.EDGE,
edge_data,
out_size,
edge_map,
out_map,
)
@property
def name(self):
......@@ -53,6 +59,7 @@ class SimpleReduceFunction(ReduceFunction):
# Generate all following reducer functions:
# sum, max, min, mean, prod
def _gen_reduce_builtin(reducer):
docstring = """Builtin reduce function that aggregates messages by {0}.
......@@ -73,10 +80,13 @@ def _gen_reduce_builtin(reducer):
>>> import torch
>>> def reduce_func(nodes):
>>> return {{'h': torch.{0}(nodes.mailbox['m'], dim=1)}}
""".format(reducer)
""".format(
reducer
)
def func(msg, out):
return SimpleReduceFunction(reducer, msg, out)
func.__name__ = str(reducer)
func.__qualname__ = str(reducer)
func.__doc__ = docstring
......
python/dgl/generators.py
View file @ 76bb5404
"""Module for various graph generator functions."""
from . import backend as F
from . import convert
from . import random
from . import convert, random
__all__ = ["rand_graph", "rand_bipartite"]
__all__ = ['rand_graph', 'rand_bipartite']
def rand_graph(num_nodes, num_edges, idtype=F.int64, device=F.cpu()):
"""Generate a random graph of the given number of nodes/edges and return.
......@@ -46,20 +46,28 @@ def rand_graph(num_nodes, num_edges, idtype=F.int64, device=F.cpu()):
ndata_schemes={}
edata_schemes={})
"""
#TODO(minjie): support RNG as one of the arguments.
# TODO(minjie): support RNG as one of the arguments.
eids = random.choice(num_nodes * num_nodes, num_edges, replace=False)
eids = F.zerocopy_to_numpy(eids)
rows = F.zerocopy_from_numpy(eids // num_nodes)
cols = F.zerocopy_from_numpy(eids % num_nodes)
rows = F.copy_to(F.astype(rows, idtype), device)
cols = F.copy_to(F.astype(cols, idtype), device)
return convert.graph((rows, cols),
num_nodes=num_nodes,
idtype=idtype, device=device)
def rand_bipartite(utype, etype, vtype,
num_src_nodes, num_dst_nodes, num_edges,
idtype=F.int64, device=F.cpu()):
return convert.graph(
(rows, cols), num_nodes=num_nodes, idtype=idtype, device=device
)
def rand_bipartite(
utype,
etype,
vtype,
num_src_nodes,
num_dst_nodes,
num_edges,
idtype=F.int64,
device=F.cpu(),
):
"""Generate a random uni-directional bipartite graph and return.
It uniformly chooses ``num_edges`` from all possible node pairs and form a graph.
......@@ -107,13 +115,18 @@ def rand_bipartite(utype, etype, vtype,
num_edges={('user', 'buys', 'game'): 10},
metagraph=[('user', 'game', 'buys')])
"""
#TODO(minjie): support RNG as one of the arguments.
eids = random.choice(num_src_nodes * num_dst_nodes, num_edges, replace=False)
# TODO(minjie): support RNG as one of the arguments.
eids = random.choice(
num_src_nodes * num_dst_nodes, num_edges, replace=False
)
eids = F.zerocopy_to_numpy(eids)
rows = F.zerocopy_from_numpy(eids // num_dst_nodes)
cols = F.zerocopy_from_numpy(eids % num_dst_nodes)
rows = F.copy_to(F.astype(rows, idtype), device)
cols = F.copy_to(F.astype(cols, idtype), device)
return convert.heterograph({(utype, etype, vtype): (rows, cols)},
{utype: num_src_nodes, vtype: num_dst_nodes},
idtype=idtype, device=device)
return convert.heterograph(
{(utype, etype, vtype): (rows, cols)},
{utype: num_src_nodes, vtype: num_dst_nodes},
idtype=idtype,
device=device,
)
python/dgl/geometry/__init__.py
View file @ 76bb5404
......@@ -8,5 +8,5 @@
This package is experimental and the interfaces may be subject
to changes in future releases.
"""
from .fps import *
from .edge_coarsening import *
from .fps import *
python/dgl/geometry/capi.py
View file @ 76bb5404
"""Python interfaces to DGL farthest point sampler."""
import numpy as np
from .._ffi.base import DGLError
from .._ffi.function import _init_api
from .. import backend as F
from .. import ndarray as nd
from .._ffi.base import DGLError
from .._ffi.function import _init_api
def _farthest_point_sampler(data, batch_size, sample_points, dist, start_idx, result):
def _farthest_point_sampler(
data, batch_size, sample_points, dist, start_idx, result
):
r"""Farthest Point Sampler
Parameters
......@@ -32,14 +35,19 @@ def _farthest_point_sampler(data, batch_size, sample_points, dist, start_idx, re
assert F.shape(data)[0] >= sample_points * batch_size
assert F.shape(data)[0] % batch_size == 0
_CAPI_FarthestPointSampler(F.zerocopy_to_dgl_ndarray(data),
batch_size, sample_points,
F.zerocopy_to_dgl_ndarray(dist),
F.zerocopy_to_dgl_ndarray(start_idx),
F.zerocopy_to_dgl_ndarray(result))
_CAPI_FarthestPointSampler(
F.zerocopy_to_dgl_ndarray(data),
batch_size,
sample_points,
F.zerocopy_to_dgl_ndarray(dist),
F.zerocopy_to_dgl_ndarray(start_idx),
F.zerocopy_to_dgl_ndarray(result),
)
def _neighbor_matching(graph_idx, num_nodes, edge_weights=None, relabel_idx=True):
def _neighbor_matching(
graph_idx, num_nodes, edge_weights=None, relabel_idx=True
):
"""
Description
-----------
......@@ -82,7 +90,11 @@ def _neighbor_matching(graph_idx, num_nodes, edge_weights=None, relabel_idx=True
if edge_weights is not None:
edge_weight_capi = F.zerocopy_to_dgl_ndarray(edge_weights)
node_label = F.full_1d(
num_nodes, -1, getattr(F, graph_idx.dtype), F.to_backend_ctx(graph_idx.ctx))
num_nodes,
-1,
getattr(F, graph_idx.dtype),
F.to_backend_ctx(graph_idx.ctx),
)
node_label_capi = F.zerocopy_to_dgl_ndarray_for_write(node_label)
_CAPI_NeighborMatching(graph_idx, edge_weight_capi, node_label_capi)
if F.reduce_sum(node_label < 0).item() != 0:
......@@ -99,4 +111,4 @@ def _neighbor_matching(graph_idx, num_nodes, edge_weights=None, relabel_idx=True
return node_label
_init_api('dgl.geometry', __name__)
_init_api("dgl.geometry", __name__)
python/dgl/geometry/edge_coarsening.py
View file @ 76bb5404
......@@ -3,7 +3,8 @@
from .. import remove_self_loop
from .capi import _neighbor_matching
__all__ = ['neighbor_matching']
__all__ = ["neighbor_matching"]
def neighbor_matching(graph, e_weights=None, relabel_idx=True):
r"""
......@@ -48,13 +49,16 @@ def neighbor_matching(graph, e_weights=None, relabel_idx=True):
>>> res = neighbor_matching(g)
tensor([0, 1, 1])
"""
assert graph.is_homogeneous, \
"The graph used in graph node matching must be homogeneous"
assert (
graph.is_homogeneous
), "The graph used in graph node matching must be homogeneous"
if e_weights is not None:
graph.edata['e_weights'] = e_weights
graph.edata["e_weights"] = e_weights
graph = remove_self_loop(graph)
e_weights = graph.edata['e_weights']
graph.edata.pop('e_weights')
e_weights = graph.edata["e_weights"]
graph.edata.pop("e_weights")
else:
graph = remove_self_loop(graph)
return _neighbor_matching(graph._graph, graph.num_nodes(), e_weights, relabel_idx)
return _neighbor_matching(
graph._graph, graph.num_nodes(), e_weights, relabel_idx
)
python/dgl/geometry/fps.py
View file @ 76bb5404
"""Farthest Point Sampler for pytorch Geometry package"""
#pylint: disable=no-member, invalid-name
# pylint: disable=no-member, invalid-name
from .. import backend as F
from ..base import DGLError
from .capi import _farthest_point_sampler
__all__ = ['farthest_point_sampler']
__all__ = ["farthest_point_sampler"]
def farthest_point_sampler(pos, npoints, start_idx=None):
......@@ -50,12 +49,16 @@ def farthest_point_sampler(pos, npoints, start_idx=None):
pos = pos.reshape(-1, C)
dist = F.zeros((B * N), dtype=pos.dtype, ctx=ctx)
if start_idx is None:
start_idx = F.randint(shape=(B, ), dtype=F.int64,
ctx=ctx, low=0, high=N-1)
start_idx = F.randint(
shape=(B,), dtype=F.int64, ctx=ctx, low=0, high=N - 1
)
else:
if start_idx >= N or start_idx < 0:
raise DGLError("Invalid start_idx, expected 0 <= start_idx < {}, got {}".format(
N, start_idx))
raise DGLError(
"Invalid start_idx, expected 0 <= start_idx < {}, got {}".format(
N, start_idx
)
)
start_idx = F.full_1d(B, start_idx, dtype=F.int64, ctx=ctx)
result = F.zeros((npoints * B), dtype=F.int64, ctx=ctx)
_farthest_point_sampler(pos, B, npoints, dist, start_idx, result)
......
python/dgl/global_config.py
View file @ 76bb5404
"""Module for global configuration operators."""
from ._ffi.function import _init_api
__all__ = ["is_libxsmm_enabled", "use_libxsmm"]
def use_libxsmm(flag):
r"""Set whether DGL uses libxsmm at runtime.
......@@ -21,6 +21,7 @@ def use_libxsmm(flag):
"""
_CAPI_DGLConfigSetLibxsmm(flag)
def is_libxsmm_enabled():
r"""Get whether the use_libxsmm flag is turned on.
......@@ -35,4 +36,5 @@ def is_libxsmm_enabled():
"""
return _CAPI_DGLConfigGetLibxsmm()
_init_api("dgl.global_config")
python/dgl/graph_index.py
View file @ 76bb5404
"""Module for graph index class definition."""
from __future__ import absolute_import
import numpy as np
import networkx as nx
import numpy as np
import scipy
from ._ffi.object import register_object, ObjectBase
from ._ffi.function import _init_api
from .base import DGLError, dgl_warning
from . import backend as F
from . import utils
from ._ffi.function import _init_api
from ._ffi.object import ObjectBase, register_object
from .base import DGLError, dgl_warning
class BoolFlag(object):
"""Bool flag with unknown value"""
BOOL_UNKNOWN = -1
BOOL_FALSE = 0
BOOL_TRUE = 1
@register_object('graph.Graph')
@register_object("graph.Graph")
class GraphIndex(ObjectBase):
"""Graph index object.
......@@ -33,6 +36,7 @@ class GraphIndex(ObjectBase):
- `dgl.graph_index.from_csr`
- `dgl.graph_index.from_coo`
"""
def __new__(cls):
obj = ObjectBase.__new__(cls)
obj._readonly = None # python-side cache of the flag
......@@ -53,13 +57,15 @@ class GraphIndex(ObjectBase):
# Pickle compatibility check
# TODO: we should store a storage version number in later releases.
if isinstance(state, tuple) and len(state) == 5:
dgl_warning("The object is pickled pre-0.4.2. Multigraph flag is ignored in 0.4.3")
dgl_warning(
"The object is pickled pre-0.4.2. Multigraph flag is ignored in 0.4.3"
)
num_nodes, _, readonly, src, dst = state
elif isinstance(state, tuple) and len(state) == 4:
# post-0.4.3.
num_nodes, readonly, src, dst = state
else:
raise IOError('Unrecognized storage format.')
raise IOError("Unrecognized storage format.")
self._cache = {}
self._readonly = readonly
......@@ -68,7 +74,8 @@ class GraphIndex(ObjectBase):
src.todgltensor(),
dst.todgltensor(),
int(num_nodes),
readonly)
readonly,
)
def add_nodes(self, num):
"""Add nodes.
......@@ -240,7 +247,9 @@ class GraphIndex(ObjectBase):
"""
u_array = u.todgltensor()
v_array = v.todgltensor()
return utils.toindex(_CAPI_DGLGraphHasEdgesBetween(self, u_array, v_array))
return utils.toindex(
_CAPI_DGLGraphHasEdgesBetween(self, u_array, v_array)
)
def predecessors(self, v, radius=1):
"""Return the predecessors of the node.
......@@ -257,8 +266,9 @@ class GraphIndex(ObjectBase):
utils.Index
Array of predecessors
"""
return utils.toindex(_CAPI_DGLGraphPredecessors(
self, int(v), int(radius)))
return utils.toindex(
_CAPI_DGLGraphPredecessors(self, int(v), int(radius))
)
def successors(self, v, radius=1):
"""Return the successors of the node.
......@@ -275,8 +285,9 @@ class GraphIndex(ObjectBase):
utils.Index
Array of successors
"""
return utils.toindex(_CAPI_DGLGraphSuccessors(
self, int(v), int(radius)))
return utils.toindex(
_CAPI_DGLGraphSuccessors(self, int(v), int(radius))
)
def edge_id(self, u, v):
"""Return the id array of all edges between u and v.
......@@ -432,7 +443,7 @@ class GraphIndex(ObjectBase):
"""
_CAPI_DGLSortAdj(self)
@utils.cached_member(cache='_cache', prefix='edges')
@utils.cached_member(cache="_cache", prefix="edges")
def edges(self, order=None):
"""Return all the edges
......@@ -606,7 +617,7 @@ class GraphIndex(ObjectBase):
e_array = e.todgltensor()
return _CAPI_DGLGraphEdgeSubgraph(self, e_array, preserve_nodes)
@utils.cached_member(cache='_cache', prefix='scipy_adj')
@utils.cached_member(cache="_cache", prefix="scipy_adj")
def adjacency_matrix_scipy(self, transpose, fmt, return_edge_ids=None):
"""Return the scipy adjacency matrix representation of this graph.
......@@ -631,8 +642,10 @@ class GraphIndex(ObjectBase):
The scipy representation of adjacency matrix.
"""
if not isinstance(transpose, bool):
raise DGLError('Expect bool value for "transpose" arg,'
' but got %s.' % (type(transpose)))
raise DGLError(
'Expect bool value for "transpose" arg,'
" but got %s." % (type(transpose))
)
if return_edge_ids is None:
dgl_warning(
......@@ -640,17 +653,24 @@ class GraphIndex(ObjectBase):
" As a result there is one 0 entry which is not eliminated."
" In the next release it will return 1s by default,"
" and 0 will be eliminated otherwise.",
FutureWarning)
FutureWarning,
)
return_edge_ids = True
rst = _CAPI_DGLGraphGetAdj(self, transpose, fmt)
if fmt == "csr":
indptr = utils.toindex(rst(0)).tonumpy()
indices = utils.toindex(rst(1)).tonumpy()
data = utils.toindex(rst(2)).tonumpy() if return_edge_ids else np.ones_like(indices)
data = (
utils.toindex(rst(2)).tonumpy()
if return_edge_ids
else np.ones_like(indices)
)
n = self.number_of_nodes()
return scipy.sparse.csr_matrix((data, indices, indptr), shape=(n, n))
elif fmt == 'coo':
return scipy.sparse.csr_matrix(
(data, indices, indptr), shape=(n, n)
)
elif fmt == "coo":
idx = utils.toindex(rst(0)).tonumpy()
n = self.number_of_nodes()
m = self.number_of_edges()
......@@ -660,7 +680,7 @@ class GraphIndex(ObjectBase):
else:
raise Exception("unknown format")
@utils.cached_member(cache='_cache', prefix='immu_gidx')
@utils.cached_member(cache="_cache", prefix="immu_gidx")
def get_immutable_gidx(self, ctx):
"""Create an immutable graph index and copy to the given device context.
......@@ -717,8 +737,10 @@ class GraphIndex(ObjectBase):
if shuffle is not required.
"""
if not isinstance(transpose, bool):
raise DGLError('Expect bool value for "transpose" arg,'
' but got %s.' % (type(transpose)))
raise DGLError(
'Expect bool value for "transpose" arg,'
" but got %s." % (type(transpose))
)
fmt = F.get_preferred_sparse_format()
rst = _CAPI_DGLGraphGetAdj(self, transpose, fmt)
if fmt == "csr":
......@@ -726,8 +748,11 @@ class GraphIndex(ObjectBase):
indices = F.copy_to(utils.toindex(rst(1)).tousertensor(), ctx)
shuffle = utils.toindex(rst(2))
dat = F.ones(indices.shape, dtype=F.float32, ctx=ctx)
spmat = F.sparse_matrix(dat, ('csr', indices, indptr),
(self.number_of_nodes(), self.number_of_nodes()))[0]
spmat = F.sparse_matrix(
dat,
("csr", indices, indptr),
(self.number_of_nodes(), self.number_of_nodes()),
)[0]
return spmat, shuffle
elif fmt == "coo":
## FIXME(minjie): data type
......@@ -736,8 +761,10 @@ class GraphIndex(ObjectBase):
idx = F.reshape(idx, (2, m))
dat = F.ones((m,), dtype=F.float32, ctx=ctx)
n = self.number_of_nodes()
adj, shuffle_idx = F.sparse_matrix(dat, ('coo', idx), (n, n))
shuffle_idx = utils.toindex(shuffle_idx) if shuffle_idx is not None else None
adj, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, n))
shuffle_idx = (
utils.toindex(shuffle_idx) if shuffle_idx is not None else None
)
return adj, shuffle_idx
else:
raise Exception("unknown format")
......@@ -783,21 +810,21 @@ class GraphIndex(ObjectBase):
eid = eid.tousertensor(ctx) # the index of the ctx will be cached
n = self.number_of_nodes()
m = self.number_of_edges()
if typestr == 'in':
if typestr == "in":
row = F.unsqueeze(dst, 0)
col = F.unsqueeze(eid, 0)
idx = F.cat([row, col], dim=0)
# FIXME(minjie): data type
dat = F.ones((m,), dtype=F.float32, ctx=ctx)
inc, shuffle_idx = F.sparse_matrix(dat, ('coo', idx), (n, m))
elif typestr == 'out':
inc, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, m))
elif typestr == "out":
row = F.unsqueeze(src, 0)
col = F.unsqueeze(eid, 0)
idx = F.cat([row, col], dim=0)
# FIXME(minjie): data type
dat = F.ones((m,), dtype=F.float32, ctx=ctx)
inc, shuffle_idx = F.sparse_matrix(dat, ('coo', idx), (n, m))
elif typestr == 'both':
inc, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, m))
elif typestr == "both":
# first remove entries for self loops
mask = F.logical_not(F.equal(src, dst))
src = F.boolean_mask(src, mask)
......@@ -812,10 +839,12 @@ class GraphIndex(ObjectBase):
x = -F.ones((n_entries,), dtype=F.float32, ctx=ctx)
y = F.ones((n_entries,), dtype=F.float32, ctx=ctx)
dat = F.cat([x, y], dim=0)
inc, shuffle_idx = F.sparse_matrix(dat, ('coo', idx), (n, m))
inc, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, m))
else:
raise DGLError('Invalid incidence matrix type: %s' % str(typestr))
shuffle_idx = utils.toindex(shuffle_idx) if shuffle_idx is not None else None
raise DGLError("Invalid incidence matrix type: %s" % str(typestr))
shuffle_idx = (
utils.toindex(shuffle_idx) if shuffle_idx is not None else None
)
return inc, shuffle_idx
def to_networkx(self):
......@@ -902,7 +931,9 @@ class GraphIndex(ObjectBase):
GraphIndex
The graph index on the given device context.
"""
return _CAPI_DGLImmutableGraphCopyTo(self, ctx.device_type, ctx.device_id)
return _CAPI_DGLImmutableGraphCopyTo(
self, ctx.device_type, ctx.device_id
)
def copyto_shared_mem(self, shared_mem_name):
"""Copy this immutable graph index to shared memory.
......@@ -939,7 +970,10 @@ class GraphIndex(ObjectBase):
int
The number of bits needed
"""
if self.number_of_edges() >= 0x80000000 or self.number_of_nodes() >= 0x80000000:
if (
self.number_of_edges() >= 0x80000000
or self.number_of_nodes() >= 0x80000000
):
return 64
else:
return 32
......@@ -961,9 +995,11 @@ class GraphIndex(ObjectBase):
"""
return _CAPI_DGLImmutableGraphAsNumBits(self, int(bits))
@register_object('graph.Subgraph')
@register_object("graph.Subgraph")
class SubgraphIndex(ObjectBase):
"""Subgraph data structure"""
@property
def graph(self):
"""The subgraph structure
......@@ -1028,16 +1064,15 @@ def from_coo(num_nodes, src, dst, readonly):
dst = utils.toindex(dst)
if readonly:
gidx = _CAPI_DGLGraphCreate(
src.todgltensor(),
dst.todgltensor(),
int(num_nodes),
readonly)
src.todgltensor(), dst.todgltensor(), int(num_nodes), readonly
)
else:
gidx = _CAPI_DGLGraphCreateMutable()
gidx.add_nodes(num_nodes)
gidx.add_edges(src, dst)
return gidx
def from_csr(indptr, indices, direction):
"""Load a graph from CSR arrays.
......@@ -1058,11 +1093,11 @@ def from_csr(indptr, indices, direction):
indptr = utils.toindex(indptr)
indices = utils.toindex(indices)
gidx = _CAPI_DGLGraphCSRCreate(
indptr.todgltensor(),
indices.todgltensor(),
direction)
indptr.todgltensor(), indices.todgltensor(), direction
)
return gidx
def from_shared_mem_graph_index(shared_mem_name):
"""Load a graph index from the shared memory.
......@@ -1078,6 +1113,7 @@ def from_shared_mem_graph_index(shared_mem_name):
"""
return _CAPI_DGLGraphCSRCreateMMap(shared_mem_name)
def from_networkx(nx_graph, readonly):
"""Convert from networkx graph.
......@@ -1107,7 +1143,7 @@ def from_networkx(nx_graph, readonly):
# nx_graph.edges(data=True) returns src, dst, attr_dict
if nx_graph.number_of_edges() > 0:
has_edge_id = 'id' in next(iter(nx_graph.edges(data=True)))[-1]
has_edge_id = "id" in next(iter(nx_graph.edges(data=True)))[-1]
else:
has_edge_id = False
......@@ -1116,7 +1152,7 @@ def from_networkx(nx_graph, readonly):
src = np.zeros((num_edges,), dtype=np.int64)
dst = np.zeros((num_edges,), dtype=np.int64)
for u, v, attr in nx_graph.edges(data=True):
eid = attr['id']
eid = attr["id"]
src[eid] = u
dst[eid] = v
else:
......@@ -1131,6 +1167,7 @@ def from_networkx(nx_graph, readonly):
dst = utils.toindex(dst)
return from_coo(num_nodes, src, dst, readonly)
def from_scipy_sparse_matrix(adj, readonly):
"""Convert from scipy sparse matrix.
......@@ -1145,7 +1182,7 @@ def from_scipy_sparse_matrix(adj, readonly):
GraphIndex
The graph index.
"""
if adj.getformat() != 'csr' or not readonly:
if adj.getformat() != "csr" or not readonly:
num_nodes = max(adj.shape[0], adj.shape[1])
adj_coo = adj.tocoo()
return from_coo(num_nodes, adj_coo.row, adj_coo.col, readonly)
......@@ -1153,6 +1190,7 @@ def from_scipy_sparse_matrix(adj, readonly):
# If the input matrix is csr, we still treat it as multigraph.
return from_csr(adj.indptr, adj.indices, "out")
def from_edge_list(elist, readonly):
"""Convert from an edge list.
......@@ -1172,6 +1210,7 @@ def from_edge_list(elist, readonly):
num_nodes = max(src.max(), dst.max()) + 1
return from_coo(num_nodes, src_ids, dst_ids, readonly)
def map_to_subgraph_nid(induced_nodes, parent_nids):
"""Map parent node Ids to the subgraph node Ids.
......@@ -1188,8 +1227,12 @@ def map_to_subgraph_nid(induced_nodes, parent_nids):
utils.Index
Node Ids in the subgraph.
"""
return utils.toindex(_CAPI_DGLMapSubgraphNID(induced_nodes.todgltensor(),
parent_nids.todgltensor()))
return utils.toindex(
_CAPI_DGLMapSubgraphNID(
induced_nodes.todgltensor(), parent_nids.todgltensor()
)
)
def transform_ids(mapping, ids):
"""Transform ids by the given mapping.
......@@ -1206,8 +1249,10 @@ def transform_ids(mapping, ids):
utils.Index
The new ids.
"""
return utils.toindex(_CAPI_DGLMapSubgraphNID(
mapping.todgltensor(), ids.todgltensor()))
return utils.toindex(
_CAPI_DGLMapSubgraphNID(mapping.todgltensor(), ids.todgltensor())
)
def disjoint_union(graphs):
"""Return a disjoint union of the input graphs.
......@@ -1230,6 +1275,7 @@ def disjoint_union(graphs):
"""
return _CAPI_DGLDisjointUnion(list(graphs))
def disjoint_partition(graph, num_or_size_splits):
"""Partition the graph disjointly.
......@@ -1252,14 +1298,13 @@ def disjoint_partition(graph, num_or_size_splits):
"""
if isinstance(num_or_size_splits, utils.Index):
rst = _CAPI_DGLDisjointPartitionBySizes(
graph,
num_or_size_splits.todgltensor())
graph, num_or_size_splits.todgltensor()
)
else:
rst = _CAPI_DGLDisjointPartitionByNum(
graph,
int(num_or_size_splits))
rst = _CAPI_DGLDisjointPartitionByNum(graph, int(num_or_size_splits))
return rst
def create_graph_index(graph_data, readonly):
"""Create a graph index object.
......@@ -1289,11 +1334,15 @@ def create_graph_index(graph_data, readonly):
try:
gidx = from_networkx(graph_data, readonly)
except Exception: # pylint: disable=broad-except
raise DGLError('Error while creating graph from input of type "%s".'
% type(graph_data))
raise DGLError(
'Error while creating graph from input of type "%s".'
% type(graph_data)
)
return gidx
def _get_halo_subgraph_inner_node(halo_subg):
return _CAPI_GetHaloSubgraphInnerNodes(halo_subg)
_init_api("dgl.graph_index")
python/dgl/heterograph_index.py
View file @ 76bb5404
"""Module for heterogeneous graph index class definition."""
from __future__ import absolute_import
import sys
import itertools
import sys
import numpy as np
import scipy
from ._ffi.object import register_object, ObjectBase
from . import backend as F
from . import utils
from ._ffi.function import _init_api
from ._ffi.object import ObjectBase, register_object
from ._ffi.streams import to_dgl_stream_handle
from .base import DGLError, dgl_warning
from .graph_index import from_coo
from . import backend as F
from . import utils
@register_object('graph.HeteroGraph')
@register_object("graph.HeteroGraph")
class HeteroGraphIndex(ObjectBase):
"""HeteroGraph index object.
......@@ -22,6 +24,7 @@ class HeteroGraphIndex(ObjectBase):
----
Do not create GraphIndex directly.
"""
def __new__(cls):
obj = ObjectBase.__new__(cls)
obj._cache = {}
......@@ -29,8 +32,8 @@ class HeteroGraphIndex(ObjectBase):
def __getstate__(self):
"""Issue: https://github.com/pytorch/pytorch/issues/32351
Need to set the tensor created in the __getstate__ function
as object attribute to avoid potential bugs
Need to set the tensor created in the __getstate__ function
as object attribute to avoid potential bugs
"""
self._pk_state = _CAPI_DGLHeteroPickle(self)
return self._pk_state
......@@ -55,11 +58,19 @@ class HeteroGraphIndex(ObjectBase):
num_src = number_of_nodes[src_ntype]
num_dst = number_of_nodes[dst_ntype]
src_id, dst_id, _ = edges_per_type
rel_graphs.append(create_unitgraph_from_coo(
1 if src_ntype == dst_ntype else 2, num_src, num_dst, src_id, dst_id,
['coo', 'csr', ' csc']))
rel_graphs.append(
create_unitgraph_from_coo(
1 if src_ntype == dst_ntype else 2,
num_src,
num_dst,
src_id,
dst_id,
["coo", "csr", " csc"],
)
)
self.__init_handle_by_constructor__(
_CAPI_DGLHeteroCreateHeteroGraph, metagraph, rel_graphs)
_CAPI_DGLHeteroCreateHeteroGraph, metagraph, rel_graphs
)
@property
def metagraph(self):
......@@ -155,7 +166,9 @@ class HeteroGraphIndex(ObjectBase):
v : utils.Index
The dst nodes.
"""
_CAPI_DGLHeteroAddEdges(self, int(etype), u.todgltensor(), v.todgltensor())
_CAPI_DGLHeteroAddEdges(
self, int(etype), u.todgltensor(), v.todgltensor()
)
self.clear_cache()
def clear(self):
......@@ -199,9 +212,11 @@ class HeteroGraphIndex(ObjectBase):
The number of bits needed.
"""
stype, dtype = self.metagraph.find_edge(etype)
if (self.number_of_edges(etype) >= 0x80000000 or
self.number_of_nodes(stype) >= 0x80000000 or
self.number_of_nodes(dtype) >= 0x80000000):
if (
self.number_of_edges(etype) >= 0x80000000
or self.number_of_nodes(stype) >= 0x80000000
or self.number_of_nodes(dtype) >= 0x80000000
):
return 64
else:
return 32
......@@ -293,7 +308,9 @@ class HeteroGraphIndex(ObjectBase):
"""
return _CAPI_DGLHeteroRecordStream(self, to_dgl_stream_handle(stream))
def shared_memory(self, name, ntypes=None, etypes=None, formats=('coo', 'csr', 'csc')):
def shared_memory(
self, name, ntypes=None, etypes=None, formats=("coo", "csr", "csc")
):
"""Return a copy of this graph in shared memory
Parameters
......@@ -318,7 +335,9 @@ class HeteroGraphIndex(ObjectBase):
assert fmt in ("coo", "csr", "csc")
ntypes = [] if ntypes is None else ntypes
etypes = [] if etypes is None else etypes
return _CAPI_DGLHeteroCopyToSharedMem(self, name, ntypes, etypes, formats)
return _CAPI_DGLHeteroCopyToSharedMem(
self, name, ntypes, etypes, formats
)
def is_multigraph(self):
"""Return whether the graph is a multigraph
......@@ -386,8 +405,9 @@ class HeteroGraphIndex(ObjectBase):
Tensor
0-1 array indicating existence
"""
return F.from_dgl_nd(_CAPI_DGLHeteroHasVertices(
self, int(ntype), F.to_dgl_nd(vids)))
return F.from_dgl_nd(
_CAPI_DGLHeteroHasVertices(self, int(ntype), F.to_dgl_nd(vids))
)
def has_edges_between(self, etype, u, v):
"""Return true if the edge exists.
......@@ -406,8 +426,11 @@ class HeteroGraphIndex(ObjectBase):
Tensor
0-1 array indicating existence
"""
return F.from_dgl_nd(_CAPI_DGLHeteroHasEdgesBetween(
self, int(etype), F.to_dgl_nd(u), F.to_dgl_nd(v)))
return F.from_dgl_nd(
_CAPI_DGLHeteroHasEdgesBetween(
self, int(etype), F.to_dgl_nd(u), F.to_dgl_nd(v)
)
)
def predecessors(self, etype, v):
"""Return the predecessors of the node.
......@@ -426,8 +449,9 @@ class HeteroGraphIndex(ObjectBase):
Tensor
Array of predecessors
"""
return F.from_dgl_nd(_CAPI_DGLHeteroPredecessors(
self, int(etype), int(v)))
return F.from_dgl_nd(
_CAPI_DGLHeteroPredecessors(self, int(etype), int(v))
)
def successors(self, etype, v):
"""Return the successors of the node.
......@@ -446,8 +470,9 @@ class HeteroGraphIndex(ObjectBase):
Tensor
Array of successors
"""
return F.from_dgl_nd(_CAPI_DGLHeteroSuccessors(
self, int(etype), int(v)))
return F.from_dgl_nd(
_CAPI_DGLHeteroSuccessors(self, int(etype), int(v))
)
def edge_ids_all(self, etype, u, v):
"""Return a triplet of arrays that contains the edge IDs.
......@@ -471,7 +496,8 @@ class HeteroGraphIndex(ObjectBase):
The edge ids.
"""
edge_array = _CAPI_DGLHeteroEdgeIdsAll(
self, int(etype), F.to_dgl_nd(u), F.to_dgl_nd(v))
self, int(etype), F.to_dgl_nd(u), F.to_dgl_nd(v)
)
src = F.from_dgl_nd(edge_array(0))
dst = F.from_dgl_nd(edge_array(1))
......@@ -496,8 +522,11 @@ class HeteroGraphIndex(ObjectBase):
Tensor
The edge ids.
"""
eid = F.from_dgl_nd(_CAPI_DGLHeteroEdgeIdsOne(
self, int(etype), F.to_dgl_nd(u), F.to_dgl_nd(v)))
eid = F.from_dgl_nd(
_CAPI_DGLHeteroEdgeIdsOne(
self, int(etype), F.to_dgl_nd(u), F.to_dgl_nd(v)
)
)
return eid
def find_edges(self, etype, eid):
......@@ -520,7 +549,8 @@ class HeteroGraphIndex(ObjectBase):
The edge ids.
"""
edge_array = _CAPI_DGLHeteroFindEdges(
self, int(etype), F.to_dgl_nd(eid))
self, int(etype), F.to_dgl_nd(eid)
)
src = F.from_dgl_nd(edge_array(0))
dst = F.from_dgl_nd(edge_array(1))
......@@ -607,9 +637,11 @@ class HeteroGraphIndex(ObjectBase):
"""
if order is None:
order = ""
elif order not in ['srcdst', 'eid']:
raise DGLError("Expect order to be one of None, 'srcdst', 'eid', "
"got {}".format(order))
elif order not in ["srcdst", "eid"]:
raise DGLError(
"Expect order to be one of None, 'srcdst', 'eid', "
"got {}".format(order)
)
edge_array = _CAPI_DGLHeteroEdges(self, int(etype), order)
src = F.from_dgl_nd(edge_array(0))
dst = F.from_dgl_nd(edge_array(1))
......@@ -633,8 +665,9 @@ class HeteroGraphIndex(ObjectBase):
Tensor
The in degree array.
"""
return F.from_dgl_nd(_CAPI_DGLHeteroInDegrees(
self, int(etype), F.to_dgl_nd(v)))
return F.from_dgl_nd(
_CAPI_DGLHeteroInDegrees(self, int(etype), F.to_dgl_nd(v))
)
def out_degrees(self, etype, v):
"""Return the out degrees of the nodes.
......@@ -653,8 +686,9 @@ class HeteroGraphIndex(ObjectBase):
Tensor
The out degree array.
"""
return F.from_dgl_nd(_CAPI_DGLHeteroOutDegrees(
self, int(etype), F.to_dgl_nd(v)))
return F.from_dgl_nd(
_CAPI_DGLHeteroOutDegrees(self, int(etype), F.to_dgl_nd(v))
)
def adjacency_matrix(self, etype, transpose, ctx):
"""Return the adjacency matrix representation of this graph.
......@@ -683,28 +717,43 @@ class HeteroGraphIndex(ObjectBase):
if shuffle is not required.
"""
if not isinstance(transpose, bool):
raise DGLError('Expect bool value for "transpose" arg,'
' but got %s.' % (type(transpose)))
raise DGLError(
'Expect bool value for "transpose" arg,'
" but got %s." % (type(transpose))
)
fmt = F.get_preferred_sparse_format()
rst = _CAPI_DGLHeteroGetAdj(self, int(etype), transpose, fmt)
# convert to framework-specific sparse matrix
srctype, dsttype = self.metagraph.find_edge(etype)
nrows = self.number_of_nodes(dsttype) if transpose else self.number_of_nodes(srctype)
ncols = self.number_of_nodes(srctype) if transpose else self.number_of_nodes(dsttype)
nrows = (
self.number_of_nodes(dsttype)
if transpose
else self.number_of_nodes(srctype)
)
ncols = (
self.number_of_nodes(srctype)
if transpose
else self.number_of_nodes(dsttype)
)
nnz = self.number_of_edges(etype)
if fmt == "csr":
indptr = F.copy_to(F.from_dgl_nd(rst(0)), ctx)
indices = F.copy_to(F.from_dgl_nd(rst(1)), ctx)
shuffle = F.copy_to(F.from_dgl_nd(rst(2)), ctx)
dat = F.ones(nnz, dtype=F.float32, ctx=ctx) # FIXME(minjie): data type
spmat = F.sparse_matrix(dat, ('csr', indices, indptr), (nrows, ncols))[0]
dat = F.ones(
nnz, dtype=F.float32, ctx=ctx
) # FIXME(minjie): data type
spmat = F.sparse_matrix(
dat, ("csr", indices, indptr), (nrows, ncols)
)[0]
return spmat, shuffle
elif fmt == "coo":
idx = F.copy_to(F.from_dgl_nd(rst(0)), ctx)
idx = F.reshape(idx, (2, nnz))
dat = F.ones((nnz,), dtype=F.float32, ctx=ctx)
adj, shuffle_idx = F.sparse_matrix(
dat, ('coo', idx), (nrows, ncols))
dat, ("coo", idx), (nrows, ncols)
)
return adj, shuffle_idx
else:
raise Exception("unknown format")
......@@ -743,27 +792,39 @@ class HeteroGraphIndex(ObjectBase):
equivalent to a consecutive array from zero to the number of edges minus one.
"""
if not isinstance(transpose, bool):
raise DGLError('Expect bool value for "transpose" arg,'
' but got %s.' % (type(transpose)))
raise DGLError(
'Expect bool value for "transpose" arg,'
" but got %s." % (type(transpose))
)
rst = _CAPI_DGLHeteroGetAdj(self, int(etype), transpose, fmt)
srctype, dsttype = self.metagraph.find_edge(etype)
nrows = self.number_of_nodes(dsttype) if transpose else self.number_of_nodes(srctype)
ncols = self.number_of_nodes(srctype) if transpose else self.number_of_nodes(dsttype)
nrows = (
self.number_of_nodes(dsttype)
if transpose
else self.number_of_nodes(srctype)
)
ncols = (
self.number_of_nodes(srctype)
if transpose
else self.number_of_nodes(dsttype)
)
nnz = self.number_of_edges(etype)
if fmt == "csr":
indptr = F.from_dgl_nd(rst(0))
indices = F.from_dgl_nd(rst(1))
data = F.from_dgl_nd(rst(2))
return nrows, ncols, indptr, indices, data
elif fmt == 'coo':
elif fmt == "coo":
idx = F.from_dgl_nd(rst(0))
row, col = F.reshape(idx, (2, nnz))
return nrows, ncols, row, col
else:
raise ValueError("unknown format")
def adjacency_matrix_scipy(self, etype, transpose, fmt, return_edge_ids=None):
def adjacency_matrix_scipy(
self, etype, transpose, fmt, return_edge_ids=None
):
"""Return the scipy adjacency matrix representation of this graph.
By default, a row of returned adjacency matrix represents the destination
......@@ -794,12 +855,14 @@ class HeteroGraphIndex(ObjectBase):
" As a result there is one 0 entry which is not eliminated."
" In the next release it will return 1s by default,"
" and 0 will be eliminated otherwise.",
FutureWarning)
FutureWarning,
)
return_edge_ids = True
if fmt == 'csr':
nrows, ncols, indptr, indices, data = \
self.adjacency_matrix_tensors(etype, transpose, fmt)
if fmt == "csr":
nrows, ncols, indptr, indices, data = self.adjacency_matrix_tensors(
etype, transpose, fmt
)
indptr = F.asnumpy(indptr)
indices = F.asnumpy(indices)
data = F.asnumpy(data)
......@@ -810,15 +873,23 @@ class HeteroGraphIndex(ObjectBase):
else:
data = np.ones_like(indices)
return scipy.sparse.csr_matrix((data, indices, indptr), shape=(nrows, ncols))
elif fmt == 'coo':
nrows, ncols, row, col = \
self.adjacency_matrix_tensors(etype, transpose, fmt)
return scipy.sparse.csr_matrix(
(data, indices, indptr), shape=(nrows, ncols)
)
elif fmt == "coo":
nrows, ncols, row, col = self.adjacency_matrix_tensors(
etype, transpose, fmt
)
row = F.asnumpy(row)
col = F.asnumpy(col)
data = np.arange(self.number_of_edges(etype)) if return_edge_ids \
else np.ones_like(row)
return scipy.sparse.coo_matrix((data, (row, col)), shape=(nrows, ncols))
data = (
np.arange(self.number_of_edges(etype))
if return_edge_ids
else np.ones_like(row)
)
return scipy.sparse.coo_matrix(
(data, (row, col)), shape=(nrows, ncols)
)
else:
raise ValueError("unknown format")
......@@ -863,25 +934,26 @@ class HeteroGraphIndex(ObjectBase):
srctype, dsttype = self.metagraph.find_edge(etype)
m = self.number_of_edges(etype)
if typestr == 'in':
if typestr == "in":
n = self.number_of_nodes(dsttype)
row = F.unsqueeze(dst, 0)
col = F.unsqueeze(eid, 0)
idx = F.copy_to(F.cat([row, col], dim=0), ctx)
# FIXME(minjie): data type
dat = F.ones((m,), dtype=F.float32, ctx=ctx)
inc, shuffle_idx = F.sparse_matrix(dat, ('coo', idx), (n, m))
elif typestr == 'out':
inc, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, m))
elif typestr == "out":
n = self.number_of_nodes(srctype)
row = F.unsqueeze(src, 0)
col = F.unsqueeze(eid, 0)
idx = F.copy_to(F.cat([row, col], dim=0), ctx)
# FIXME(minjie): data type
dat = F.ones((m,), dtype=F.float32, ctx=ctx)
inc, shuffle_idx = F.sparse_matrix(dat, ('coo', idx), (n, m))
elif typestr == 'both':
assert srctype == dsttype, \
"'both' is supported only if source and destination type are the same"
inc, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, m))
elif typestr == "both":
assert (
srctype == dsttype
), "'both' is supported only if source and destination type are the same"
n = self.number_of_nodes(srctype)
# first remove entries for self loops
mask = F.logical_not(F.equal(src, dst))
......@@ -897,9 +969,9 @@ class HeteroGraphIndex(ObjectBase):
x = -F.ones((n_entries,), dtype=F.float32, ctx=ctx)
y = F.ones((n_entries,), dtype=F.float32, ctx=ctx)
dat = F.cat([x, y], dim=0)
inc, shuffle_idx = F.sparse_matrix(dat, ('coo', idx), (n, m))
inc, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, m))
else:
raise DGLError('Invalid incidence matrix type: %s' % str(typestr))
raise DGLError("Invalid incidence matrix type: %s" % str(typestr))
return inc, shuffle_idx
def node_subgraph(self, induced_nodes, relabel_nodes):
......@@ -982,7 +1054,9 @@ class HeteroGraphIndex(ObjectBase):
order = csr(2)
rev_csr = _CAPI_DGLHeteroGetAdj(self, int(etype), True, "csr")
rev_order = rev_csr(2)
return utils.toindex(order, self.dtype), utils.toindex(rev_order, self.dtype)
return utils.toindex(order, self.dtype), utils.toindex(
rev_order, self.dtype
)
def formats(self, formats=None):
"""Get a graph index with the specified sparse format(s) or query
......@@ -1014,16 +1088,13 @@ class HeteroGraphIndex(ObjectBase):
created = []
not_created = []
if formats is None:
for fmt in ['coo', 'csr', 'csc']:
for fmt in ["coo", "csr", "csc"]:
if fmt in formats_allowed:
if fmt in formats_created:
created.append(fmt)
else:
not_created.append(fmt)
return {
'created': created,
'not created': not_created
}
return {"created": created, "not created": not_created}
else:
if isinstance(formats, str):
formats = [formats]
......@@ -1044,9 +1115,11 @@ class HeteroGraphIndex(ObjectBase):
"""
return _CAPI_DGLHeteroReverse(self)
@register_object('graph.HeteroSubgraph')
@register_object("graph.HeteroSubgraph")
class HeteroSubgraphIndex(ObjectBase):
"""Hetero-subgraph data structure"""
@property
def graph(self):
"""The subgraph structure
......@@ -1089,6 +1162,7 @@ class HeteroSubgraphIndex(ObjectBase):
# Creators
#################################################################
def create_metagraph_index(ntypes, canonical_etypes):
"""Return a GraphIndex instance for a metagraph given the node types and canonical
edge types.
......@@ -1129,8 +1203,17 @@ def create_metagraph_index(ntypes, canonical_etypes):
metagraph = from_coo(len(ntypes), meta_edges_src, meta_edges_dst, True)
return metagraph, ntypes, etypes, relations
def create_unitgraph_from_coo(num_ntypes, num_src, num_dst, row, col,
formats, row_sorted=False, col_sorted=False):
def create_unitgraph_from_coo(
num_ntypes,
num_src,
num_dst,
row,
col,
formats,
row_sorted=False,
col_sorted=False,
):
"""Create a unitgraph graph index from COO format
Parameters
......@@ -1160,12 +1243,27 @@ def create_unitgraph_from_coo(num_ntypes, num_src, num_dst, row, col,
if isinstance(formats, str):
formats = [formats]
return _CAPI_DGLHeteroCreateUnitGraphFromCOO(
int(num_ntypes), int(num_src), int(num_dst),
F.to_dgl_nd(row), F.to_dgl_nd(col),
formats, row_sorted, col_sorted)
def create_unitgraph_from_csr(num_ntypes, num_src, num_dst, indptr, indices, edge_ids,
formats, transpose=False):
int(num_ntypes),
int(num_src),
int(num_dst),
F.to_dgl_nd(row),
F.to_dgl_nd(col),
formats,
row_sorted,
col_sorted,
)
def create_unitgraph_from_csr(
num_ntypes,
num_src,
num_dst,
indptr,
indices,
edge_ids,
formats,
transpose=False,
):
"""Create a unitgraph graph index from CSR format
Parameters
......@@ -1194,11 +1292,20 @@ def create_unitgraph_from_csr(num_ntypes, num_src, num_dst, indptr, indices, edg
if isinstance(formats, str):
formats = [formats]
return _CAPI_DGLHeteroCreateUnitGraphFromCSR(
int(num_ntypes), int(num_src), int(num_dst),
F.to_dgl_nd(indptr), F.to_dgl_nd(indices), F.to_dgl_nd(edge_ids),
formats, transpose)
def create_heterograph_from_relations(metagraph, rel_graphs, num_nodes_per_type):
int(num_ntypes),
int(num_src),
int(num_dst),
F.to_dgl_nd(indptr),
F.to_dgl_nd(indices),
F.to_dgl_nd(edge_ids),
formats,
transpose,
)
def create_heterograph_from_relations(
metagraph, rel_graphs, num_nodes_per_type
):
"""Create a heterograph from metagraph and graphs of every relation.
Parameters
......@@ -1218,7 +1325,9 @@ def create_heterograph_from_relations(metagraph, rel_graphs, num_nodes_per_type)
return _CAPI_DGLHeteroCreateHeteroGraph(metagraph, rel_graphs)
else:
return _CAPI_DGLHeteroCreateHeteroGraphWithNumNodes(
metagraph, rel_graphs, num_nodes_per_type.todgltensor())
metagraph, rel_graphs, num_nodes_per_type.todgltensor()
)
def create_heterograph_from_shared_memory(name):
"""Create a heterograph from shared memory with the given name.
......@@ -1239,6 +1348,7 @@ def create_heterograph_from_shared_memory(name):
g, ntypes, etypes = _CAPI_DGLHeteroCreateFromSharedMem(name)
return g, list(ntypes), list(etypes)
def joint_union(metagraph, gidx_list):
"""Return a joint union of the input heterographs.
......@@ -1256,6 +1366,7 @@ def joint_union(metagraph, gidx_list):
"""
return _CAPI_DGLHeteroJointUnion(metagraph, gidx_list)
def disjoint_union(metagraph, graphs):
"""Return a disjoint union of the input heterographs.
......@@ -1273,6 +1384,7 @@ def disjoint_union(metagraph, graphs):
"""
return _CAPI_DGLHeteroDisjointUnion_v2(metagraph, graphs)
def disjoint_partition(graph, bnn_all_types, bne_all_types):
"""Partition the graph disjointly.
......@@ -1290,10 +1402,16 @@ def disjoint_partition(graph, bnn_all_types, bne_all_types):
list of HeteroGraphIndex
Heterographs unbatched.
"""
bnn_all_types = utils.toindex(list(itertools.chain.from_iterable(bnn_all_types)))
bne_all_types = utils.toindex(list(itertools.chain.from_iterable(bne_all_types)))
bnn_all_types = utils.toindex(
list(itertools.chain.from_iterable(bnn_all_types))
)
bne_all_types = utils.toindex(
list(itertools.chain.from_iterable(bne_all_types))
)
return _CAPI_DGLHeteroDisjointPartitionBySizes_v2(
graph, bnn_all_types.todgltensor(), bne_all_types.todgltensor())
graph, bnn_all_types.todgltensor(), bne_all_types.todgltensor()
)
def slice_gidx(graph, num_nodes, start_nid, num_edges, start_eid):
"""Slice a chunk of the graph.
......@@ -1317,20 +1435,28 @@ def slice_gidx(graph, num_nodes, start_nid, num_edges, start_eid):
The sliced graph.
"""
return _CAPI_DGLHeteroSlice(
graph, num_nodes.todgltensor(), start_nid.todgltensor(),
num_edges.todgltensor(), start_eid.todgltensor())
graph,
num_nodes.todgltensor(),
start_nid.todgltensor(),
num_edges.todgltensor(),
start_eid.todgltensor(),
)
#################################################################
# Data structure used by C APIs
#################################################################
@register_object("graph.FlattenedHeteroGraph")
class FlattenedHeteroGraph(ObjectBase):
"""FlattenedHeteroGraph object class in C++ backend."""
@register_object("graph.HeteroPickleStates")
class HeteroPickleStates(ObjectBase):
"""Pickle states object class in C++ backend."""
@property
def version(self):
"""Version number
......@@ -1368,10 +1494,12 @@ class HeteroPickleStates(ObjectBase):
def __getstate__(self):
"""Issue: https://github.com/pytorch/pytorch/issues/32351
Need to set the tensor created in the __getstate__ function
as object attribute to avoid potential bugs
Need to set the tensor created in the __getstate__ function
as object attribute to avoid potential bugs
"""
self._pk_arrays = [F.zerocopy_from_dgl_ndarray(arr) for arr in self.arrays]
self._pk_arrays = [
F.zerocopy_from_dgl_ndarray(arr) for arr in self.arrays
]
return self.version, self.meta, self._pk_arrays
def __setstate__(self, state):
......@@ -1379,12 +1507,18 @@ class HeteroPickleStates(ObjectBase):
version, meta, arrays = state
arrays = [F.zerocopy_to_dgl_ndarray(arr) for arr in arrays]
self.__init_handle_by_constructor__(
_CAPI_DGLCreateHeteroPickleStates, version, meta, arrays)
_CAPI_DGLCreateHeteroPickleStates, version, meta, arrays
)
else:
metagraph, num_nodes_per_type, adjs = state
num_nodes_per_type = F.zerocopy_to_dgl_ndarray(num_nodes_per_type)
self.__init_handle_by_constructor__(
_CAPI_DGLCreateHeteroPickleStatesOld, metagraph, num_nodes_per_type, adjs)
_CAPI_DGLCreateHeteroPickleStatesOld,
metagraph,
num_nodes_per_type,
adjs,
)
def _forking_rebuild(pk_state):
version, meta, arrays = pk_state
......@@ -1394,6 +1528,7 @@ def _forking_rebuild(pk_state):
graph_index._forking_pk_state = pk_state
return graph_index
def _forking_reduce(graph_index):
# Because F.from_dgl_nd(F.to_dgl_nd(x)) loses the information of shared memory
# file descriptor (because DLPack does not keep it), without caching the tensors
......@@ -1403,7 +1538,7 @@ def _forking_reduce(graph_index):
# should be rare though because (1) DataLoader will create all the formats if num_workers > 0
# anyway, and (2) we require the users to explicitly create all formats before calling
# mp.spawn().
if hasattr(graph_index, '_forking_pk_state'):
if hasattr(graph_index, "_forking_pk_state"):
return _forking_rebuild, (graph_index._forking_pk_state,)
states = _CAPI_DGLHeteroForkingPickle(graph_index)
arrays = [F.from_dgl_nd(arr) for arr in states.arrays]
......@@ -1415,10 +1550,11 @@ def _forking_reduce(graph_index):
return _forking_rebuild, (graph_index._forking_pk_state,)
if not (F.get_preferred_backend() == 'mxnet' and sys.version_info.minor <= 6):
if not (F.get_preferred_backend() == "mxnet" and sys.version_info.minor <= 6):
# Python 3.6 MXNet crashes with the following statement; remove until we no longer support
# 3.6 (which is EOL anyway).
from multiprocessing.reduction import ForkingPickler
ForkingPickler.register(HeteroGraphIndex, _forking_reduce)
_init_api("dgl.heterograph_index")
python/dgl/init.py
View file @ 76bb5404
......@@ -3,9 +3,12 @@ from __future__ import absolute_import
from . import backend as F
__all__ = ['base_initializer', 'zero_initializer']
__all__ = ["base_initializer", "zero_initializer"]
def base_initializer(shape, dtype, ctx, id_range): # pylint: disable=unused-argument
def base_initializer(
shape, dtype, ctx, id_range
): # pylint: disable=unused-argument
"""The function signature for feature initializer.
Any customized feature initializer should follow this signature (see
......@@ -44,7 +47,10 @@ def base_initializer(shape, dtype, ctx, id_range): # pylint: disable=unused-arg
"""
raise NotImplementedError
def zero_initializer(shape, dtype, ctx, id_range): # pylint: disable=unused-argument
def zero_initializer(
shape, dtype, ctx, id_range
): # pylint: disable=unused-argument
"""Zero feature initializer
Examples
......
python/dgl/merge.py
View file @ 76bb5404
"""Utilities for merging graphs."""
import dgl
from . import backend as F
from .base import DGLError
__all__ = ['merge']
__all__ = ["merge"]
def merge(graphs):
r"""Merge a sequence of graphs together into a single graph.
......@@ -62,7 +64,7 @@ def merge(graphs):
"""
if len(graphs) == 0:
raise DGLError('The input list of graphs cannot be empty.')
raise DGLError("The input list of graphs cannot be empty.")
ref = graphs[0]
ntypes = ref.ntypes
......@@ -87,9 +89,15 @@ def merge(graphs):
if len(keys) == 0:
edges_data = None
else:
edges_data = {k: F.cat([f[k] for f in edata_frames], dim=0) for k in keys}
merged_us = F.copy_to(F.astype(F.cat(unmerged_us, dim=0), ref.idtype), ref.device)
merged_vs = F.copy_to(F.astype(F.cat(unmerged_vs, dim=0), ref.idtype), ref.device)
edges_data = {
k: F.cat([f[k] for f in edata_frames], dim=0) for k in keys
}
merged_us = F.copy_to(
F.astype(F.cat(unmerged_us, dim=0), ref.idtype), ref.device
)
merged_vs = F.copy_to(
F.astype(F.cat(unmerged_vs, dim=0), ref.idtype), ref.device
)
merged.add_edges(merged_us, merged_vs, edges_data, etype)
# Add node data and isolated nodes from next_graph to merged.
......@@ -98,12 +106,16 @@ def merge(graphs):
merged_ntype_id = merged.get_ntype_id(ntype)
next_ntype_id = next_graph.get_ntype_id(ntype)
next_ndata = next_graph._node_frames[next_ntype_id]
node_diff = (next_graph.num_nodes(ntype=ntype) -
merged.num_nodes(ntype=ntype))
node_diff = next_graph.num_nodes(ntype=ntype) - merged.num_nodes(
ntype=ntype
)
n_extra_nodes = max(0, node_diff)
merged.add_nodes(n_extra_nodes, ntype=ntype)
next_nodes = F.arange(
0, next_graph.num_nodes(ntype=ntype), merged.idtype, merged.device
0,
next_graph.num_nodes(ntype=ntype),
merged.idtype,
merged.device,
)
merged._node_frames[merged_ntype_id].update_row(
next_nodes, next_ndata
......
python/dgl/mock_sparse/__init__.py
View file @ 76bb5404
"""dgl sparse class."""
from .diag_matrix import *
from .sp_matrix import *
from .elementwise_op import *
from .matmul import *
from .reduction import * # pylint: disable=W0622
from .sddmm import *
from .reduction import * # pylint: disable=W0622
from .sp_matrix import *
from .unary_diag import *
from .unary_sp import *
from .matmul import *
python/dgl/mock_sparse/sp_matrix.py
View file @ 76bb5404
"""DGL sparse matrix module."""
from typing import Optional, Tuple
import torch
__all__ = [
......
python/dgl/multiprocessing/__init__.py
View file @ 76bb5404
......@@ -6,11 +6,12 @@
# make fork() and openmp work together.
from .. import backend as F
if F.get_preferred_backend() == 'pytorch':
if F.get_preferred_backend() == "pytorch":
# Wrap around torch.multiprocessing...
from torch.multiprocessing import *
# ... and override the Process initializer.
from .pytorch import *
else:
# Just import multiprocessing module.
from multiprocessing import * # pylint: disable=redefined-builtin
from multiprocessing import * # pylint: disable=redefined-builtin
python/dgl/multiprocessing/pytorch.py
View file @ 76bb5404
"""PyTorch multiprocessing wrapper."""
from functools import wraps
import random
import traceback
from _thread import start_new_thread
from functools import wraps
import torch
import torch.multiprocessing as mp
from ..utils import create_shared_mem_array, get_shared_mem_array
def thread_wrapped_func(func):
"""
Wraps a process entry point to make it work with OpenMP.
"""
@wraps(func)
def decorated_function(*args, **kwargs):
queue = mp.Queue()
def _queue_result():
exception, trace, res = None, None, None
try:
......@@ -31,18 +35,31 @@ def thread_wrapped_func(func):
else:
assert isinstance(exception, Exception)
raise exception.__class__(trace)
return decorated_function
# pylint: disable=missing-docstring
class Process(mp.Process):
# pylint: disable=dangerous-default-value
def __init__(self, group=None, target=None, name=None, args=(), kwargs={}, *, daemon=None):
def __init__(
self,
group=None,
target=None,
name=None,
args=(),
kwargs={},
*,
daemon=None
):
target = thread_wrapped_func(target)
super().__init__(group, target, name, args, kwargs, daemon=daemon)
def _get_shared_mem_name(id_):
return "shared" + str(id_)
def call_once_and_share(func, shape, dtype, rank=0):
"""Invoke the function in a single process of the PyTorch distributed process group,
and share the result with other processes.
......@@ -61,7 +78,7 @@ def call_once_and_share(func, shape, dtype, rank=0):
current_rank = torch.distributed.get_rank()
dist_buf = torch.LongTensor([1])
if torch.distributed.get_backend() == 'nccl':
if torch.distributed.get_backend() == "nccl":
# Use .cuda() to transfer it to the correct device. Should be OK since
# PyTorch recommends the users to call set_device() after getting inside
# torch.multiprocessing.spawn()
......@@ -88,6 +105,7 @@ def call_once_and_share(func, shape, dtype, rank=0):
return result
def shared_tensor(shape, dtype=torch.float32):
"""Create a tensor in shared memory accessible by all processes within the same
``torch.distributed`` process group.
......@@ -106,4 +124,6 @@ def shared_tensor(shape, dtype=torch.float32):
Tensor
The shared tensor.
"""
return call_once_and_share(lambda: torch.empty(*shape, dtype=dtype), shape, dtype)
return call_once_and_share(
lambda: torch.empty(*shape, dtype=dtype), shape, dtype
)
python/dgl/ndarray.py
View file @ 76bb5404
......@@ -9,17 +9,28 @@ from __future__ import absolute_import as _abs
import ctypes
import functools
import operator
import numpy as _np
from ._ffi.object import register_object, ObjectBase
from ._ffi.function import _init_api
from ._ffi.ndarray import DGLContext, DGLDataType, NDArrayBase
from ._ffi.ndarray import context, empty, empty_shared_mem, from_dlpack, numpyasarray
from ._ffi.ndarray import _set_class_ndarray
from . import backend as F
from ._ffi.function import _init_api
from ._ffi.ndarray import (
DGLContext,
DGLDataType,
NDArrayBase,
_set_class_ndarray,
context,
empty,
empty_shared_mem,
from_dlpack,
numpyasarray,
)
from ._ffi.object import ObjectBase, register_object
class NDArray(NDArrayBase):
"""Lightweight NDArray class for DGL framework."""
def __len__(self):
return functools.reduce(operator.mul, self.shape, 1)
......@@ -35,7 +46,10 @@ class NDArray(NDArrayBase):
-------
NDArray
"""
return empty_shared_mem(name, True, self.shape, self.dtype).copyfrom(self)
return empty_shared_mem(name, True, self.shape, self.dtype).copyfrom(
self
)
def cpu(dev_id=0):
"""Construct a CPU device
......@@ -52,6 +66,7 @@ def cpu(dev_id=0):
"""
return DGLContext(1, dev_id)
def gpu(dev_id=0):
"""Construct a CPU device
......@@ -67,6 +82,7 @@ def gpu(dev_id=0):
"""
return DGLContext(2, dev_id)
def array(arr, ctx=cpu(0)):
"""Create an array from source arr.
......@@ -87,6 +103,7 @@ def array(arr, ctx=cpu(0)):
arr = _np.array(arr)
return empty(arr.shape, arr.dtype, ctx).copyfrom(arr)
def zerocopy_from_numpy(np_data):
"""Create an array that shares the given numpy data.
......@@ -104,6 +121,7 @@ def zerocopy_from_numpy(np_data):
handle = ctypes.pointer(arr)
return NDArray(handle, is_view=True)
def cast_to_signed(arr):
"""Cast this NDArray from unsigned integer to signed one.
......@@ -124,8 +142,9 @@ def cast_to_signed(arr):
"""
return _CAPI_DGLArrayCastToSigned(arr)
def get_shared_mem_array(name, shape, dtype):
""" Get a tensor from shared memory with specific name
"""Get a tensor from shared memory with specific name
Parameters
----------
......@@ -141,12 +160,15 @@ def get_shared_mem_array(name, shape, dtype):
F.tensor
The tensor got from shared memory.
"""
new_arr = empty_shared_mem(name, False, shape, F.reverse_data_type_dict[dtype])
new_arr = empty_shared_mem(
name, False, shape, F.reverse_data_type_dict[dtype]
)
dlpack = new_arr.to_dlpack()
return F.zerocopy_from_dlpack(dlpack)
def create_shared_mem_array(name, shape, dtype):
""" Create a tensor from shared memory with the specific name
"""Create a tensor from shared memory with the specific name
Parameters
----------
......@@ -162,12 +184,15 @@ def create_shared_mem_array(name, shape, dtype):
F.tensor
The created tensor.
"""
new_arr = empty_shared_mem(name, True, shape, F.reverse_data_type_dict[dtype])
new_arr = empty_shared_mem(
name, True, shape, F.reverse_data_type_dict[dtype]
)
dlpack = new_arr.to_dlpack()
return F.zerocopy_from_dlpack(dlpack)
def exist_shared_mem_array(name):
""" Check the existence of shared-memory array.
"""Check the existence of shared-memory array.
Parameters
----------
......@@ -181,23 +206,27 @@ def exist_shared_mem_array(name):
"""
return _CAPI_DGLExistSharedMemArray(name)
class SparseFormat:
"""Format code"""
ANY = 0
COO = 1
CSR = 2
CSC = 3
FORMAT2STR = {
0 : 'ANY',
1 : 'COO',
2 : 'CSR',
3 : 'CSC',
0: "ANY",
1: "COO",
2: "CSR",
3: "CSC",
}
@register_object('aten.SparseMatrix')
@register_object("aten.SparseMatrix")
class SparseMatrix(ObjectBase):
"""Sparse matrix object class in C++ backend."""
@property
def format(self):
"""Sparse format enum
......@@ -250,17 +279,26 @@ class SparseMatrix(ObjectBase):
return _CAPI_DGLSparseMatrixGetFlags(self)
def __getstate__(self):
return self.format, self.num_rows, self.num_cols, self.indices, self.flags
return (
self.format,
self.num_rows,
self.num_cols,
self.indices,
self.flags,
)
def __setstate__(self, state):
fmt, nrows, ncols, indices, flags = state
indices = [F.zerocopy_to_dgl_ndarray(idx) for idx in indices]
self.__init_handle_by_constructor__(
_CAPI_DGLCreateSparseMatrix, fmt, nrows, ncols, indices, flags)
_CAPI_DGLCreateSparseMatrix, fmt, nrows, ncols, indices, flags
)
def __repr__(self):
return 'SparseMatrix(fmt="{}", shape=({},{}))'.format(
SparseFormat.FORMAT2STR[self.format], self.num_rows, self.num_cols)
SparseFormat.FORMAT2STR[self.format], self.num_rows, self.num_cols
)
_set_class_ndarray(NDArray)
_init_api("dgl.ndarray")
......@@ -270,5 +308,5 @@ _init_api("dgl.ndarray.uvm", __name__)
# other backend tensors.
NULL = {
"int64": array(_np.array([], dtype=_np.int64)),
"int32": array(_np.array([], dtype=_np.int32))
"int32": array(_np.array([], dtype=_np.int32)),
}
python/dgl/network.py
View file @ 76bb5404
......@@ -2,13 +2,14 @@
from __future__ import absolute_import
import time
from enum import Enum
from collections import namedtuple
from enum import Enum
import dgl.backend as F
from ._ffi.function import _init_api
from ._deprecate.nodeflow import NodeFlow
from . import utils
from ._deprecate.nodeflow import NodeFlow
from ._ffi.function import _init_api
_init_api("dgl.network")
......@@ -19,12 +20,11 @@ _WAIT_TIME_SEC = 3 # 3 seconds
def _network_wait():
"""Sleep for a few seconds
"""
"""Sleep for a few seconds"""
time.sleep(_WAIT_TIME_SEC)
def _create_sender(net_type, msg_queue_size=2*1024*1024*1024):
def _create_sender(net_type, msg_queue_size=2 * 1024 * 1024 * 1024):
"""Create a Sender communicator via C api
Parameters
......@@ -34,11 +34,11 @@ def _create_sender(net_type, msg_queue_size=2*1024*1024*1024):
msg_queue_size : int
message queue size (2GB by default)
"""
assert net_type in ('socket', 'mpi'), 'Unknown network type.'
assert net_type in ("socket", "mpi"), "Unknown network type."
return _CAPI_DGLSenderCreate(net_type, msg_queue_size)
def _create_receiver(net_type, msg_queue_size=2*1024*1024*1024):
def _create_receiver(net_type, msg_queue_size=2 * 1024 * 1024 * 1024):
"""Create a Receiver communicator via C api
Parameters
......@@ -48,7 +48,7 @@ def _create_receiver(net_type, msg_queue_size=2*1024*1024*1024):
msg_queue_size : int
message queue size (2GB by default)
"""
assert net_type in ('socket', 'mpi'), 'Unknown network type.'
assert net_type in ("socket", "mpi"), "Unknown network type."
return _CAPI_DGLReceiverCreate(net_type, msg_queue_size)
......@@ -64,8 +64,7 @@ def _finalize_sender(sender):
def _finalize_receiver(receiver):
"""Finalize Receiver Communicator
"""
"""Finalize Receiver Communicator"""
_CAPI_DGLFinalizeReceiver(receiver)
......@@ -83,7 +82,7 @@ def _add_receiver_addr(sender, ip_addr, port, recv_id):
recv_id : int
Receiver ID
"""
assert recv_id >= 0, 'recv_id cannot be a negative number.'
assert recv_id >= 0, "recv_id cannot be a negative number."
_CAPI_DGLSenderAddReceiver(sender, ip_addr, int(port), int(recv_id))
......@@ -112,7 +111,7 @@ def _receiver_wait(receiver, ip_addr, port, num_sender):
num_sender : int
total number of Sender
"""
assert num_sender >= 0, 'num_sender cannot be a negative number.'
assert num_sender >= 0, "num_sender cannot be a negative number."
_CAPI_DGLReceiverWait(receiver, ip_addr, int(port), int(num_sender))
......@@ -131,19 +130,22 @@ def _send_nodeflow(sender, nodeflow, recv_id):
recv_id : int
Receiver ID
"""
assert recv_id >= 0, 'recv_id cannot be a negative number.'
assert recv_id >= 0, "recv_id cannot be a negative number."
gidx = nodeflow._graph
node_mapping = nodeflow._node_mapping.todgltensor()
edge_mapping = nodeflow._edge_mapping.todgltensor()
layers_offsets = utils.toindex(nodeflow._layer_offsets).todgltensor()
flows_offsets = utils.toindex(nodeflow._block_offsets).todgltensor()
_CAPI_SenderSendNodeFlow(sender,
int(recv_id),
gidx,
node_mapping,
edge_mapping,
layers_offsets,
flows_offsets)
_CAPI_SenderSendNodeFlow(
sender,
int(recv_id),
gidx,
node_mapping,
edge_mapping,
layers_offsets,
flows_offsets,
)
def _send_sampler_end_signal(sender, recv_id):
"""Send an epoch-end signal to remote Receiver.
......@@ -155,9 +157,10 @@ def _send_sampler_end_signal(sender, recv_id):
recv_id : int
Receiver ID
"""
assert recv_id >= 0, 'recv_id cannot be a negative number.'
assert recv_id >= 0, "recv_id cannot be a negative number."
_CAPI_SenderSendSamplerEndSignal(sender, int(recv_id))
def _recv_nodeflow(receiver, graph):
"""Receive sampled subgraph (NodeFlow) from remote sampler.
......@@ -183,8 +186,8 @@ def _recv_nodeflow(receiver, graph):
class KVMsgType(Enum):
"""Type of kvstore message
"""
"""Type of kvstore message"""
FINAL = 1
INIT = 2
PUSH = 3
......@@ -215,6 +218,7 @@ c_ptr : void*
c pointer of message
"""
def _send_kv_msg(sender, msg, recv_id):
"""Send kvstore message.
......@@ -230,12 +234,8 @@ def _send_kv_msg(sender, msg, recv_id):
if msg.type == KVMsgType.PULL:
tensor_id = F.zerocopy_to_dgl_ndarray(msg.id)
_CAPI_SenderSendKVMsg(
sender,
int(recv_id),
msg.type.value,
msg.rank,
msg.name,
tensor_id)
sender, int(recv_id), msg.type.value, msg.rank, msg.name, tensor_id
)
elif msg.type in (KVMsgType.INIT, KVMsgType.GET_SHAPE_BACK):
tensor_shape = F.zerocopy_to_dgl_ndarray(msg.shape)
_CAPI_SenderSendKVMsg(
......@@ -244,20 +244,14 @@ def _send_kv_msg(sender, msg, recv_id):
msg.type.value,
msg.rank,
msg.name,
tensor_shape)
tensor_shape,
)
elif msg.type in (KVMsgType.IP_ID, KVMsgType.GET_SHAPE):
_CAPI_SenderSendKVMsg(
sender,
int(recv_id),
msg.type.value,
msg.rank,
msg.name)
sender, int(recv_id), msg.type.value, msg.rank, msg.name
)
elif msg.type in (KVMsgType.FINAL, KVMsgType.BARRIER):
_CAPI_SenderSendKVMsg(
sender,
int(recv_id),
msg.type.value,
msg.rank)
_CAPI_SenderSendKVMsg(sender, int(recv_id), msg.type.value, msg.rank)
else:
tensor_id = F.zerocopy_to_dgl_ndarray(msg.id)
data = F.zerocopy_to_dgl_ndarray(msg.data)
......@@ -268,7 +262,8 @@ def _send_kv_msg(sender, msg, recv_id):
msg.rank,
msg.name,
tensor_id,
data)
data,
)
def _recv_kv_msg(receiver):
......@@ -288,7 +283,9 @@ def _recv_kv_msg(receiver):
rank = _CAPI_ReceiverGetKVMsgRank(msg_ptr)
if msg_type == KVMsgType.PULL:
name = _CAPI_ReceiverGetKVMsgName(msg_ptr)
tensor_id = F.zerocopy_from_dgl_ndarray(_CAPI_ReceiverGetKVMsgID(msg_ptr))
tensor_id = F.zerocopy_from_dgl_ndarray(
_CAPI_ReceiverGetKVMsgID(msg_ptr)
)
msg = KVStoreMsg(
type=msg_type,
rank=rank,
......@@ -296,11 +293,14 @@ def _recv_kv_msg(receiver):
id=tensor_id,
data=None,
shape=None,
c_ptr=msg_ptr)
c_ptr=msg_ptr,
)
return msg
elif msg_type in (KVMsgType.INIT, KVMsgType.GET_SHAPE_BACK):
name = _CAPI_ReceiverGetKVMsgName(msg_ptr)
tensor_shape = F.zerocopy_from_dgl_ndarray(_CAPI_ReceiverGetKVMsgShape(msg_ptr))
tensor_shape = F.zerocopy_from_dgl_ndarray(
_CAPI_ReceiverGetKVMsgShape(msg_ptr)
)
msg = KVStoreMsg(
type=msg_type,
rank=rank,
......@@ -308,7 +308,8 @@ def _recv_kv_msg(receiver):
id=None,
data=None,
shape=tensor_shape,
c_ptr=msg_ptr)
c_ptr=msg_ptr,
)
return msg
elif msg_type in (KVMsgType.IP_ID, KVMsgType.GET_SHAPE):
name = _CAPI_ReceiverGetKVMsgName(msg_ptr)
......@@ -319,7 +320,8 @@ def _recv_kv_msg(receiver):
id=None,
data=None,
shape=None,
c_ptr=msg_ptr)
c_ptr=msg_ptr,
)
return msg
elif msg_type in (KVMsgType.FINAL, KVMsgType.BARRIER):
msg = KVStoreMsg(
......@@ -329,11 +331,14 @@ def _recv_kv_msg(receiver):
id=None,
data=None,
shape=None,
c_ptr=msg_ptr)
c_ptr=msg_ptr,
)
return msg
else:
name = _CAPI_ReceiverGetKVMsgName(msg_ptr)
tensor_id = F.zerocopy_from_dgl_ndarray(_CAPI_ReceiverGetKVMsgID(msg_ptr))
tensor_id = F.zerocopy_from_dgl_ndarray(
_CAPI_ReceiverGetKVMsgID(msg_ptr)
)
data = F.zerocopy_from_dgl_ndarray(_CAPI_ReceiverGetKVMsgData(msg_ptr))
msg = KVStoreMsg(
type=msg_type,
......@@ -342,25 +347,34 @@ def _recv_kv_msg(receiver):
id=tensor_id,
data=data,
shape=None,
c_ptr=msg_ptr)
c_ptr=msg_ptr,
)
return msg
raise RuntimeError('Unknown message type: %d' % msg_type.value)
raise RuntimeError("Unknown message type: %d" % msg_type.value)
def _clear_kv_msg(msg):
"""Clear data of kvstore message
"""
"""Clear data of kvstore message"""
F.sync()
if msg.c_ptr is not None:
_CAPI_DeleteKVMsg(msg.c_ptr)
def _fast_pull(name, id_tensor,
machine_count, group_count, machine_id, client_id,
partition_book, g2l, local_data,
sender, receiver):
""" Pull message
def _fast_pull(
name,
id_tensor,
machine_count,
group_count,
machine_id,
client_id,
partition_book,
g2l,
local_data,
sender,
receiver,
):
"""Pull message
Parameters
----------
......@@ -393,17 +407,33 @@ def _fast_pull(name, id_tensor,
target tensor
"""
if g2l is not None:
res_tensor = _CAPI_FastPull(name, machine_id, machine_count, group_count, client_id,
F.zerocopy_to_dgl_ndarray(id_tensor),
F.zerocopy_to_dgl_ndarray(partition_book),
F.zerocopy_to_dgl_ndarray(local_data),
sender, receiver, 'has_g2l',
F.zerocopy_to_dgl_ndarray(g2l))
res_tensor = _CAPI_FastPull(
name,
machine_id,
machine_count,
group_count,
client_id,
F.zerocopy_to_dgl_ndarray(id_tensor),
F.zerocopy_to_dgl_ndarray(partition_book),
F.zerocopy_to_dgl_ndarray(local_data),
sender,
receiver,
"has_g2l",
F.zerocopy_to_dgl_ndarray(g2l),
)
else:
res_tensor = _CAPI_FastPull(name, machine_id, machine_count, group_count, client_id,
F.zerocopy_to_dgl_ndarray(id_tensor),
F.zerocopy_to_dgl_ndarray(partition_book),
F.zerocopy_to_dgl_ndarray(local_data),
sender, receiver, 'no_g2l')
res_tensor = _CAPI_FastPull(
name,
machine_id,
machine_count,
group_count,
client_id,
F.zerocopy_to_dgl_ndarray(id_tensor),
F.zerocopy_to_dgl_ndarray(partition_book),
F.zerocopy_to_dgl_ndarray(local_data),
sender,
receiver,
"no_g2l",
)
return F.zerocopy_from_dgl_ndarray(res_tensor)
python/dgl/nn/__init__.py
View file @ 76bb5404
......@@ -14,20 +14,22 @@ with "[NN] XXX module".
"""
import importlib
import sys
import os
import sys
from ..backend import backend_name
from ..utils import expand_as_pair
# [BarclayII] Not sure what's going on with pylint.
# Possible issue: https://github.com/PyCQA/pylint/issues/2648
from . import functional # pylint: disable=import-self
from . import functional # pylint: disable=import-self
from ..backend import backend_name
from ..utils import expand_as_pair
def _load_backend(mod_name):
mod = importlib.import_module('.%s' % mod_name, __name__)
mod = importlib.import_module(".%s" % mod_name, __name__)
thismod = sys.modules[__name__]
for api, obj in mod.__dict__.items():
setattr(thismod, api, obj)
_load_backend(backend_name)
python/dgl/nn/mxnet/conv/__init__.py
View file @ 76bb5404
"""MXNet modules for graph convolutions."""
# pylint: disable= no-member, arguments-differ, invalid-name
from .graphconv import GraphConv
from .relgraphconv import RelGraphConv
from .tagconv import TAGConv
from .gatconv import GATConv
from .sageconv import SAGEConv
from .gatedgraphconv import GatedGraphConv
from .chebconv import ChebConv
from .agnnconv import AGNNConv
from .appnpconv import APPNPConv
from .chebconv import ChebConv
from .densechebconv import DenseChebConv
from .densegraphconv import DenseGraphConv
from .densesageconv import DenseSAGEConv
from .densechebconv import DenseChebConv
from .edgeconv import EdgeConv
from .gatconv import GATConv
from .gatedgraphconv import GatedGraphConv
from .ginconv import GINConv
from .gmmconv import GMMConv
from .graphconv import GraphConv
from .nnconv import NNConv
from .relgraphconv import RelGraphConv
from .sageconv import SAGEConv
from .sgconv import SGConv
from .tagconv import TAGConv
__all__ = ['GraphConv', 'TAGConv', 'RelGraphConv', 'GATConv',
'SAGEConv', 'GatedGraphConv', 'ChebConv', 'AGNNConv',
'APPNPConv', 'DenseGraphConv', 'DenseSAGEConv', 'DenseChebConv',
'EdgeConv', 'GINConv', 'GMMConv', 'NNConv', 'SGConv']
__all__ = [
"GraphConv",
"TAGConv",
"RelGraphConv",
"GATConv",
"SAGEConv",
"GatedGraphConv",
"ChebConv",
"AGNNConv",
"APPNPConv",
"DenseGraphConv",
"DenseSAGEConv",
"DenseChebConv",
"EdgeConv",
"GINConv",
"GMMConv",
"NNConv",
"SGConv",
]
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment