I’m about to start adding batched sparse QR factorizations with cusolver to PyTorch. What’s the best way to handle a 3D sparse batched tensor (a batch of 2D matrices)? The batched QR factorization and solves also require that the tensors have the same sparsity pattern across the batch. I’d like to propose adding support for 3D sparse batch tensors with the same sparsity pattern using an interface like this.
nBatch = 2; nx =4
i = torch.LongTensor([range(nx), range(nx)])
v = torch.randn(nBatch, nx)
sz = torch.Size([nBatch, nx, nx])
A = torch.sparse.FloatTensor(i, v, sz)

What do you all think? For the sake of getting the batched QR factorizations and solves working, I’m going to pass dense tensors around without using the sparse wrapper. Eventually (ideally soon), I’d like for these functions to take sparse tensors as input.

Here’s the github issue for the factorizations for reference:

  
      github.com/pytorch/pytorch
  
  
    
  
	  
  

  
    
      Batched sparse QR factorizations and solves with cusolver 40
    

    
      
        opened 
        
          
        
        Apr 3, 2017
      
      


      
        
          
          bamos
        
      
    
  



  These should significantly speed up qpth in some cases when sparse matrices can be used. Here is a cusolver reference on...



    feature
    low priority
    module: sparse
    todo
    topic: linear algebra
    triaged