How to sort rows of a sparse matrix / make sure sparse matrix-vector multiplication is working optimally

[Ebanflo42]

I get the following warning while running a simulation with JAX that involves sparse matrix - dense vector multiplication:

coo.py:462: CuSparseEfficiencyWarning: coo_matvec GPU lowering requires matrices with sorted rows or sorted cols. To sort the rows in your matrix, use e.g. mat = mat._sort_rows(). Falling back to the default implementation.
  warnings.warn("coo_matvec GPU lowering requires matrices with sorted rows or sorted cols. "

The method _sort_rows() does not seem to exist on COO matrices, or it is not exported somehow. What exactly does it mean to sort the rows in this case and is there a way to do it quickly with the existing API?

[jakevdp]

Thanks for the question - that should read _sort_indices() instead of _sort_rows(). I'll update the warning message. (Edit: fixed in #15243)

[Ebanflo42]

Thanks for the super quick response!