Repository: pytorch
Issue: [feature request] Varlen indexing function for lookup and concat of varlen BPE tokens from a tensor vocab (i.e. `detokenize(...)` and arrays of strings)
Labels: triaged,module: nestedtensor

### 🚀 The feature, motivation and pitch

In BPE decoding we often have a tensor of token ids and a vocab storing an array of strings where elements may have a variable length. Decoding in the basic way corresponds to looping over the ids, retrieving the token offset and token size and then indexing into the vocab and concatting. In practice, this is useful for embedding BPE detokenizer into the hermetic ONNX model and maybe slightly speeding up BPE detokenization (e.g. for ASR).

If we have NJT representing an array of strings, the hack below is simply indexing + flattening. Maybe this can give an impulse to representing strings / arrays of strings in bare PyTorch (maybe as NJT) and how these can get lowered to ONNX / other simpler opsets if they do not support arrayss of strings. Sometimes it might be god to expose publicly the lower-level NJT-like working functions (accepting/filling offsets/lengths). 

I propose to add some function for doing this in a loopless way, this function already exists in ONNX as https://github.com/onnx/onnx/blob/main/docs/Operators-ml.md#aionnxmlarrayfeatureextractor (although not sure if this supports string-valued tensors cc @cpuhrsch @jbschlosser @bhosmer @drisspg @soulitzer @davidberard98 @YuqingJ @justinchuby, but it appears so )

This is also very related to handling/encoding of varlen tensors like nested/jagged tensors. Ideally, it could be represented as some function indexing into a NJT which could present an array of strings. Especially if we want to do it in a batched way, accepting a list of token id arrays as input and producing an array of varlen byte strings as output.

Even if such function higher-level function is not added, some lower level / "plumbing" varlen+concat functions might be useful which would help constructing index arrays: batched varlen arange, or batched slice.

This is related to:
- Standardized ways of representing/storing strings in tensor:
  - https://github.com/pytorch/pytorch/issues/101699
- repeat_interleave which can repeat with different lengths and concat
- batched arange (where start/end vary across the batch and even `end - start` can vary)
- batched varlen slicing + cat
- nested/jagged tensors (because both strings and tokens and even unicode characters are typically varlen)

I implemented it with crazy hacks using `repeat_interleave`, `cumsum` and `gather` (in form of tensor indexing). But not sure if this gets exported well to ONNX. Hopefully this works now also for legacy torch.onnx.export:
- https://github.com/pytorch/pytorch/issues/128505

```python
# works only with a single, non-batched tensor of token ids

import torch

def bpedecode_loop(token_ids, token_utf8bytes, token_lens):
    inds = torch.cat((torch.zeros_like(token_lens[:1]), token_lens.cumsum(-1)))
    return torch.cat([token_utf8bytes[inds[i]:inds[i_p_1]] for i, i_p_1 in zip(token_ids, token_ids + 1)])

def bpedecode_vec(token_ids, token_utf8bytes, token_lens):
    inds_start = torch.cat((torch.zeros_like(token_lens[:1]), token_lens[:-1].cumsum(-1)))
    inds_end = inds_start + token_lens

    starts = inds_start[token_ids]
    ends = inds_end[token_ids]
    lens = token_lens[token_ids]
    ones = torch.ones_like(token_ids)

    starts_shifted = torch.cat([starts[:1], starts[1:] - ends[:-1] + 1])
    repeats = torch.stack([ones, lens - 1], dim = -1).flatten()
    i = torch.stack([starts_shifted, ones], dim = -1).flatten()
    I = i.repeat_interleave(repeats).cumsum(-1)
    return token_utf8bytes[I]

if __name__ == '__main__':
    token_ids = torch.tensor([1, 0, 1, 3], dtype = torch.int64)
    token_utf8bytes = torch.tensor([1, 17, 31, 2, 2, 2, 2, 3, 7], dtype = torch.uint8)
    token_lens = torch.tensor([1, 2, 4, 2], dtype = torch.int64)

    print('loop:', bpedecode_loop(token_ids, token_utf8bytes, token_lens))
    print(' vec:', bpedecode_vec (token_ids, token_utf8bytes, token_lens))
```

### Alternatives

_No response_

### Additional context

_No response_