spdiagmm

hypercoil.functional.sparse.spdiagmm(lhs: Tensor | BCOO, rhs: Tensor | BCOO, lhs_diag: bool = False) → BCOO

Matrix multiplication of a top-k format sparse matrix with a diagonal matrix. Returns a sparse matrix in the top-k format.

The diagonal matrix should be formatted as a vector or batch of vectors whose final dimension is equal to the matching inner dimension of the top-k sparse matrix. All other dimensions should be broadcastable.

A diagonal matrix that is in matrix form can be converted into a vector by calling a diagonal function that is appropriately vmapped.