Repository: pytorch
Issue: `embedding_bag` has unexpected behavior when given `offsets` that are not monotonically increasing
Labels: module: nn,triaged,module: embedding

## 🐛 Bug

When `embedding_bag` is given a 1-D `input`, it requires the `offsets` argument to tell it where the bags begin and end. Interestingly, it's possible to give an `offsets` arg that is not monotonically increasing, and the behavior we get from this is surprising on the CPU, and on the CUDA side it just crashes.

The documentation leaves this case ambiguous, stating simply "`offsets` determines the starting index position of each bag (sequence) in `input`".

## To Reproduce

Steps to reproduce the behavior:

```python
>>> weights = torch.randn(4, 3)
>>> indices = torch.tensor([0, 1, 2, 3])
>>> weights
tensor([[-0.0331, -0.2412,  0.5930],
        [-0.1633,  0.7328, -0.2875],
        [ 1.1365, -0.8380, -2.0821],
        [-0.5085,  2.7105,  2.1375]])
```
This is the normal, intended way to use `offsets`. The indices 0 and 1 are placed in the first bag, and indices 2 and 3 are placed in the second bag
```python
>>> torch.nn.functional.embedding_bag(indices, weights, torch.tensor([0, 2]), mode='max')
tensor([[-0.0331,  0.7328,  0.5930],
        [ 1.1365,  2.7105,  2.1375]])
```
Here's the equivalent 2-D input call:
```python
>>> torch.nn.functional.embedding_bag(torch.tensor([[0, 1], [2, 3]]), weights, mode='max')
tensor([[-0.0331,  0.7328,  0.5930],
        [ 1.1365,  2.7105,  2.1375]])
```
But now, let's try an `offsets` tensor that's not monotonically increasing
```python
>>> torch.nn.functional.embedding_bag(indices, weights, torch.tensor([0, 2, 1]), mode='max')
tensor([[-0.0331, -0.2412,  0.5930],
        [-0.1633,  0.7328, -0.2875],
        [ 1.1365,  2.7105,  2.1375]])
```
It turns out that the CPU version of mode='max' seems to be sorting the `offsets`. Or at least it appears so in this one case. Here's the presorted equivalent:
```python
>>> torch.nn.functional.embedding_bag(indices, weights, torch.tensor([0, 1, 2]), mode='max')
tensor([[-0.0331, -0.2412,  0.5930],
        [-0.1633,  0.7328, -0.2875],
        [ 1.1365,  2.7105,  2.1375]])
```
However, the CUDA version just gives a device-side assert
```
>>> torch.nn.functional.embedding_bag(indices.cuda(), weights.cuda(), torch.tensor([0, 2, 1]).cuda(), mode='max')
...
RuntimeError: CUDA error: device-side assert triggered
```
The CPU version of mode='sum' does not appear to sort the `offsets` first
```python
>>> torch.nn.functional.embedding_bag(indices, weights, torch.tensor([0, 2, 1]), mode='sum')
tensor([[-0.1964,  0.4916,  0.3055],
        [ 0.0000,  0.0000,  0.0000],
        [ 0.4646,  2.6053, -0.2322]])
```
The `offsets` sorted version does not give the same result
```python
>>> torch.nn.functional.embedding_bag(indices, weights, torch.tensor([0, 1, 2]), mode='sum')
tensor([[-0.0331, -0.2412,  0.5930],
        [-0.1633,  0.7328, -0.2875],
        [ 0.6279,  1.8725,  0.0553]])
```
Instead, this is an equivalent call, showing that the `offsets=[0, 2, 1]` call actually created bags with indices [0, 1], [], and [1, 2, 3]. In other words, each pair of offsets indicates where the bag begins and ends, and if the endpoint is equal or less than the beginning, the bag is empty.
```python
>>> torch.nn.functional.embedding_bag(torch.tensor([0, 1, 1, 2, 3]), weights, torch.tensor([0, 2, 2]), mode='sum')
tensor([[-0.1964,  0.4916,  0.3055],
        [ 0.0000,  0.0000,  0.0000],
        [ 0.4646,  2.6053, -0.2322]])
```
And again, mode='max' does NOT behave the same way as mode='sum'. Note that the following call does not match with the previous mode='max' result (copied again here for clarity).
```python
>>> torch.nn.functional.embedding_bag(torch.tensor([0, 1, 1, 2, 3]), weights, torch.tensor([0, 2, 2]), mode='max')
tensor([[-0.0331,  0.7328,  0.5930],
        [ 0.0000,  0.0000,  0.0000],
        [ 1.1365,  2.7105,  2.1375]])
# previous
>>> torch.nn.functional.embedding_bag(indices, weights, torch.tensor([0, 2, 1]), mode='max')
tensor([[-0.0331, -0.2412,  0.5930],
        [-0.1633,  0.7328, -0.2875],
        [ 1.1365,  2.7105,  2.1375]])
```
CUDA mode='sum' also has a device-side assert
```python
>>> torch.nn.functional.embedding_bag(indices.cuda(), weights.cuda(), torch.tensor([0, 2, 1]).cuda(), mode='sum')
...
RuntimeError: CUDA error: device-side assert triggered
```

## Expected behavior

I would expect the behavior of `offsets` to be consistent across all `mode` settings and across different devices. I think we have three main choices:

- Throw an error if `offsets` is not monotonically increasing. CUDA seems to be already doing this, although it is a device-side assert which is not ideal.
- Always sort `offsets` first (or whatever mode='max' is doing at the moment)
- Each subsequent pair of  `offsets` indicates the inclusive starting and exclusive ending index. If the ending index is less than or equal to the starting index, the bag is empty.

In my opinion, throwing an error would be the best choice. The CUDA impl is already doing that. Plus, it's the most conceptually simple behavior. I don't see any benefit for allowing `offsets` to not be monotonically increasing. Although if someone has a use case, we can certainly consider supporting it.

Of course the documentation will need to be updated as well to reflect the decision.

## Environment


```
PyTorch version: 1.9.0a0+gitd490e01
Is debug build: False
CUDA used to build PyTorch: 11.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (crosstool-NG 1.24.0.133_b0863d8_dirty) 9.3.0
Clang version: Could not collect
CMake version: version 3.19.7

Python version: 3.9 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 11.0.221
GPU models and configuration: 
GPU 0: TITAN RTX
GPU 1: TITAN RTX

Nvidia driver version: 460.27.04
cuDNN version: /usr/local/cuda-10.2.89/targets/x86_64-linux/lib/libcudnn.so.7
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.2
[pip3] pytorch-sphinx-theme==0.0.24
[pip3] torch==1.9.0a0+git256b670
[conda] magma-cuda110             2.5.2                         1    pytorch
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] mkl-include               2020.4             h726a3e6_304    conda-forge
[conda] numpy                     1.20.2           py39hdbf815f_0    conda-forge
[conda] pytorch-sphinx-theme      0.0.24                    dev_0    <develop>
[conda] torch                     1.9.0a0+git256b670           dev_0    <develop>
```


cc @albanD @mruberry @jbschlosser