Inconsistent DirectSolver interface.

[anadodik]

Hi everyone! I've noticed a small inconsistency in the sparse solver interface, namely when it comes to using it with a 2D tensor where the last dimension is 1.

Given a 3x3 lhs sparse matrix and a 3x2 rhs dense matrix:

a = torch.sparse_coo_tensor([[0, 1, 2], [0, 1, 2]], [1.0, 1.0, 1.0], device="cuda").to_sparse_csr()
b = torch.randn([3, 2], device="cuda")

we can set up a solver for with:

nvsparse.DirectSolver(a, b.mT.contiguous().mT)

This works fine. However, for a 3x1 rhs:

b = torch.randn([3,1], device="cuda")

setting up the solver in the same way breaks with:

TypeError("The RHS tensor([[1.],\n        [1.],\n        [1.]], device='cuda:0') must be a matrix or vector with col-major layout, and for implicitly-batched RHS (N-D >= 3), each matrix sample must have col-major layout (the second dimension from the end must have unit stride.")

It would be great not to have to special-case this.

Best,
Ana