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 matricesmat1 andmat2 at the locationsspecified by the sparsity pattern ofinput. The matrixinput is added to the final result.

Mathematically this performs the following operation:

$$\text{out}=\alpha (\text{mat1}\mathrm{@}\text{mat2})\ast \text{spy}(\text{input})+\beta \text{ input}$$out=α (mat1@mat2)∗spy(input)+β input

where$\text{spy}(\text{input})$spy(input) is the sparsity pattern matrix ofinput,alphaandbeta are the scaling factors.$\text{spy}(\text{input})$spy(input) 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

• input (Tensor) – a sparse CSR matrix of shape(m, n) to be added and used to computethe sampled matrix multiplication

• mat1 (Tensor) – a dense matrix of shape(m, k) to be multiplied

• mat2 (Tensor) – a dense matrix of shape(k, n) to be multiplied

Keyword Arguments

• beta (Number,optional) – multiplier forinput ($\beta$β)

• alpha (Number,optional) – multiplier for$mat1\mathrm{@}mat2$mat1@mat2 ($\alpha$α)

• out (Tensor,optional) – output tensor. Ignored ifNone. 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)