Description
🚀 Feature
Currently, sparse tensors only support real floating point tensors. Complex support should be added for completion.
Motivation
I'm trying to use complex tensors for compressed sensing and am running out of GPU memory (4GiB); even though my tensors are sparse (~1%) I cannot use optimized routines, since PyTorch does not currently support them. I have resorted to converting them to scipy sparse matrices and going without autograd.
Pitch
torch.sparse.CFloatTensor, torch.sparse.CDoubleTensor added as complex extensions of existing real data types.
Status
Complex support was added in this PR #57125 for CPU/CUDA.
- complex support (torch.cfloat and torch.cdouble) to torch.sparse_coo_tensor constructors
- complex support to coalesce function
- complex support to to_dense function
- complex support to to_sparse function
- complex support to sparse_add function
- complex support to addmm function
- unit tests
Note: That PR contains only complex support for torch.sparse_coo_tensor forward/backward function and the related ops used with this function (coalesce, to_dense, to_sparse, and sparse_add, addmm). The following PRs should cover other sparse operations to have a more complex sparse support, specifically related with the use of specific APIs for accelerated linear algebra.
CSR sparse tensor: ops not implemented yet:
- transpose, complex conjugate, torch.pow, torch.exp, torch.block_diag, torch.diagflat, torch.eye
- matmult ops: bmm, sparse@sparse
- Add complex support for CSR Sparse Tensors
cc @ezyang @anjali411 @dylanbespalko @mruberry @aocsa @nikitaved @pearu