torch.sparse.sampled_addmm¶
- torch.sparse.sampled_addmm(input, mat1, mat2, *, beta=1., alpha=1., out=None) Tensor ¶
Performs a matrix multiplication of the dense matrices
mat1
andmat2
at the locations specified by the sparsity pattern ofinput
. The matrixinput
is added to the final result.Mathematically this performs the following operation:
where is the sparsity pattern matrix of
input
,alpha
andbeta
are the scaling factors. has value 1 at the positions whereinput
has non-zero values, and 0 elsewhere.Note
input
must be a sparse CSR tensor.mat1
andmat2
must be dense tensors.- Parameters
- Keyword Arguments
beta (Number, optional) – multiplier for
input
()alpha (Number, optional) – multiplier for ()
out (Tensor, optional) – output tensor. Ignored if None. Default: None.
Examples:
>>> input = torch.eye(3, device='cuda').to_sparse_csr() >>> mat1 = torch.randn(3, 5, device='cuda') >>> mat2 = torch.randn(5, 3, device='cuda') >>> torch.sparse.sampled_addmm(input, mat1, mat2) tensor(crow_indices=tensor([0, 1, 2, 3]), col_indices=tensor([0, 1, 2]), values=tensor([ 0.2847, -0.7805, -0.1900]), device='cuda:0', size=(3, 3), nnz=3, layout=torch.sparse_csr) >>> torch.sparse.sampled_addmm(input, mat1, mat2).to_dense() tensor([[ 0.2847, 0.0000, 0.0000], [ 0.0000, -0.7805, 0.0000], [ 0.0000, 0.0000, -0.1900]], device='cuda:0') >>> torch.sparse.sampled_addmm(input, mat1, mat2, beta=0.5, alpha=0.5) tensor(crow_indices=tensor([0, 1, 2, 3]), col_indices=tensor([0, 1, 2]), values=tensor([ 0.1423, -0.3903, -0.0950]), device='cuda:0', size=(3, 3), nnz=3, layout=torch.sparse_csr)