matmul

iskra.sparse.matmul(a: SparseTensor, b: SparseTensor) → SparseTensor[SOURCE]

iskra.sparse.matmul(a: Tensor, b: SparseTensor) → SparseTensor

iskra.sparse.matmul(a: SparseTensor, b: Tensor) → SparseTensor

iskra.sparse.matmul(a: Tensor, b: Tensor) → Tensor

Matrix-matrix product with support for sparse COO tensors.

Unlike PyTorch’s default, differentiating through this function produces sparse gradients through the sparse inputs. SparseTensor.__matmul__ wraps around this function, meaning that, unlike PyTorch, a @ b is not a foot-gun in iskra.

Warns:

• In case of a sparse-dense product, the sparse matrix is internally

• converted to the CSR format. This limitation stems from PyTorch’s implementation

• of sparse COO matrices, which would produce dense gradients by default.

• You might want to consider converting it to CSR yourself before calling matmul.

Parameters:

• a (SparseTensor | Tensor) – First tensor to be multiplied.

• b (SparseTensor | Tensor) – Second tensor to be multiplied.

Returns:

​ (SparseTensor | Tensor) – Product a @ b.