torch.sparse.mm

torch.sparse.mm()

Выполняет матричное умножение разреженной матрицы mat1 и (разреженной или с шагом) матрицы mat2. Аналогично torch.mm(), если mat1 является тензором $(n \times m)$, mat2 является тензором $(m \times p)$, out будет тензором $(n \times p)$. Когда mat1 является тензором COO, он должен иметь sparse_dim = 2. При использовании COO-тензоров в качестве входных данных, эта функция также поддерживает обратное вычисление для обоих входных данных.

Поддерживает форматы хранения CSR и COO.

Примечание

Эта функция не поддерживает вычисление производных по отношению к CSR-матрицам.

Эта функция также дополнительно принимает необязательный reduce аргумент, который позволяет указать необязательную операцию сокращения, математически выполняющую следующую операцию:

$$z_{ij} = \bigoplus_{k = 0}^{K - 1} x_{ik} y_{kj}$$

где $\bigoplus$ определяет операцию сокращения. reduce реализована только для формата хранения CSR на устройстве CPU.

Параметры

• mat1 (Tensor) – первая разреженная матрица, подлежащая умножению
• mat2 (Tensor) – вторая матрица, подлежащая умножению, которая может быть разреженной или плотной
• reduce (str, optional) – операция сокращения для применения к не уникальным индексам ("sum", "mean", "amax", "amin"). По умолчанию "sum".

Форма:

Формат выходного тензора этой функции следующий: - разреженная х разреженная -> разреженная - разреженная х плотная -> плотная

Пример:

>>> a = torch.tensor([[1., 0, 2], [0, 3, 0]]).to_sparse().requires_grad_()
>>> a
tensor(indices=tensor([[0, 0, 1],
                       [0, 2, 1]]),
       values=tensor([1., 2., 3.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo, requires_grad=True)
>>> b = torch.tensor([[0, 1.], [2, 0], [0, 0]], requires_grad=True)
>>> b
tensor([[0., 1.],
        [2., 0.],
        [0., 0.]], requires_grad=True)
>>> y = torch.sparse.mm(a, b)
>>> y
tensor([[0., 1.],
        [6., 0.]], grad_fn=<SparseAddmmBackward0>)
>>> y.sum().backward()
>>> a.grad
tensor(indices=tensor([[0, 0, 1],
                       [0, 2, 1]]),
       values=tensor([1., 0., 2.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
>>> c = a.detach().to_sparse_csr()
>>> c
tensor(crow_indices=tensor([0, 2, 3]),
       col_indices=tensor([0, 2, 1]),
       values=tensor([1., 2., 3.]), size=(2, 3), nnz=3,
       layout=torch.sparse_csr)
>>> y1 = torch.sparse.mm(c, b, 'sum')
>>> y1
tensor([[0., 1.],
        [6., 0.]], grad_fn=<SparseMmReduceImplBackward0>)
>>> y2 = torch.sparse.mm(c, b, 'max')
>>> y2
tensor([[0., 1.],
        [6., 0.]], grad_fn=<SparseMmReduceImplBackward0>)