chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,2 @@
|
||||
from .sparse import *
|
||||
from .tensor import *
|
||||
@@ -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)
|
||||
@@ -0,0 +1 @@
|
||||
"""Sparse optimizer is not supported for mxnet"""
|
||||
@@ -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
|
||||
@@ -0,0 +1,2 @@
|
||||
from .sparse import *
|
||||
from .tensor import *
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -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])
|
||||
@@ -0,0 +1,6 @@
|
||||
import os
|
||||
|
||||
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
|
||||
|
||||
from .sparse import *
|
||||
from .tensor import *
|
||||
@@ -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)
|
||||
@@ -0,0 +1 @@
|
||||
"""Sparse optimizer is not supported for tensorflow"""
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user