Function _csrmask

python/dgl/_sparse_ops.py:860–886

Return the weights of A at the locations identical to the sparsity pattern of B. If a non-zero entry in B does not exist in A, DGL returns 0 for that location instead. Note that the edge weights of the graph must be scalar, i.e. :attr:`A_weights` must be a 1D vector. In sc

(A, A_weights, B)

def _csrmask(A, A_weights, B):
    """Return the weights of A at the locations identical to the sparsity pattern
    of B.

    If a non-zero entry in B does not exist in A, DGL returns 0 for that location
    instead.

    Note that the edge weights of the graph must be scalar, i.e. :attr:`A_weights`
    must be a 1D vector.

    In scipy notation this is identical to ``A[B != 0]``.

    Parameters
    ----------
    A : HeteroGraphIndex
        The input graph index as left operand.
    A_weights : Tensor
        The edge weights of graph A as 1D tensor.
    B : HeteroGraphIndex
        The input graph index as right operand.

    Returns
    -------
    B_weights : Tensor
        The output weights.
    """
    return F.from_dgl_nd(_CAPI_DGLCSRMask(A, F.to_dgl_nd(A_weights), B))


###################################################################################################