Differentiable sparse linear solver with cupy backend - "unsupported tensor layout: Sparse" in gradcheck

username: tvercaut
cooked: <p>I am trying to create a differentiable sparse linear solver for pytorch using cupy as a backend. I am mimicking the approach discussed <a href="https://blog.flaport.net/solving-sparse-linear-systems-in-pytorch.html" rel="noopener nofollow ugc">here</a> but relying on cupy rather than scipy to take advantage of the GPU. Initial tests seem to run fine on <a href="https://colab.research.google.com/gist/tvercaut/66fe2fc73b56592cbbf48d86ba0fa3e0/sparseinverse-torchcupy-interop-test.ipynb" rel="noopener nofollow ugc">colab</a> except that I get an error when using <code>gradcheck</code>. Before digging further I would like to know if I am missing something.</p>
<p>Here is my code:</p>
<pre><code class="lang-python">import torch
print('PyTorch version:',torch.__version__)

torchdevice = torch.device('cpu')
if torch.cuda.is_available():
  torchdevice = torch.device('cuda')
  print('Default GPU is ' + torch.cuda.get_device_name(torch.device('cuda')))
print('Running on ' + str(torchdevice))

# PyTorch also doesn't support any sparse solver
# -&gt; workaround by using cupy with a custom backward pass
# https://docs.cupy.dev/en/stable/reference/scipy_sparse_linalg.html#solving-linear-problems
# https://blog.flaport.net/solving-sparse-linear-systems-in-pytorch.html
# Note that the default cupy version clashes with the default pytorch version
# in colab -&gt; use cupy 10
!pip install cupy-cuda111==10.0.0
import cupy as cp
import cupyx.scipy.sparse.linalg
print('CuPy version:',cp.__version__)
print('Running on ',cp.array([1]).data.device)

# Convenience function to map a torch COO tensor in a cupy one
def coo_torch2cupy(A):
  A = A.data.coalesce()
  Avals_cp = cp.asarray(A.values())
  Aidx_cp = cp.asarray(A.indices())
  return cp.sparse.coo_matrix((Avals_cp, Aidx_cp))

# Custom PyTorch sparse solver exploiting a CuPy backend
# See https://blog.flaport.net/solving-sparse-linear-systems-in-pytorch.html
class SparseSolve(torch.autograd.Function):
  @staticmethod
  def forward(ctx, A, b):
    # Sanity check
    if A.ndim != 2 or (A.shape[0] != A.shape[1]):
      raise ValueError("A should be a square 2D matrix.")
    # Transfer data to CuPy
    A_cp = coo_torch2cupy(A)
    b_cp = cp.asarray(b.data)
    # Solver the sparse system
    if (b.ndim == 1) or (A.shape[1] == 1):
      # cp.sparse.linalg.spsolve only works if b is a vector but is fully on GPU
      x_cp = cp.sparse.linalg.spsolve(A_cp, b_cp)
    else:
      # Make use of a factorisation (only the solver is then on the GPU)
      factorisedsolver = cp.sparse.linalg.factorized(A_cp)
      x_cp = factorisedsolver(b_cp)
    # Transfer (dense) result back to PyTorch
    x = torch.as_tensor(x_cp, device=torchdevice)
    # Not sure if the following is needed / helpful
    if A.requires_grad or b.requires_grad:
      x.requires_grad = True
    # Save context for backward pass
    ctx.save_for_backward(A, b, x)
    return x

  @staticmethod
  def backward(ctx, grad):
    # Recover context
    A, b, x = ctx.saved_tensors
    # Compute gradient with respect to b
    gradb = SparseSolve.apply(A.t(), grad)
    # The gradient with respect to the (dense) matrix A would be something like
    # -gradb @ x.T but we are only interested in the gradient with respect to
    # the (non-zero) values of A
    gradAidx = A.indices()
    mgradbselect = -gradb.index_select(0,gradAidx[0,:])
    xselect = x.index_select(0,gradAidx[1,:])
    mgbx = mgradbselect * xselect
    if x.dim() == 1:
      gradAvals = mgbx
    else:
      gradAvals = torch.sum( mgbx, dim=1 )
    gradAs = torch.sparse_coo_tensor(gradAidx, gradAvals, A.shape)
    return gradAs, gradb

sparsesolve = SparseSolve.apply

# Test matrix-vector solver
Aref = torch.randn(3,3, dtype=torch.float64, requires_grad=False, device=torchdevice).to_sparse()
Aref.requires_grad=False
bref = torch.randn(3, dtype=torch.float64, requires_grad=False, device=torchdevice)

A = Aref.detach().clone().requires_grad_(True)
b = bref.detach().clone().requires_grad_(True)

# Solve
x = sparsesolve(A,b)

# random scalar function to mimick a loss
loss = x.sum()
loss.backward()

print('x',x)
with torch.no_grad(): print('allclose:',torch.allclose(A @ x, b))
print('A.grad',A.grad)
print('b.grad',b.grad)

# Compare with dense op
A = Aref.detach().clone().to_dense().requires_grad_(True)
b = bref.detach().clone().requires_grad_(True)
x = torch.linalg.solve(A,b)
loss = x.sum()
loss.backward()
print('x',x)
with torch.no_grad(): print('allclose:',torch.allclose(A @ x, b))
print('A.grad',A.grad)
print('b.grad',b.grad)

# Test matrix-matrix solver
Aref = torch.randn(3,3, dtype=torch.float64, requires_grad=False, device=torchdevice).to_sparse()
Aref.requires_grad=False
bref = torch.randn(3,2, dtype=torch.float64, requires_grad=False, device=torchdevice)

A = Aref.detach().clone().requires_grad_(True)
b = bref.detach().clone().requires_grad_(True)

# Solve
x = sparsesolve(A,b)

# random scalar function to mimick a loss
loss = x.sum()
loss.backward()

print('x',x)
with torch.no_grad(): print('allclose:',torch.allclose(A @ x, b))
print('A.grad',A.grad)
print('b.grad',b.grad)

# Compare with dense op
A = Aref.detach().clone().to_dense().requires_grad_(True)
b = bref.detach().clone().requires_grad_(True)
x = torch.linalg.solve(A,b)
loss = x.sum()
loss.backward()
print('x',x)
print('allclose:',torch.allclose(A @ x, b))
print('A.grad',A.grad)
print('b.grad',b.grad)

# Now try some gradcheck
A = torch.randn(3,3, dtype=torch.float64, device=torchdevice).to_sparse()
A.requires_grad=True
b = torch.randn(3, dtype=torch.float64, requires_grad=True, device=torchdevice)
torch.autograd.gradcheck(sparsesolve, [A, b], check_sparse_nnz=True, raise_exception=True)
</code></pre>

username: tvercaut
cooked: <p>Not sure why this is needed but adding <code>check_undefined_grad=False</code> when calling <code>gradcheck</code> solves the issue. This was similarly noted here: <a href="https://github.com/flaport/torch_sparse_solve/issues/7" class="inline-onebox" rel="noopener nofollow ugc">Occur bugs when doing backward process. · Issue #7 · flaport/torch_sparse_solve · GitHub</a></p>

username: tvercaut
cooked: <p>As a follow-up, there is ongoing effort to integrate a sparse solver natively in torch so this workaround shouldn’t be necessary anymore in the future:</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/69538">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/69538" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/69538" target="_blank" rel="noopener nofollow ugc">cuSOLVER backend for Sparse CSR direct linear solvers</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2021-12-07" data-time="17:01:58" data-timezone="UTC">05:01PM - 07 Dec 21 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/IvanYashchuk" target="_blank" rel="noopener nofollow ugc">
          <img alt="IvanYashchuk" src="https://avatars.githubusercontent.com/u/19621411?v=4" class="onebox-avatar-inline" width="20" height="20">
          IvanYashchuk
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: sparse
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: cuda
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: linear algebra
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🚀 Feature

cuSolverSP (part of cuSOLVER) provides linear solver, least-squa<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">re solver routines for matrices in CSR format. We can use it as a backend for `torch.linalg.solve` and `torch.linalg.lstsq`.

We are working on improving coverage of Sparse CSR matrices in PyTorch. Solving systems of linear equations is one of the fundamental operations in computational linear algebra.

From cuSOLVER docs:
&gt; The high-level API is designed for ease-of-use, so it allocates any required memory under the hood automatically. If the host or GPU system memory is not enough, an error is returned.

Unfortunately, it means these functions do not allow passing externally allocated buffers, and paths with inputs on GPU are not implemented (there remarks for relevant functions "only CPU (Host) path is provided").

- [ ] `torch.linalg.solve`
    Use [`cusolverSp&lt;t&gt;csrlsvlu`](https://docs.nvidia.com/cuda/cusolver/index.html#cusolver-lt-t-gt-csrlsvlu). 

- [ ] `torch.linalg.lstsq`
    Use [`cusolverSp&lt;t&gt;csrlsqvqr`](https://docs.nvidia.com/cuda/cusolver/index.html#cusolver-lt-t-gt-csrlsqvqr)



cc @nikitaved @pearu @cpuhrsch @IvanYashchuk @ngimel @jianyuh @mruberry @walterddr @xwang233 @Lezcano</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>