Repository: pytorch
Issue: [torchscript] nn.Embedding(_weight), tensor.masked_fill_, and torch.autograd.grad causes wrong gradient
Labels: oncall: jit

## 🐛 Bug

## To Reproduce

The following script reproduces the problem:

```python
import numpy as np
import torch
import torch.nn as nn

class _Attention(nn.Module):
    def __init__(self,
                 input_size=128,
                 attention_hidden_size=128):
        super().__init__()
        self.W1 = nn.Linear(input_size, attention_hidden_size)
        self.W2 = nn.Linear(attention_hidden_size, 1, bias=False)
        self.softmax = nn.Softmax(dim=-1)

    def score(self, values):
        h_t = torch.tanh(self.W1(values))
        score_t = self.W2(h_t)
        return score_t.squeeze(2)

    def forward(self, values, mask):
        score_t = self.score(values)
        # masked_fill_ course the error, torch.where works fine
        score_t.masked_fill_(torch.logical_not(mask), -np.inf)

        alpha_t = self.softmax(score_t)
        context_t = (alpha_t.unsqueeze(2) * values).sum(1)

        return context_t

class GradientImportanceMeasure(nn.Module):
    def __init__(self, hidden_size=128, num_of_classes=2):
        super().__init__()
        # If _weight is not used here, iteration: 1 (second iteration) produces a RuntimeError:
        # > RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.
        self.embedding = nn.Embedding(100, hidden_size,
                                      padding_idx=0,
                                      _weight=torch.Tensor(np.random.randn(100, hidden_size)))
        self.attention = _Attention(hidden_size, hidden_size)
        self.decoder = nn.Linear(hidden_size, num_of_classes)

    def forward(self, sentence, mask) -> torch.Tensor:
        embedding = self.embedding(sentence)
        embedding.requires_grad_(True)
        z = self.attention(embedding, mask)
        y = self.decoder(z)[0, 1]

        yc_wrt_embedding, = torch.autograd.grad([y], (embedding, ))
        if yc_wrt_embedding is None:
            raise ValueError('Could not compute gradient')

        return yc_wrt_embedding

if __name__ == "__main__":
    torch.manual_seed(0)
    measure_slow = GradientImportanceMeasure()
    torch.manual_seed(0)
    measure_fast = torch.jit.script(GradientImportanceMeasure())

    sentence = torch.tensor([
        [   1,   24,  89,  70,    5, 40,  10,  10,    2]
    ])
    mask = torch.tensor([
        [False,  True,  True,  True,  True,  True,  True,  True, False]
    ])

    for i in range(2):
        print(f'iteration: {i}')
        np.testing.assert_allclose(
            measure_slow(sentence.detach().clone(), mask),
            measure_fast(sentence.detach().clone(), mask)
        )
```

The output is:

```
iteration: 0
Traceback (most recent call last):
  File "test-ultra-simplified-model.py", line 67, in <module>
    np.testing.assert_allclose(
  File "/Users/Andreas/.pyenv/versions/3.8.6/lib/python3.8/site-packages/numpy/testing/_private/utils.py", line 1527, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "/Users/Andreas/.pyenv/versions/3.8.6/lib/python3.8/site-packages/numpy/testing/_private/utils.py", line 840, in assert_array_compare
    raise AssertionError(msg)
AssertionError:
Not equal to tolerance rtol=1e-07, atol=0

Mismatched elements: 896 / 1152 (77.8%)
Max absolute difference: 0.01038937
Max relative difference: 230.3212
 x: array([[[ 0.      ,  0.      ,  0.      , ...,  0.      ,  0.      ,
          0.      ],
        [-0.003255, -0.008186, -0.006048, ..., -0.008964, -0.00644 ,...
 y: array([[[ 0.000000e+00,  0.000000e+00,  0.000000e+00, ...,
          0.000000e+00,  0.000000e+00,  0.000000e+00],
        [-1.406996e-05, -9.762709e-03, -6.487017e-03, ...,...
```

If `_weight` is not used in `nn.Embedding()` the output is:

```
iteration: 0
iteration: 1
Traceback (most recent call last):
  File "test-ultra-simplified-model.py", line 68, in <module>
    measure_fast(sentence.detach().clone(), mask)
  File "/Users/Andreas/.pyenv/versions/3.8.6/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1015, in _call_impl
    return forward_call(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "test-ultra-simplified-model.py", line 42, in forward
        embedding = self.embedding(sentence)
        embedding.requires_grad_(True)
        z = self.attention(embedding, mask)
            ~~~~~~~~~~~~~~ <--- HERE
        y = self.decoder(z)[0, 1]

  File "test-ultra-simplified-model.py", line 22, in forward
        score_t = self.score(values)
        # masked_fill_ course the error, torch.where does not
        score_t.masked_fill_(torch.logical_not(mask), -np.inf)
        ~~~~~~~~~~~~~~~~~~~~ <--- HERE

        alpha_t = self.softmax(score_t)
RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.
```

note that the RuntimeError occurs at the second iteration.

## Expected behavior

The `RuntimeError` always occurs at the first iteration. Or, even better `torch.jit.script` detects the issue at compile time.

## Environment

```

Collecting environment information...
PyTorch version: 1.9.0.dev20210324
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 10.13.6 (x86_64)
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.12.1

Python version: 3.8 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.4
[pip3] pytorch-lightning==1.2.0
[pip3] torch==1.9.0.dev20210324
[pip3] torchtext==0.8.1
[pip3] torchvision==0.8.2
[conda] Could not collect
```

Note this issue also exist for pytorch 1.8.0.


cc @gmagogsfm