Repository: pytorch
Issue: Slice operation on "ragged" dimension in NestedTensor
Labels: triaged,enhancement,module: nestedtensor

### 🚀 The feature, motivation and pitch

## Motivation

In preproc we often wants to operates over variable-width list, such as token ids in text domain, or sparse features in recommendation domain; one common operation is to slice over each list (e.g. only need first k elements). One way is to use Arrow's List type: 
```python
>>> import torcharrow as ta
>>> id_list = ta.column([[0, 1, 2, 3], [4, 5, 6, 7, 8], [9, 10]])
>>> id_list
0  [0, 1, 2, 3]
1  [4, 5, 6, 7, 8]
2  [9, 10]
dtype: List(int64), length: 3, null_count: 0

>>> id_list.list.slice(stop=3)
0  [0, 1, 2]
1  [4, 5, 6]
2  [9, 10]
dtype: List(Int64(nullable=True)), length: 3, null_count: 0
```

I was thinking nested tensor may also work well for this use case (especially when doing preproc after Tensor collate). But looks like slice is not yet supported on ragged dimension?

```python
>>> import torch
>>> a, b, c = torch.arange(4), torch.arange(5) + 4, torch.arange(2) + 9
>>> id_list = torch.nested_tensor([a, b, c])
>>> id_list
nested_tensor([
  tensor([0, 1, 2, 3]),
  tensor([4, 5, 6, 7, 8]),
  tensor([9, 10])
])
>>> id_list[:, :3]
raceback (most recent call last):
  File "<stdin>", line 1, in <module>
NotImplementedError: Could not run 'aten::slice.Tensor' with arguments from the 'NestedTensorCPU' 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::slice.Tensor' is only available for these backends: [CPU, CUDA, HIP, XLA, MPS, IPU, XPU, HPU, VE, Lazy, Meta, PrivateUse1, PrivateUse2, PrivateUse3, FPGA, ORT, Vulkan, Metal, QuantizedCPU, QuantizedCUDA, QuantizedHIP, QuantizedXLA, QuantizedMPS, QuantizedIPU, QuantizedXPU, QuantizedHPU, QuantizedVE, QuantizedLazy, QuantizedMeta, QuantizedPrivateUse1, QuantizedPrivateUse2, QuantizedPrivateUse3, CustomRNGKeyId, MkldnnCPU, SparseCPU, SparseCUDA, SparseHIP, SparseXLA, SparseMPS, SparseIPU, SparseXPU, SparseHPU, SparseVE, SparseLazy, SparseMeta, SparsePrivateUse1, SparsePrivateUse2, SparsePrivateUse3, SparseCsrCPU, SparseCsrCUDA, BackendSelect, Python, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMeta, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradNestedTensor, Tracer, AutocastCPU, AutocastCUDA, Batched, VmapMode, PythonTLSSnapshot].
......
```

Wondering if there is any plan to support this? Thanks!


### Alternatives

_No response_

### Additional context

Variable width data is often modelled as the flattened value and the offset tensor. For the above (simplified 1D) case, one way is to model it as the following internal representation (which is the Arrow Layout, other layout variations exist, such as use the `lengths`):

```python
values=tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10]),
offsets=tensor([ 0,  4,  9, 11]),

# Logically, represent the following variable-width data: 
#
# 0  [0, 1, 2, 3]
# 1  [4, 5, 6, 7, 8]
# 2  [9, 10]
# dtype: List(int64), length: 3
```

So we kind wants to to a "batched slice" over `values` over the ranges `(0, 3), (4, 7), (9, 11)`. The ranges is kind of like `offsets, offsets + 3` (needs to capped by the end of each list.

General n-D Tensor slice support is more complicated, but the similar idea may still work? 

The request originally posted in the NestedTensor repo: https://github.com/pytorch/nestedtensor/issues/473 . But now realized new feature about NestedTensor should be posted in PyTorch repo. 

Thanks!


cc @cpuhrsch @jbschlosser @bhosmer