I have a giant sparse matrix S. Each column has 1 entry whose value is 1. The rest of the values are zeros. I frequently have to multiply S with another matrix A (S@A). It’s really just a “selector”, so it’s trivial to implement in numpy.

Does torch have some APIs that can help me with this? The result of S@A will be used extensively in gpu (such as further multiplication, solving linear systems, etc). so I’d like to implement S@A in GPU too.