Sparse Tensor implementation (sparsify last dimension only)

richardbaxter: <p>I would like to be able to sparsify a single (last) dimension of an N-dimensional dense tensor, as I need to be able to manipulate the sparse tensor using ordinary pytorch operations. An example of an ordinary operation I need to be able to perform to the tensor is <code>normaliseTensor</code> (<code>min</code>/<code>max</code> are unavailable);</p>
<pre><code class="lang-auto">def normaliseTensor(tensor):
    min_vals, _ = torch.min(tensor, dim=-1, keepdim=True)
    max_vals, _ = torch.max(tensor, dim=-1, keepdim=True)
    epsilon = 1e-8  # Small epsilon value
    tensor = (tensor - min_vals) / (max_vals - min_vals + epsilon)
    return tensor
</code></pre>
<p><code>min(dim)</code>/<code>max(dim)</code> could be easily implemented using a hybrid dense/sparse tensor (where only the last dimension is sparsified).</p>
<p>Ideally the user should be able to select a dimension along which a tensor is sparsified (e.g. the last dimension of an N-dimensional dense tensor) and all the sparse matrix operations/optimisations should work under the hood with no change to the pytorch interface (they should look like dense tensors in pytorch) - hardware permitting.</p>
<p>Perhaps there is a way the existing pytorch sparse tensor types can be manipulated to emulate this functionality?</p>

richardbaxter: <p><strong>A. batched sparse tensors</strong></p>
<p>The problem could be solved with a batched sparse tensor implementation (a hybrid tensor where the first dimension is dense and the other(s) are sparse);</p>
<ol>
<li>To multiply a dense tensor (of N-1 dim) with a batched sparse representation (2 dim; 1 dense, 1 sparse) of the underlying hybrid tensor (N dim; N-1 dense, 1 sparse) it could be temporarily reshaped to the sparse batch size (1 dim).</li>
<li>The hybrid dense-sparse tensor could be reshaped back to its original state (N-1 dim) from its batched sparse representation (2 dim; 1 dense, 1 sparse) after the sparse dimension operations have been performed and the sparse dimension has been removed.</li>
</ol>
<p>It appears that batched sparse tensors have never been implemented in pytorch;</p><aside class="quote" data-post="1" data-topic="58938">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/e5b9ba/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/creating-a-batch-of-sparse-tensors/58938">Creating A Batch of Sparse Tensors</a> 
  </div>
  <blockquote>
    Let say I want to create a 3-dimensional tensor, where the first dimension is the batch size, 
a[batch_index][sparse_dim1][sparse_dim2].

Thus I need a hybrid sparse tensor representation. The first dimension is a dense dimension. The second and the third dimensions are sparse dimensions. 
However, from the <a href="https://pytorch.org/docs/stable/sparse.html" rel="noopener nofollow ugc">documentation</a>, it seems that the dense dimension must be after the sparse dimension? Then how I represent the hybrid tensor?
  </blockquote>
</aside>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/75721">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/75721" 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 fill-rule="evenodd" 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/75721" target="_blank" rel="noopener nofollow ugc">Support batch indexing with sparse tensors with torch.sparse</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2022-04-13" data-time="03:04:41" data-timezone="UTC">03:04AM - 13 Apr 22 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/BichenWuUCB" target="_blank" rel="noopener nofollow ugc">
          <img alt="BichenWuUCB" src="https://avatars.githubusercontent.com/u/8868254?v=4" class="onebox-avatar-inline" width="20" height="20">
          BichenWuUCB
        </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;">
          triaged
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">### 🚀 The feature, motivation and pitch

I am working on a problem that requires<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden"> looking up sparse tensor values based on a batch of indices. The problem can be abstracted by the following example:

```
import torch
a = torch.sparse_coo_tensor(indices=[[0, 2, 3, 6]], values=[0, 1, 2, 3], size=(10000000,))
index = [1, 2, 3]
print(a[index])
```
If `a` is a dense tensor (vector), then given a list of indices, `a[index]` will return a tensor with the i-th element as `a[index[i]]`. However, on sparse tensors, this operator is not supported, and will give the following error:
```
NotImplementedError: Could not run 'aten::index.Tensor' with arguments from the 'SparseCPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::index.Tensor' is only available for these backends: [CPU, CUDA, QuantizedCPU, BackendSelect, Python, Named, Conjugate, Negative, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradLazy, AutogradXPU, AutogradMLC, AutogradHPU, AutogradNestedTensor, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, Tracer, UNKNOWN_TENSOR_TYPE_ID, Autocast, Batched, VmapMode].

CPU: registered at aten/src/ATen/RegisterCPU.cpp:18433 [kernel]
CUDA: registered at aten/src/ATen/RegisterCUDA.cpp:26496 [kernel]
QuantizedCPU: registered at aten/src/ATen/RegisterQuantizedCPU.cpp:1068 [kernel]
BackendSelect: fallthrough registered at ../aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
Python: registered at ../aten/src/ATen/core/PythonFallbackKernel.cpp:47 [backend fallback]
Named: registered at ../aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]
Conjugate: registered at ../aten/src/ATen/ConjugateFallback.cpp:18 [backend fallback]
Negative: registered at ../aten/src/ATen/native/NegateFallback.cpp:18 [backend fallback]
ADInplaceOrView: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:64 [backend fallback]
AutogradOther: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradCPU: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradCUDA: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradXLA: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradLazy: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradXPU: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradMLC: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradHPU: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradNestedTensor: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradPrivateUse1: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradPrivateUse2: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
AutogradPrivateUse3: registered at ../torch/csrc/autograd/generated/VariableType_1.cpp:9548 [autograd kernel]
Tracer: registered at ../torch/csrc/autograd/generated/TraceType_1.cpp:10664 [kernel]
UNKNOWN_TENSOR_TYPE_ID: fallthrough registered at ../aten/src/ATen/autocast_mode.cpp:466 [backend fallback]
Autocast: fallthrough registered at ../aten/src/ATen/autocast_mode.cpp:305 [backend fallback]
Batched: registered at ../aten/src/ATen/BatchingRegistrations.cpp:1016 [backend fallback]
VmapMode: fallthrough registered at ../aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]
```

To by-pass this, I will have to convert `a` to a dense tensor as 
```
print(a.to_dense()[index])
```
However, this is not memory-friendly, especially if `a` is very sparse and the the full size of `a` is huge.

So I wonder there exists a way to batch-index sparse tensors -- the same way as dense tensors, or it is possible to support this in near future.

Thanks!

### Alternatives

_No response_

### Additional context

_No response_

cc @nikitaved @pearu @cpuhrsch</span></p>
  </div>

  </article>

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

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

<p><strong>B. hybrid dense-sparse tensor implementations</strong></p>
<p>A possible implementation of hybrid dense-sparse tensor would be a dense tensor that contains an integer pointer to a sparse tensor (each sparse tensor contains elements with value/index as in the current pytorch sparse tensor implementation).</p>
<p>When the user requests a particular subset of the tensor (i.e. any operation that requires indexing dense dimensions/batch samples of the tensor e.g. t[3] or torch.min(t, dim=1)), it performs an effective gather (single or multi-index lookup) operation based on the indices to gather the relevant sparse tensors (elements).</p>
<p>Are there specific developers working on (hybrid) sparse tensor implementations?</p>