wehub-resource-sync
574 lines
13 KiB
Python
574 lines
13 KiB
Python
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
|