wehub-resource-sync
2103 lines
46 KiB
Python
2103 lines
46 KiB
Python
"""This file defines the unified tensor framework interface required by DGL.
|
|
|
|
The principles of this interface:
|
|
* There should be as few interfaces as possible.
|
|
* The interface is used by DGL system so it is more important to have
|
|
clean definition rather than convenient usage.
|
|
* Default arguments should be avoided.
|
|
* Keyword or positional arguments should be avoided.
|
|
* Argument type should be easier to understand.
|
|
|
|
It is recommended the frameworks implement all the interfaces. However, it is
|
|
also OK to skip some. The generated backend module has an ``is_enabled`` function
|
|
that returns whether the interface is supported by the framework or not.
|
|
"""
|
|
|
|
###############################################################################
|
|
# Tensor, data type and context interfaces
|
|
|
|
|
|
def data_type_dict():
|
|
"""Returns a dictionary from data type string to the data type.
|
|
|
|
The dictionary should include at least:
|
|
bfloat16
|
|
float16
|
|
float32
|
|
float64
|
|
uint8
|
|
int8
|
|
int16
|
|
int32
|
|
int64
|
|
bool
|
|
|
|
This function will be called only *once* during the initialization fo the
|
|
backend module. The returned dictionary will become the attributes of the
|
|
backend module.
|
|
|
|
Examples
|
|
--------
|
|
>>> import torch as th
|
|
>>> def data_type_dict():
|
|
>>> return { 'float16' : th.float16, 'float32' : th.float32, ... }
|
|
|
|
After the module is initialized.
|
|
|
|
>>> import backend as F
|
|
>>> F.float16 # this will point to torch.float16
|
|
|
|
Returns
|
|
-------
|
|
dict of str to data type
|
|
The data type dict.
|
|
"""
|
|
pass
|
|
|
|
|
|
def cpu():
|
|
"""Return a context object for CPU device."""
|
|
pass
|
|
|
|
|
|
def tensor(data, dtype=None):
|
|
"""Create a tensor given the data and data type.
|
|
|
|
If the input is already a tensor and has the same dtype,
|
|
directly return.
|
|
|
|
Scalar input is converted to a array of one element instead of
|
|
a 0-dim tensor to avoid certain issues with some backends.
|
|
|
|
Parameters
|
|
----------
|
|
data : int, iterable, Tensor
|
|
The interface should at least support list and numpy array.
|
|
The data is copied to a newly-allocated tensor.
|
|
dtype : data type, optional
|
|
It should be one of the values in the data type dict.
|
|
If is none, the type should be inferred from data.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def as_scalar(data):
|
|
"""Returns a scalar whose value is copied from this array.
|
|
|
|
Parameters
|
|
----------
|
|
data : Tensor
|
|
The input data
|
|
|
|
Returns
|
|
-------
|
|
scalar
|
|
The scalar value in the tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
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.
|
|
|
|
Returns
|
|
-------
|
|
string
|
|
the name of the preferred sparse matrix format.
|
|
"""
|
|
pass
|
|
|
|
|
|
def sparse_matrix(data, index, shape, force_format=False):
|
|
"""Create a sparse matrix.
|
|
|
|
NOTE: Please make sure that the data and index tensors are not
|
|
copied. This is critical to the performance.
|
|
|
|
Parameters
|
|
----------
|
|
data : Tensor
|
|
Data tensor. It should be of shape (nnz,).
|
|
index : tuple
|
|
This is used to support different sparse formats.
|
|
For COO format:
|
|
index=('coo', coord), where coord is of shape (2, nnz).
|
|
coord[0,:] should be the row index and coord[1,:] should be
|
|
the column index.
|
|
For CSR format:
|
|
index=('csr', indices, indptr), where indices is of shape (nnz,)
|
|
and indptr is of shape (nrows+1,). See ``scipy.sparse.csr_matrix``
|
|
for more documents on what each array means.
|
|
shape : tuple of int
|
|
The shape.
|
|
force_format : bool
|
|
If true, the returned sparse matrix must be stored in the same
|
|
format as the given index.
|
|
|
|
Returns
|
|
-------
|
|
SparseMatrix
|
|
The framework-specific sparse matrix. It can be stored in any format
|
|
unless force_format is True.
|
|
Tensor
|
|
The data convert index due to sparse format change.
|
|
None if no conversion is needed.
|
|
"""
|
|
pass
|
|
|
|
|
|
def sparse_matrix_indices(spmat):
|
|
"""Return the indices of the given sparse matrix.
|
|
|
|
Parameters
|
|
----------
|
|
spmat : SparseMatrix
|
|
The framework-specific sparse matrix.
|
|
|
|
Returns
|
|
-------
|
|
index : tuple
|
|
This is used to support different sparse formats.
|
|
For COO format:
|
|
index=('coo', coord), where coord is of shape (2, nnz).
|
|
coord[0,:] should be the row index and coord[1,:] should be
|
|
the column index.
|
|
For CSR format:
|
|
index=('csr', indices, indptr), where indices is of shape (nnz,)
|
|
and indptr is of shape (nrows+1,). See ``scipy.sparse.csr_matrix``
|
|
for more documents on what each array means.
|
|
"""
|
|
pass
|
|
|
|
|
|
def is_tensor(obj):
|
|
"""Returns true if the given object is a framework-specific tensor."""
|
|
pass
|
|
|
|
|
|
def shape(input):
|
|
"""Return the shape of the tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
tuple of int
|
|
The tensor shape.
|
|
"""
|
|
pass
|
|
|
|
|
|
def dtype(input):
|
|
"""Return the data type of the tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
data type
|
|
It should be one of the values in the data type dict.
|
|
"""
|
|
pass
|
|
|
|
|
|
def ndim(input):
|
|
"""Return the number of dimensions of the tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of dimensions
|
|
"""
|
|
pass
|
|
|
|
|
|
def context(input):
|
|
"""Return the context/device of the input tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
Context object
|
|
A framework-specific context object.
|
|
"""
|
|
pass
|
|
|
|
|
|
def device_type(ctx):
|
|
"""Return a str representing device type.
|
|
|
|
Parameters
|
|
----------
|
|
ctx : Device context object.
|
|
Device context.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
"""
|
|
pass
|
|
|
|
|
|
def device_id(ctx):
|
|
"""Return device index.
|
|
|
|
For CPU, the index does not matter. For GPU, the index means which GPU
|
|
device on the machine.
|
|
|
|
Parameters
|
|
----------
|
|
ctx : Device context object.
|
|
Device context.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The device index.
|
|
"""
|
|
pass
|
|
|
|
|
|
def to_backend_ctx(dglctx):
|
|
"""Convert a DGL context object to a backend context.
|
|
|
|
Parameters
|
|
----------
|
|
dglctx : dgl.ndarray.DGLContext
|
|
DGL context object. See _ffi.runtime_types for definition.
|
|
|
|
Returns
|
|
-------
|
|
ctx : framework-specific context object.
|
|
"""
|
|
pass
|
|
|
|
|
|
def astype(input, ty):
|
|
"""Convert the input tensor to the given data type.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
ty : data type
|
|
It should be one of the values in the data type dict.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def asnumpy(input):
|
|
"""Convert the input tensor to numpy array.
|
|
|
|
The data is copied.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
Numpy array.
|
|
"""
|
|
pass
|
|
|
|
|
|
def copy_to(input, ctx, **kwargs):
|
|
"""Copy the given tensor to the context.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor
|
|
ctx :
|
|
A framework-specific context object.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The tensor on the given context.
|
|
"""
|
|
pass
|
|
|
|
|
|
def is_pinned(input):
|
|
"""Check whether the tensor is in pinned memory.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The tensor.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
Whether the tensor is in pinned memory.
|
|
"""
|
|
pass
|
|
|
|
|
|
###############################################################################
|
|
# Tensor functions on feature data
|
|
# --------------------------------
|
|
# These functions are performance critical, so it's better to have efficient
|
|
# implementation in each framework.
|
|
|
|
|
|
def sum(input, dim, keepdims=False):
|
|
"""Reduce sum the input tensor along the given dim.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
dim : int
|
|
The reduce dim.
|
|
keepdims : bool
|
|
Whether to keep the summed dimension.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def floor_div(in1, in2):
|
|
"""Element-wise integer division and rounds each quotient towards zero.
|
|
|
|
Parameters
|
|
----------
|
|
in1 : Tensor
|
|
The input tensor
|
|
in2 : Tensor or integer
|
|
The input
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
|
|
|
|
def reduce_sum(input):
|
|
"""Returns the sum of all elements in the input tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor with shape (1,)
|
|
"""
|
|
pass
|
|
|
|
|
|
def cumsum(input, dim):
|
|
"""Return the cumulative sum of the elements along a given axis.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
dim : int
|
|
The cumulative dimension.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def mean(input, dim):
|
|
"""Reduce average the input tensor along the given dim.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
dim : int
|
|
The reduce dim.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def reduce_mean(input):
|
|
"""Returns the average of all elements in the input tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor with shape (1,)
|
|
"""
|
|
pass
|
|
|
|
|
|
def max(input, dim):
|
|
"""Reduce max the input tensor along the given dim.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
dim : int
|
|
The reduce dim.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def reduce_max(input):
|
|
"""Returns the max of all elements in the input tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor with shape (1,)
|
|
"""
|
|
pass
|
|
|
|
|
|
def min(input, dim):
|
|
"""Reduce min the input tensor along the given dim.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
dim : int
|
|
The reduce dim.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def reduce_min(input):
|
|
"""Returns the min of all elements in the input tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor with shape (1,)
|
|
"""
|
|
pass
|
|
|
|
|
|
def argsort(input, dim, descending):
|
|
"""Return the indices that would sort the input along the given dim.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
dim : int
|
|
The dim to sort along.
|
|
descending : bool
|
|
Controls the sorting order (False: ascending, True: descending)
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
|
|
|
|
def topk(input, k, dim, descending=True):
|
|
"""Return the k largest elements of the given input tensor along the given dimension.
|
|
|
|
If descending is False then the k smallest elements are returned.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
k : int
|
|
The number of elements.
|
|
dim : int
|
|
The dim to sort along.
|
|
descending : bool
|
|
Controls whether to return largest/smallest elements.
|
|
"""
|
|
pass
|
|
|
|
|
|
def argtopk(input, k, dim, descending=True):
|
|
"""Return the indices of the k largest elements of the given input tensor
|
|
along the given dimension.
|
|
|
|
If descending is False then the k smallest elements are returned.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
k : int
|
|
The number of elements.
|
|
dim : int
|
|
The dimension to sort along.
|
|
descending : bool
|
|
Controls whether to return largest/smallest elements.
|
|
"""
|
|
pass
|
|
|
|
|
|
def exp(input):
|
|
"""Returns a new tensor with the exponential of the elements of the input tensor `input`.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The output tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def inverse(input):
|
|
"""Returns the inverse matrix of a square matrix if it exists.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input square matrix.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The output tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def sqrt(input):
|
|
"""Returns a new tensor with the square root of the elements of the input tensor `input`.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The output tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def softmax(input, dim=-1):
|
|
"""Apply the softmax function on given dimension.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
dim : int
|
|
The dimension along which to compute softmax.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The output tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def cat(seq, dim):
|
|
"""Concat the sequence of tensors in the given dimension.
|
|
|
|
Parameters
|
|
----------
|
|
seq : list of Tensor
|
|
The tensor sequence.
|
|
dim : int
|
|
The concat dim.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def stack(seq, dim):
|
|
"""Stack the sequence of tensors along the given dimension.
|
|
|
|
Parameters
|
|
----------
|
|
seq : list of Tensor
|
|
The tensor sequence.
|
|
dim : int
|
|
The concat dim.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def split(input, sizes_or_sections, dim):
|
|
"""Split the input tensor into chunks.
|
|
|
|
If ``sizes_or_sections`` is an integer, then the tensor will
|
|
be splitted into equal pieces.
|
|
|
|
If ``sizes_or_sections`` is a list, then the tensor will be
|
|
splitted into segments.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
Tensor to split.
|
|
sizes_or_sections : int, list[int]
|
|
Split sizes or sections.
|
|
dim : int
|
|
The dimension to split on.
|
|
|
|
Returns
|
|
-------
|
|
list of Tensor
|
|
The splitted tensors.
|
|
"""
|
|
pass
|
|
|
|
|
|
def repeat(input, repeats, dim):
|
|
"""Repeats elements of an array.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
Input data array
|
|
repeats : int, Tensor
|
|
The number of repetitions for each element
|
|
dim : int
|
|
The dim along which to repeat values.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The obtained tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def gather_row(data, row_index):
|
|
"""Slice out the data given the row index.
|
|
|
|
Parameters
|
|
----------
|
|
data : Tensor
|
|
The data tensor
|
|
row_index : Tensor
|
|
A 1-D integer tensor containing which rows to be sliced out.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The sliced data. The first dimension should equal to ``len(row_index)``.
|
|
"""
|
|
pass
|
|
|
|
|
|
def slice_axis(data, axis, begin, end):
|
|
"""Slice along a given axis.
|
|
Returns an array slice along a given axis starting from :attr:`begin` index to :attr:`end` index.
|
|
|
|
Parameters
|
|
----------
|
|
data : Tensor
|
|
The data tensor.
|
|
axis : int
|
|
The axis along to slice the tensor.
|
|
begin : int
|
|
Indicates the begin index.
|
|
end : int
|
|
Indicates the end index.
|
|
Returns:
|
|
--------
|
|
Tensor
|
|
The sliced tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def take(data, indices, dim):
|
|
"""Takes elements from an input array along the given dim.
|
|
|
|
Parameters
|
|
----------
|
|
data : Tensor
|
|
The data tensor.
|
|
indices : Tensor
|
|
The indices tensor.
|
|
dim : Tensor
|
|
The dimension to gather along.
|
|
"""
|
|
pass
|
|
|
|
|
|
def narrow_row(x, start, stop):
|
|
"""Narrow down the tensor along the first dimension.
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
The input tensor.
|
|
start : int
|
|
The start index (inclusive).
|
|
stop : int
|
|
The stop index (exclusive).
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The narrowed tensor
|
|
|
|
Notes
|
|
-----
|
|
The returned tensor could be a view of the original tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def scatter_row(data, row_index, value):
|
|
"""Write the value into the data tensor using the row index.
|
|
|
|
This is an out-place write so it can work with autograd.
|
|
|
|
Parameters
|
|
----------
|
|
data : Tensor
|
|
The data tensor to be updated.
|
|
row_index : Tensor
|
|
A 1-D integer tensor containing which rows to be updated.
|
|
value : Tensor
|
|
The new value.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The new data.
|
|
"""
|
|
pass
|
|
|
|
|
|
def index_add_inplace(data, row_idx, value):
|
|
"""Add the values into the data tensor using the row index inplace.
|
|
|
|
If two row indices are the same, the corresponding values are sum up before
|
|
adding to the data tensor.
|
|
|
|
Examples
|
|
--------
|
|
>>> import torch as th
|
|
>>> arr = th.zeros((10))
|
|
>>> F. index_add_inplace(arr, th.tensor([0, 1, 1]), th.tensor([1.0, 1.0, 1.0]))
|
|
>>> arr
|
|
tensor([1., 2., 0., 0., 0., 0., 0., 0., 0., 0.])
|
|
|
|
Parameters
|
|
----------
|
|
data : Tensor
|
|
The data tensor to be updated.
|
|
row_index : Tensor
|
|
A 1-D integer tensor containing which rows to be updated.
|
|
value : Tensor
|
|
The new value.
|
|
"""
|
|
pass
|
|
|
|
|
|
def scatter_row_inplace(data, row_index, value):
|
|
"""Write the value into the data tensor using the row index inplace.
|
|
|
|
This is an inplace write so it will break the autograd.
|
|
|
|
Parameters
|
|
----------
|
|
data : Tensor
|
|
The data tensor to be updated.
|
|
row_index : Tensor
|
|
A 1-D integer tensor containing which rows to be updated.
|
|
value : Tensor
|
|
The new value.
|
|
"""
|
|
pass
|
|
|
|
|
|
def squeeze(input, dim):
|
|
"""Remove the given dimension of size 1.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
dim : int
|
|
The dimension to be squeezed.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def unsqueeze(input, dim):
|
|
"""Add the given dimension of size 1.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
dim : int
|
|
The dimension to be unsqueezed.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def reshape(input, shape):
|
|
"""Reshape the tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
shape : tuple of int
|
|
The new shape.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The reshaped tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def swapaxes(input, axis1, axis2):
|
|
"""Interchange the two given axes of a tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor.
|
|
axis1, axis2 : int
|
|
The two axes.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The transposed tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def empty(shape, dtype, ctx):
|
|
"""Create a tensor filled with uninitialized data.
|
|
|
|
Parameters
|
|
----------
|
|
shape : tuple of int
|
|
The tensor shape.
|
|
dtype : data type
|
|
It should be one of the values in the data type dict.
|
|
ctx : context
|
|
The device of the result tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The emtpy tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def zeros(shape, dtype, ctx):
|
|
"""Create a zero tensor.
|
|
|
|
Parameters
|
|
----------
|
|
shape : tuple of int
|
|
The tensor shape.
|
|
dtype : data type
|
|
It should be one of the values in the data type dict.
|
|
ctx : context
|
|
The device of the result tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The zero tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def zeros_like(input):
|
|
"""Create a zero tensor with the same shape, dtype and context of the
|
|
given tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result
|
|
"""
|
|
pass
|
|
|
|
|
|
def ones(shape, dtype, ctx):
|
|
"""Create a one tensor.
|
|
|
|
Parameters
|
|
----------
|
|
shape : tuple of int
|
|
The tensor shape.
|
|
dtype : data type
|
|
It should be one of the values in the data type dict.
|
|
ctx : context
|
|
The device of the result tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The one tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def uniform(shape, dtype, ctx, low, high):
|
|
"""Create a tensor with random value in a uniform
|
|
distribution between low (inclusive) and high (exclusive).
|
|
|
|
Parameters
|
|
----------
|
|
shape : tuple of int
|
|
The tensor shape.
|
|
dtype : data type
|
|
It should be one of the values in the data type dict.
|
|
ctx : context
|
|
The device of the result tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The random tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def randint(shape, dtype, ctx, low, high):
|
|
"""Create a tensor with random value in a uniform integer
|
|
distribution between low (inclusive) and high (exclusive)
|
|
|
|
Parameters
|
|
----------
|
|
shape : tuple of int
|
|
The tensor shape.
|
|
dtype : data type
|
|
It should be one of the values in the data type dict.
|
|
ctx : context
|
|
The device of the result tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The random tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def pad_packed_tensor(input, lengths, value, l_min=None):
|
|
r"""Pads a packed batch of variable length tensors with given value.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor with shape :math:`(N, *)`
|
|
lengths : list or tensor
|
|
The array of tensor lengths (of the first dimension) :math:`L`.
|
|
It should satisfy :math:`\sum_{i=1}^{B}L_i = N`,
|
|
where :math:`B` is the length of :math:`L`.
|
|
value : float
|
|
The value to fill in the tensor.
|
|
l_min : int or None, defaults to None.
|
|
The minimum length each tensor need to be padded to, if set to None,
|
|
then there is no minimum length requirement.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The obtained tensor with shape :math:`(B, \max(\max_i(L_i), l_{min}), *)`
|
|
"""
|
|
pass
|
|
|
|
|
|
def pack_padded_tensor(input, lengths):
|
|
r"""Packs a tensor containing padded sequence of variable length.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor with shape :math:`(B, L, *)`, where :math:`B` is
|
|
the batch size and :math:`L` is the maximum length of the batch.
|
|
lengths : list or tensor
|
|
The array of tensor lengths (of the first dimension) :math:`L`.
|
|
:math:`\max_i(L_i)` should equal :math:`L`.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The obtained tensor with shape :math:`(N, *)` where
|
|
:math:`N = \sum_{i=1}^{B}L_i`
|
|
"""
|
|
pass
|
|
|
|
|
|
def boolean_mask(input, mask):
|
|
"""Selects elements in x according to the given mask from the first
|
|
dimension.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor
|
|
mask : Boolean Tensor
|
|
The mask
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result
|
|
"""
|
|
pass
|
|
|
|
|
|
def equal(x, y):
|
|
"""Compares whether the elements are equal.
|
|
|
|
Parameters
|
|
----------
|
|
x, y : Tensor
|
|
The two tensors
|
|
|
|
Returns
|
|
-------
|
|
Boolean or integer tensor
|
|
The result, with the same shape as input.
|
|
"""
|
|
pass
|
|
|
|
|
|
def allclose(x, y, rtol=1e-4, atol=1e-4):
|
|
"""Compares whether all elements are close.
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
First tensor
|
|
y : Tensor
|
|
Second tensor
|
|
rtol : float, optional
|
|
Relative tolerance
|
|
atol : float, optional
|
|
Absolute tolerance
|
|
"""
|
|
|
|
|
|
def logical_not(input):
|
|
"""Perform a logical not operation. Equivalent to np.logical_not
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result
|
|
"""
|
|
pass
|
|
|
|
|
|
def logical_and(input1, input2):
|
|
pass
|
|
|
|
|
|
def clone(input):
|
|
"""Return a clone of the input tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
Input tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A clone tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def clamp(data, min_val, max_val):
|
|
"""Clamp all elements in :attr:`input` into the range [min_val, max_val]
|
|
and return a resulting tensor.
|
|
|
|
Parameters
|
|
----------
|
|
data : Tensor
|
|
Input tensor
|
|
min_val : Scalar
|
|
Min value.
|
|
max_val : Scalar
|
|
Max value.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result.
|
|
"""
|
|
pass
|
|
|
|
|
|
def replace_inf_with_zero(x):
|
|
"""Returns a new tensor replacing infinity and negative infinity with zeros.
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
The input
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result
|
|
"""
|
|
pass
|
|
|
|
|
|
def count_nonzero(input):
|
|
"""Return the count of non-zero values in the tensor input.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The tensor to be counted
|
|
|
|
Returns
|
|
-------
|
|
Integer
|
|
The result
|
|
"""
|
|
pass
|
|
|
|
|
|
###############################################################################
|
|
# Tensor functions used *only* on index tensor
|
|
# ----------------
|
|
# These operators are light-weighted, so it is acceptable to fallback to
|
|
# numpy operators if currently missing in the framework. Ideally in the future,
|
|
# DGL should contain all the operations on index, so this set of operators
|
|
# should be gradually removed.
|
|
|
|
|
|
def unique(input, return_inverse=False, return_counts=False):
|
|
"""Returns the unique scalar elements in a tensor.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
Must be a 1-D tensor.
|
|
return_inverse : bool, optional
|
|
Whether to also return the indices for where elements in the original
|
|
input ended up in the returned unique list.
|
|
return_counts : bool, optional
|
|
Whether to also return the counts for each unique element.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A 1-D tensor containing unique elements.
|
|
Tensor, optional
|
|
A 1-D tensor containing the new positions of the elements in the input.
|
|
It is returned if return_inverse is True.
|
|
Tensor, optional
|
|
A 1-D tensor containing the number of occurrences for each unique value or tensor.
|
|
It is returned if return_counts is True.
|
|
"""
|
|
pass
|
|
|
|
|
|
def full_1d(length, fill_value, dtype, ctx):
|
|
"""Create a 1D tensor full of the fill_value.
|
|
|
|
Parameters
|
|
----------
|
|
shape : int
|
|
The length of the vector.
|
|
fill_value : int
|
|
The filled value.
|
|
dtype : data type
|
|
It should be one of the values in the data type dict.
|
|
ctx : context
|
|
The device of the result tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A result 1D tensor
|
|
"""
|
|
pass
|
|
|
|
|
|
def nonzero_1d(input):
|
|
"""Return the nonzero index of the given 1D input.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
Must be a 1D tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A 1D integer tensor containing the nonzero indices.
|
|
"""
|
|
pass
|
|
|
|
|
|
def sort_1d(input):
|
|
"""Sort a 1D tensor (in ascending order) and also return the original index.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The tensor to be sorted.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
Sorted tensor.
|
|
Tensor
|
|
Index tensor of the elements in the original input.
|
|
"""
|
|
pass
|
|
|
|
|
|
def arange(start, stop, dtype, ctx):
|
|
"""Create a 1D range int64 tensor.
|
|
|
|
Parameters
|
|
----------
|
|
start : int
|
|
The range start.
|
|
stop : int
|
|
The range stop.
|
|
dtype: str
|
|
The dtype of result tensor.
|
|
ctx : Device context object.
|
|
Device context.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def rand_shuffle(arr):
|
|
"""Random shuffle the data in the first dimension of the array.
|
|
|
|
The shuffled data is stored in a new array.
|
|
|
|
Parameters
|
|
----------
|
|
arr : Tensor
|
|
The data tensor
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result tensor
|
|
"""
|
|
pass
|
|
|
|
|
|
def zerocopy_to_dlpack(input):
|
|
"""Create a dlpack tensor that shares the input memory.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor
|
|
|
|
Returns
|
|
-------
|
|
dlpack capsule
|
|
A dlpack capsule that can be used by other framework.
|
|
"""
|
|
pass
|
|
|
|
|
|
def zerocopy_from_dlpack(dlpack_tensor):
|
|
"""Create a tensor that shares the dlpack_tensor.
|
|
|
|
Parameters
|
|
----------
|
|
dlpack_tensor : dlpack capsule
|
|
The dlpack tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def zerocopy_to_numpy(input):
|
|
"""Create a numpy ndarray that shares the input memory.
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
The input tensor
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
A numpy ndarray.
|
|
"""
|
|
pass
|
|
|
|
|
|
def zerocopy_from_numpy(np_array):
|
|
"""Create a tensor that shares the numpy array.
|
|
|
|
Parameters
|
|
----------
|
|
np_array : numpy.ndarray
|
|
The numpy ndarray.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
A framework-specific tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def zerocopy_to_dgl_ndarray(input):
|
|
"""Zerocopy a framework-specific Tensor to dgl.ndarray.NDArray
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
|
|
Returns
|
|
-------
|
|
dgl.ndarray.NDArray
|
|
"""
|
|
pass
|
|
|
|
|
|
def zerocopy_to_dgl_ndarray_for_write(input):
|
|
"""Zerocopy a framework-specific Tensor to dgl.ndarray.NDArray
|
|
that is ready for write (required in MXNet).
|
|
|
|
Parameters
|
|
----------
|
|
input : Tensor
|
|
|
|
Returns
|
|
-------
|
|
dgl.ndarray.NDArray
|
|
"""
|
|
pass
|
|
|
|
|
|
def zerocopy_from_dgl_ndarray(input):
|
|
"""Zerocopy a dgl.ndarray.NDArray to framework-specific Tensor
|
|
|
|
Parameters
|
|
----------
|
|
input : dgl.ndarray.NDArray
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
"""
|
|
pass
|
|
|
|
|
|
###############################################################################
|
|
# Custom Operators for graph level computations.
|
|
|
|
# Note: These operators are supposed to be implemented using DGL-provided
|
|
# kernels (see kernel.py), and plug into tensor framework using custom op
|
|
# extensions.
|
|
|
|
|
|
def binary_reduce(
|
|
reducer,
|
|
binary_op,
|
|
graph,
|
|
lhs,
|
|
rhs,
|
|
lhs_data,
|
|
rhs_data,
|
|
out_size,
|
|
lhs_map,
|
|
rhs_map,
|
|
out_map,
|
|
):
|
|
"""Perform binary operation between given data and reduce based on graph
|
|
structure.
|
|
|
|
Parameters
|
|
----------
|
|
reducer : str
|
|
Type of reduction: 'sum', 'max', 'min', 'mean', 'prod', 'none' (no
|
|
reduction)
|
|
binary_op : str
|
|
Binary operation to perform, can be 'add', 'mul', 'sub', 'div'
|
|
graph : GraphIndex
|
|
The graph
|
|
lhs : int
|
|
The lhs target (src, dst, edge)
|
|
rhs : int
|
|
The rhs target (src, dst, edge)
|
|
lhs_data : Tensor
|
|
The lhs data
|
|
rhs_data : Tensor
|
|
The rhs data
|
|
out_size : int
|
|
Size of first dimension of output data
|
|
lhs_map : tuple
|
|
Two lhs id mapping arrays, one for forward pass, the other for backward
|
|
rhs_map : tuple
|
|
Two rhs id mapping arrays, one for forward pass, the other for backward
|
|
out_map : tuple
|
|
Two out id mapping arrays, one for forward pass, the other for backward
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result.
|
|
"""
|
|
pass
|
|
|
|
|
|
def copy_reduce(reducer, graph, target, in_data, out_size, in_map, out_map):
|
|
"""Copy target data and perform reduce based on graph structure.
|
|
|
|
Parameters
|
|
----------
|
|
reducer : str
|
|
Type of reduction: be 'sum', 'max', 'min', 'mean', 'prod', 'none' (no
|
|
reduction)
|
|
graph : GraphIndex
|
|
The graph
|
|
target : int
|
|
The input target (src, dst, edge)
|
|
in_data : Tensor
|
|
The input data
|
|
out_size : int
|
|
Size of first dimension of output data
|
|
in_map : tuple
|
|
Two input id mapping arrays, one for forward, the other for backward
|
|
out_map : tuple
|
|
Two output id mapping arrays, one for forward, the other for backward
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result.
|
|
"""
|
|
pass
|
|
|
|
|
|
def gspmm(gidx, op, reduce_op, lhs_data, rhs_data):
|
|
r"""Generalized Sparse Matrix Multiplication interface.
|
|
It fuses two steps into one kernel.
|
|
(1) Computes messages by :attr:`op` source node and edge features.
|
|
(2) Aggregate the messages by :attr:`reduce_op` as the features on destination nodes.
|
|
|
|
.. math::
|
|
x_v = \psi_{(u, v, e)\in \mathcal{G}}(\rho(x_u, x_e))
|
|
|
|
where :math:`x_v` is the returned feature on destination nodes, and :math`x_u`,
|
|
:math:`x_e` refers to :attr:`u`, :attr:`e` respectively. :math:`\rho` means binary
|
|
operator :attr:`op` and :math:`\psi` means reduce operator :attr:`reduce_op`,
|
|
:math:`\mathcal{G}` is the graph we apply gspmm on: :attr:`g`.
|
|
|
|
Note that this function does not handle gradients.
|
|
|
|
Parameters
|
|
----------
|
|
gidx : HeteroGraphIndex
|
|
The input graph.
|
|
op : str
|
|
The binary op's name, could be ``add``, ``sub``, ``mul``, ``div``,
|
|
``copy_lhs``, ``copy_rhs``.
|
|
reduce_op : str
|
|
Reduce operator, could be ``sum``, ``max``, ``min``.
|
|
lhs_data : tensor or None
|
|
The left operand, could be None if it's not required by the op.
|
|
rhs_data : tensor or None
|
|
The right operand, could be None if it's not required by the op.
|
|
|
|
Returns
|
|
-------
|
|
tensor
|
|
The result tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def gspmm_hetero(g, op, reduce_op, lhs_len, *lhs_and_rhs_tuple):
|
|
r"""Generalized Sparse Matrix Multiplication interface on heterogenenous graph.
|
|
All the relation types of the heterogeneous graph will be processed together.
|
|
It fuses two steps into one kernel.
|
|
(1) Computes messages by :attr:`op` source node and edge features.
|
|
(2) Aggregate the messages by :attr:`reduce_op` as the features on destination nodes.
|
|
|
|
.. math::
|
|
x_v = \psi_{(u, v, e)\in \mathcal{G}}(\rho(x_u, x_e))
|
|
|
|
where :math:`x_v` is the returned feature on destination nodes, and :math`x_u`,
|
|
:math:`x_e` refers to :attr:`u`, :attr:`e` respectively. :math:`\rho` means binary
|
|
operator :attr:`op` and :math:`\psi` means reduce operator :attr:`reduce_op`,
|
|
:math:`\mathcal{G}` is the graph we apply gspmm on: :attr:`g`.
|
|
|
|
Note that this function does not handle gradients.
|
|
|
|
Parameters
|
|
----------
|
|
g : HeteroGraph
|
|
The input graph.
|
|
op : str
|
|
The binary op's name, could be ``add``, ``sub``, ``mul``, ``div``,
|
|
``copy_lhs``, ``copy_rhs``.
|
|
reduce_op : str
|
|
Reduce operator, could be ``sum``, ``max``, ``min``.
|
|
lhs_len : int
|
|
Length of the lhs data
|
|
lhs_and_rhs_tuple : tuple of tensors
|
|
lhs_data and rhs_data are concatenated to one tuple. lhs_data is
|
|
also a tuple of tensors of size number of ntypes. Same is true for
|
|
rhs_data.
|
|
The tensor(s) in the tuple could be None
|
|
|
|
Returns
|
|
-------
|
|
tuple of tensor
|
|
The resulting tuple of tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def gsddmm(gidx, op, lhs_data, rhs_data, lhs_target="u", rhs_target="v"):
|
|
r"""Generalized Sampled-Dense-Dense Matrix Multiplication interface.
|
|
It computes edge features by :attr:`op` lhs features and rhs features.
|
|
|
|
.. math::
|
|
x_{e} = \phi(x_{lhs}, x_{rhs}), \forall (u,e,v)\in \mathcal{G}
|
|
|
|
where :math:`x_{e}` is the returned feature on edges and :math:`x_u`,
|
|
:math:`x_v` refers to :attr:`u`, :attr:`v` respectively. :math:`\phi`
|
|
is the binary operator :attr:`op`, and :math:`\mathcal{G}` is the graph
|
|
we apply gsddmm on: :attr:`g`. $lhs$ and $rhs$ are one of $u,v,e$'s.
|
|
|
|
Parameters
|
|
----------
|
|
gidx : HeteroGraphIndex
|
|
The input graph.
|
|
op : str
|
|
Binary operator, could be ``add``, ``sub``, ``mul``, ``div``, ``dot``,
|
|
``copy_lhs``, ``copy_rhs``.
|
|
lhs_data : tensor or None
|
|
The left operand, could be None if it's not required by op.
|
|
rhs_data : tensor or None
|
|
The right operand, could be None if it's not required by op.
|
|
lhs_target: str
|
|
Choice of `u`(source), `e`(edge) or `v`(destination) for left operand.
|
|
rhs_target: str
|
|
Choice of `u`(source), `e`(edge) or `v`(destination) for right operand.
|
|
|
|
Returns
|
|
-------
|
|
tensor
|
|
The result tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def gsddmm_hetero(
|
|
g, op, lhs_len, lhs_target="u", rhs_target="v", *lhs_and_rhs_tuple
|
|
):
|
|
r"""Generalized Sampled-Dense-Dense Matrix Multiplication interface on
|
|
heterogenenous graph. All the relation types of the heterogeneous graph
|
|
will be processed together.
|
|
It computes edge features by :attr:`op` lhs features and rhs features.
|
|
|
|
.. math::
|
|
x_{e} = \phi(x_{lhs}, x_{rhs}), \forall (u,e,v)\in \mathcal{G}
|
|
|
|
where :math:`x_{e}` is the returned feature on edges and :math:`x_u`,
|
|
:math:`x_v` refers to :attr:`u`, :attr:`v` respectively. :math:`\phi`
|
|
is the binary operator :attr:`op`, and :math:`\mathcal{G}` is the graph
|
|
we apply gsddmm on: :attr:`g`. $lhs$ and $rhs$ are one of $u,v,e$'s.
|
|
|
|
Parameters
|
|
----------
|
|
gidx : HeteroGraphIndex
|
|
The input graph.
|
|
op : str
|
|
Binary operator, could be ``add``, ``sub``, ``mul``, ``div``, ``dot``,
|
|
``copy_lhs``, ``copy_rhs``.
|
|
lhs_len : int
|
|
Length of the lhs data
|
|
lhs_target: str
|
|
Choice of `u`(source), `e`(edge) or `v`(destination) for left operand.
|
|
rhs_target: str
|
|
Choice of `u`(source), `e`(edge) or `v`(destination) for right operand.
|
|
lhs_and_rhs_tuple : tuple of tensors
|
|
lhs_data and rhs_data are concatenated to one tuple. lhs_data is
|
|
also a tuple of tensors of size number of ntypes. Same is true for
|
|
rhs_data.
|
|
The tensor(s) in the tuple could be None
|
|
|
|
Returns
|
|
-------
|
|
tuple of tensor
|
|
The resulting tuple of tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def edge_softmax(gidx, logits, eids, norm_by):
|
|
r"""Compute edge softmax.
|
|
|
|
For a node :math:`i`, edge softmax is an operation of computing
|
|
|
|
.. math::
|
|
a_{ij} = \frac{\exp(z_{ij})}{\sum_{j\in\mathcal{N}(i)}\exp(z_{ij})}
|
|
|
|
where :math:`z_{ij}` is a signal of edge :math:`j\rightarrow i`, also
|
|
called logits in the context of softmax. :math:`\mathcal{N}(i)` is
|
|
the set of nodes that have an edge to :math:`i`.
|
|
|
|
By default edge softmax is normalized by destination nodes(i.e. :math:`ij`
|
|
are incoming edges of `i` in the formula above). We also support edge
|
|
softmax normalized by source nodes(i.e. :math:`ij` are outgoing edges of
|
|
`i` in the formula). The previous case correspond to softmax in GAT and
|
|
Transformer, and the later case correspond to softmax in Capsule network.
|
|
|
|
Parameters
|
|
----------
|
|
gidx : HeteroGraphIndex
|
|
The graph to perfor edge softmax on.
|
|
logits : torch.Tensor
|
|
The input edge feature
|
|
eids : torch.Tensor or ALL, optional
|
|
Edges on which to apply edge softmax. If ALL, apply edge
|
|
softmax on all edges in the graph. Default: ALL.
|
|
norm_by : str, could be `src` or `dst`
|
|
Normalized by source nodes or destination nodes. Default: `dst`.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
Softmax value
|
|
"""
|
|
pass
|
|
|
|
|
|
def edge_softmax_hetero(gidx, eids, norm_by, *logits):
|
|
r"""Compute edge softmax.
|
|
|
|
For a node :math:`i`, edge softmax is an operation of computing
|
|
|
|
.. math::
|
|
a_{ij} = \frac{\exp(z_{ij})}{\sum_{j\in\mathcal{N}(i)}\exp(z_{ij})}
|
|
|
|
where :math:`z_{ij}` is a signal of edge :math:`j\rightarrow i`, also
|
|
called logits in the context of softmax. :math:`\mathcal{N}(i)` is
|
|
the set of nodes that have an edge to :math:`i`.
|
|
|
|
By default edge softmax is normalized by destination nodes(i.e. :math:`ij`
|
|
are incoming edges of `i` in the formula above). We also support edge
|
|
softmax normalized by source nodes(i.e. :math:`ij` are outgoing edges of
|
|
`i` in the formula). The previous case correspond to softmax in GAT and
|
|
Transformer, and the later case correspond to softmax in Capsule network.
|
|
|
|
Parameters
|
|
----------
|
|
gidx : HeteroGraphIndex
|
|
The graph to perfor edge softmax on.
|
|
eids : dict of tensors
|
|
Each tensor has the edges on which to apply edge softmax for a
|
|
corresponsing relation type.
|
|
logits : tuple of tensors
|
|
The input edge features of different relation types.
|
|
norm_by : str, could be `src` or `dst`
|
|
Normalized by source nodes or destination nodes. Default: `dst`.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
Softmax value
|
|
"""
|
|
pass
|
|
|
|
|
|
def segment_reduce(op, x, offsets):
|
|
"""Segment reduction operator.
|
|
|
|
It aggregates the value tensor along the first dimension by segments.
|
|
The argument ``offsets`` specifies the start offset of each segment (and
|
|
the upper bound of the last segment). Zero-length segments are allowed.
|
|
|
|
.. math::
|
|
y_i = \Phi_{j=\mathrm{offsets}_i}^{\mathrm{offsets}_{i+1}-1} x_j
|
|
|
|
where :math:`\Phi` is the reduce operator.
|
|
|
|
Parameters
|
|
----------
|
|
op : str
|
|
Aggregation method. Can be ``sum``, ``max``, ``min``.
|
|
x : Tensor
|
|
Value to aggregate.
|
|
offsets : Tensor
|
|
The start offsets of segments.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
Aggregated tensor of shape ``(len(offsets) - 1, value.shape[1:])``.
|
|
"""
|
|
pass
|
|
|
|
|
|
def scatter_add(x, idx, m):
|
|
|
|
"""Scatter add (on first dimension) operator.
|
|
|
|
Math: y[idx[i], *] += x[i, *]
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
The input feature.
|
|
idx : Tensor
|
|
The indices array.
|
|
m : int
|
|
The length of output.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The output tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def csrmm(A, A_weights, B, B_weights, num_vtypes):
|
|
"""Compute weighted adjacency matrix multiplication.
|
|
|
|
Notes
|
|
-----
|
|
Both A and B must allow creation of CSR representations, and must be simple graphs
|
|
(i.e. having at most one edge between two nodes).
|
|
|
|
The output unit graph has no format restriction.
|
|
|
|
Parameters
|
|
----------
|
|
A : HeteroGraphIndex
|
|
The unit graph as left operand.
|
|
A_weights : Tensor
|
|
The edge weights of A. Must be a 1D vector.
|
|
B : HeteroGraphIndex
|
|
The unit graph as right operand.
|
|
B_weights : Tensor
|
|
The edge weights of B. Must be a 1D vector.
|
|
num_vtypes : int
|
|
The number of node types of the output graph. Must be either 1 or 2.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The output unit graph.
|
|
Tensor
|
|
The output edge weights.
|
|
"""
|
|
pass
|
|
|
|
|
|
def csrsum(gidxs, weights):
|
|
"""Compute weighted adjacency matrix summation.
|
|
|
|
Notes
|
|
-----
|
|
All unit graphs must allow creation of CSR representations, and must be simple graphs
|
|
(i.e. having at most one edge between two nodes).
|
|
|
|
The output unit graph has no format restriction.
|
|
|
|
Parameters
|
|
----------
|
|
gidxs : list[HeteroGraphIndex]
|
|
The unit graphs.
|
|
weights : list[Tensor]
|
|
The edge weights of each graph. Must be 1D vectors.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The output unit graph.
|
|
Tensor
|
|
The output edge weights.
|
|
"""
|
|
pass
|
|
|
|
|
|
def csrmask(A, A_weights, B):
|
|
"""Retrieve the values in the weighted adjacency matrix of graph :attr:`A` at the
|
|
non-zero positions of graph :attr:`B`'s adjacency matrix.
|
|
|
|
In scipy, this is equivalent to ``A[B != 0]``.
|
|
|
|
Notes
|
|
-----
|
|
Both A and B must allow creation of CSR representations, and must be simple graphs
|
|
(i.e. having at most one edge between two nodes).
|
|
|
|
Parameters
|
|
----------
|
|
A : HeteroGraphIndex
|
|
The unit graph as left operand.
|
|
A_weights : Tensor
|
|
The edge weights of A. Must be a 1D vector.
|
|
B : HeteroGraphIndex
|
|
The unit graph as right operand.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The output tensor.
|
|
"""
|
|
pass
|
|
|
|
|
|
def gather_mm(A, B, idx_a, idx_b):
|
|
r"""Dense Matrix Multiplication interface. It multiplies 2D dense tensor A
|
|
and 3D dense tensor B according to their relation types. A is unsorted and
|
|
the relation type is fetched from idx_b.
|
|
|
|
Parameters
|
|
----------
|
|
A : tensor
|
|
2-D tensor of shape (N, D1)
|
|
B : tensor
|
|
3-D tensor of shape (R, D1, D2)
|
|
idx_a : Tensor, optional
|
|
If specified, must be a 1-D integer tensor of shape (K,).
|
|
idx_b : Tensor, optional
|
|
If specified, must be a 1-D integer tensor of shape (K,).
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The output dense matrix of shape (N, D2)
|
|
"""
|
|
pass
|
|
|
|
|
|
def segment_mm(A, B, seglen_A):
|
|
r"""Dense Matrix Multiplication interface. It multiplies dense tensor A
|
|
and dense tensor B according to relation types. A is sorted and concatenated
|
|
according to relation types.
|
|
|
|
Parameters
|
|
----------
|
|
A : tensor
|
|
2-D tensor of shape (N, D1)
|
|
B : tensor
|
|
3-D tensor of shape (R, D1, D2)
|
|
seglen_A : Tensor
|
|
An integer tensor of shape (R,). Each element is the length of segments
|
|
of input ``A``. The summation of all elements must be equal to N.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The output dense matrix of shape (N, D2)
|
|
"""
|
|
pass
|
|
|
|
|
|
###############################################################################
|
|
# Other interfaces
|
|
# ----------------
|
|
# These are not related to tensors. Some of them are temporary workarounds that
|
|
# should be included in DGL in the future.
|
|
|
|
|
|
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.
|
|
"""
|
|
pass
|
|
|
|
|
|
def attach_grad(tensor):
|
|
"""Attach gradients to the input tensor"""
|
|
pass
|
|
|
|
|
|
def backward(x, head_gradient=None):
|
|
"""Invoke backward computation with an optional head gradient."""
|
|
pass
|
|
|
|
|
|
def grad(x):
|
|
"""Fetches the gradient from the tensor after backward computation."""
|
|
pass
|
|
|
|
|
|
def is_no_grad(x):
|
|
"""Test if the input tensor has gradient"""
|
|
pass
|
|
|
|
|
|
def is_recording():
|
|
"""Test if the execution is recording gradients."""
|
|
pass
|
|
|
|
|
|
class record_grad(object):
|
|
"""Context manager that records the gradients"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, exc_type, exc_value, exc_traceback):
|
|
pass
|
|
|
|
|
|
class no_grad(object):
|
|
"""Context manager that explicitly disables gradient computation"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, exc_type, exc_value, exc_traceback):
|
|
pass
|
|
|
|
|
|
class NodeEmbedding(object):
|
|
"""Sparse node embeddings"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, exc_type, exc_value, exc_traceback):
|
|
pass
|