torch.sparse.sampled_addmm

torch.sparse.sampled_addmm(input, mat1, mat2, *, beta=1.0, alpha=1.0, out=None) → Tensor

Выполняет матричное умножение плотных матриц mat1 и mat2 в позициях, определённых структурой разреженности input. К итоговому результату добавляется матрица input.

Математически это выполняет следующую операцию:

$$\text{out} = \alpha\ (\text{mat1} \mathbin{@} \text{mat2})*\text{spy}(\text{input}) + \beta\ \text{input}$$

где $\text{spy}(\text{input})$ — матрица разреженности input, alpha и beta — коэффициенты масштабирования. $\text{spy}(\text{input})$ имеет значение 1 в позициях, где input имеет ненулевые значения, и 0 в противном случае.

Примечание

input должна быть разреженной тензором CSR. mat1 и mat2 должны быть плотными тензорами. Эта функция реализована только для тензоров на устройствах CUDA.

Параметры:

• input (Tensor) – разреженная матрица CSR формы (m, n) для добавления и вычисления выборочного матричного умножения
• mat1 (Tensor) – плотная матрица формы (m, k) для умножения
• mat2 (Tensor) – плотная матрица формы (k, n) для умножения

Ключевые аргументы:

• beta (Число, необязательно) – множитель для input ($\beta$)
• alpha (Число, необязательно) – множитель для $mat1 @ mat2$ ($\alpha$)
• out (Tensor, необязательно) – тензор результата. Игнорируется, если None. По умолчанию: None.

Примеры:

>>> 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)