sparse: Sparse matrix operations#

Utilities for operations on various BCOO sparse matrix formats.

Top-k BCOO format#

One common BCOO sparse format that is useful in many applications, such as connectopic mapping, is the top-k format. Sparse matrices in top-k format have the following properties:

  • Each row has no more than k non-zero entries.

  • The indices of nonzero entries are shared across all batch elements.

  • The indexing tensor has shape (..., n_rows, k, 1) where ... indicates a number of leading singleton dimensions equal to the number of channel_dims + 1. (n_rows can be substituted for a singleton dimension as well, in which case all rows have the same nonzero indices.)

  • The data tensor has shape (batch_size, *channel_dims, n_rows, k).