"""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