dgl.sparse.SparseMatrix.softmaxο
- SparseMatrix.softmax(dim: int = 1) SparseMatrix ο
Applies softmax to the non-zero elements of the sparse matrix on the dimension :attr:
dim
. dim = 0 or 1 indicates column-wise or row-wise softmax respectively.If
input.val
takes shape(nnz, D)
, then the output matrixoutput
andoutput.val
take the same shape asinput
andinput.val
.output.val[:, i]
is calculated based oninput.val[:, i]
.- Parameters:
input (SparseMatrix) β The input sparse matrix
- Returns:
The output sparse matrix
- Return type:
Examples
Case1: row-wise softmax on matrix with values of shape (nnz)
>>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]]) >>> val = torch.tensor([0., 1., 2., 3.]) >>> A = dglsp.spmatrix(indices, val) >>> dglsp.softmax(A) SparseMatrix(indices=tensor([[0, 0, 1, 2], [1, 2, 2, 0]]), values=tensor([0.2689, 0.7311, 1.0000, 1.0000]), shape=(3, 3), nnz=4)
Case2: row-wise softmax on matrix with values of shape (nnz, D)
>>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]]) >>> val = torch.tensor([[0., 7.], [1., 3.], [2., 2.], [3., 1.]]) >>> A = dglsp.spmatrix(indices, val) >>> dglsp.softmax(A) SparseMatrix(indices=tensor([[0, 0, 1, 2], [1, 2, 2, 0]]), values=tensor([[0.2689, 0.9820], [0.7311, 0.0180], [1.0000, 1.0000], [1.0000, 1.0000]]), shape=(3, 3), nnz=4, val_size=(2,))
Case3: column-wise softmax on matrix with values of shape (nnz)
>>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]]) >>> val = torch.tensor([0., 1., 2., 3.]) >>> A = dglsp.spmatrix(indices, val) >>> dglsp.softmax(A, 0) SparseMatrix(indices=tensor([[0, 0, 1, 2], [1, 2, 2, 0]]), values=tensor([1.0000, 0.2689, 0.7311, 1.0000]), shape=(3, 3), nnz=4)