spspmm

hypercoil.functional.sparse.spspmm(lhs: BCOO, rhs: BCOO, indices: Tensor | None = None, n_blocks: int = 1)

Sparse-sparse matrix multiplication of top-k format sparse matrices.

This function is a wrapper around the JAX sparse matrix multiplication function bcoo_dot_general. It is a convenience function that also provides the option to separate the matrix multiplication into blocks for serialisation and reduced memory usage. It additionally provides the option to use a pre-computed index tensor and return a top-k format sparse matrix containing the results of the matrix multiplication at only the specified indices.

Warning

Note that this implementation of the matrix multiplication operation returns \(A B^\intercal\) for LHS \(A\) and RHS \(B\).

Parameters:

lhsTopKTensor

Left-hand side sparse matrix in the top-k format.

rhsTopKTensor

Right-hand side sparse matrix in the top-k format.

indicesTensor or None (default: None)

Indices of the matrix product to return. If not specified, the entire matrix is returned.

n_blocksint (default: 1)

Number of blocks to split the matrix multiplication into for serialisation. If set to 1, the matrix multiplication is performed directly.

Returns:

TopKTensor or Tensor

Result of the matrix multiplication.