dgl.sparse.sp_broadcast_v

dgl.sparse.sp_broadcast_v(A: dgl.sparse.sparse_matrix.SparseMatrix, v: torch.Tensor, op: str)dgl.sparse.sparse_matrix.SparseMatrix[source]

Broadcast operator for sparse matrix and vector.

v is broadcasted to the shape of A and then the operator is applied on the non-zero values of A.

There are two cases regarding the shape of v:

1. v is a vector of shape (1, A.shape[1]) or (A.shape[1]). In this case, v is broadcasted on the row dimension of A.

2. v is a vector of shape (A.shape[0], 1). In this case, v is broadcasted on the column dimension of A.

If A.val takes shape (nnz, D), then v will be broadcasted on the D dimension.

Parameters
  • A (SparseMatrix) – Sparse matrix

  • v (torch.Tensor) – Vector

  • op (str) – Operator in [“add”, “sub”, “mul”, “truediv”]

Returns

Sparse matrix

Return type

SparseMatrix

Examples

>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([10, 20, 30])
>>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
>>> v = torch.tensor([1, 2, 3, 4])
>>> dglsp.sp_broadcast_v(A, v, "add")
SparseMatrix(indices=tensor([[1, 0, 2],
                             [0, 3, 2]]),
             values=tensor([11, 24, 33]),
             shape=(3, 4), nnz=3)
>>> v = torch.tensor([1, 2, 3]).view(-1, 1)
>>> dglsp.sp_broadcast_v(A, v, "add")
SparseMatrix(indices=tensor([[1, 0, 2],
                             [0, 3, 2]]),
             values=tensor([12, 21, 33]),
             shape=(3, 4), nnz=3)
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([[10, 20], [30, 40], [50, 60]])
>>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
>>> v = torch.tensor([1, 2, 3]).view(-1, 1)
>>> dglsp.sp_broadcast_v(A, v, "sub")
SparseMatrix(indices=tensor([[1, 0, 2],
                             [0, 3, 2]]),
             values=tensor([[ 8, 18],
                            [29, 39],
                            [47, 57]]),
             shape=(3, 4), nnz=3, val_size=(2,))