Repository: pytorch
Issue: In-place operations fail to work after moving tensor to MPS (after linalg.lstsq operation)
Labels: triaged,module: mps

### 🐛 Describe the bug

I am attempting to use `torch.linalg.lstsq` to initialize a tensor that I subsequently train using backprop, on MPS. Since lstsq gives the error `The operator 'aten::linalg_lstsq.out' is not currently implemented for the MPS device.`, I am converting my tensors to cpu for that operation, and then moving them back.

However, the result of lstsq appears to break subsequent in-place operations on MPS, as demonstrated by the following min repro code:

I am running PyTorch 2.1.0 on a 2021 M1 Max.

```python
import torch


def inplace_op(tens):
    return tens.addcmul_(torch.ones_like(tens) * 0.5, torch.ones_like(tens) * 0.5, value=1)  # Fails
    # return tens.addcdiv_(torch.ones_like(tens) * 0.5, torch.ones_like(tens) * 0.5, value=1)  # Fails
    # return tens.add_(torch.ones_like(tens) * 0.5)  # Works
    # return tens.mul_(torch.ones_like(tens) * 0.5)  # Works


def run_test_on(device, dim_a, dim_b, workaround_id=None):
    print(f"======= Running with device: {device}, workaround id: {workaround_id} =======")

    # Create dummy tensors for running through lstsq
    A = torch.ones(size=(dim_a, dim_b)).to("cpu")
    B = torch.ones(size=(dim_a, dim_a)).to("cpu")

    # Run, grab the result
    lstsq_info = torch.linalg.lstsq(A, B)
    to_modify_in_place = lstsq_info.solution

    # Apply a "work-around", to see what impact it has on behavior
    if workaround_id == 0:
        # This work-around attempts to completely re-create the tensor, by returning the data to numpy form and converting back... which does not work.
        to_modify_in_place = torch.tensor(to_modify_in_place.detach().cpu().numpy())

    elif workaround_id == 1:
        # If you attempt the work-around that does work (=2), but with "zeros_like" instead of "zeros", it still fails
        to_modify_in_place = torch.zeros_like(to_modify_in_place) + to_modify_in_place.detach().cpu().numpy()

    elif workaround_id == 2:
        # The work-around of, instead, adding the numpy data to a correctly-sized tensor does work, however.
        to_modify_in_place = torch.zeros(to_modify_in_place.shape) + to_modify_in_place.detach().cpu().numpy()

    # Print the original tensor, and the tensor after it has been modified by an in-place operation (which should add 0.25 to each element)
    print(f"[Pre-to] orig: {to_modify_in_place}, \n post: {inplace_op(to_modify_in_place)} \n")

    # Move the tensor to our device...
    to_modify_in_place = to_modify_in_place.to(device)

    # And run the in-place operation again. It *should* add 0.25 to each element again, but for MPS (without workaround 1), it fails
    print(f"[Post-to] orig: {to_modify_in_place}, \n post: {inplace_op(to_modify_in_place)} \n")

    return to_modify_in_place


def run():
    dim_a = 2
    dim_b = 2
    cpu_result = run_test_on("cpu", dim_a, dim_b)

    for workaround_id in [None, 0, 1, 2]:
        mps_result = run_test_on("mps", dim_a, dim_b, workaround_id=workaround_id)

        print(f"CPU = MPS? {(cpu_result == mps_result.cpu()).all()} \n\n")


if __name__ == "__main__":
    run()
```

This has the following result on my machine (broken up for convenient description):

With CPU, we see no issue: 0.25 is added on each call to the in-place function:

```
======= Running with device: cpu, workaround id: None =======
[Pre-to] orig: tensor([[0.5000, 0.5000],
        [0.5000, 0.5000]]), 
 post: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]]) 

[Post-to] orig: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]]), 
 post: tensor([[1., 1.],
        [1., 1.]]) 
```




When we convert to MPS, however, we see that the second in-place operations breaks, failing to add 0.25 after moving the tensor to MPS:

```
======= Running with device: mps, workaround id: None =======
[Pre-to] orig: tensor([[0.5000, 0.5000],
        [0.5000, 0.5000]]), 
 post: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]]) 

[Post-to] orig: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]], device='mps:0'), 
 post: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]], device='mps:0') 

CPU = MPS? False 
```


To understand the problem (and to unblock myself), I tried a few work-arounds that might be insightful? The first one was simply to do `to_modify_in_place = torch.tensor(to_modify_in_place.detach().cpu().numpy())`. I.e. move the data to numpy and re-wrap it. We can see that this one fails as well:

```
======= Running with device: mps, workaround id: 0 =======
[Pre-to] orig: tensor([[0.5000, 0.5000],
        [0.5000, 0.5000]]), 
 post: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]]) 

[Post-to] orig: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]], device='mps:0'), 
 post: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]], device='mps:0') 

CPU = MPS? False 

```

I next attempted to do `to_modify_in_place = torch.zeros_like(to_modify_in_place) + to_modify_in_place.detach().cpu().numpy()`, which also fails:

```
======= Running with device: mps, workaround id: 1 =======
[Pre-to] orig: tensor([[0.5000, 0.5000],
        [0.5000, 0.5000]]), 
 post: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]]) 

[Post-to] orig: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]], device='mps:0'), 
 post: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]], device='mps:0') 

CPU = MPS? False 
```

However, `to_modify_in_place = torch.zeros(to_modify_in_place.shape) + to_modify_in_place.detach().cpu().numpy()` succeeds:

```
======= Running with device: mps, workaround id: 2 =======
[Pre-to] orig: tensor([[0.5000, 0.5000],
        [0.5000, 0.5000]]), 
 post: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]]) 

[Post-to] orig: tensor([[0.7500, 0.7500],
        [0.7500, 0.7500]], device='mps:0'), 
 post: tensor([[1., 1.],
        [1., 1.]], device='mps:0') 

CPU = MPS? True 
```

It is interesting to me that `zeros_like` on the broken tensor creates a tensor that is still broken, but `zeros` using the shape (and with both at that point on CPU) does not. I am unsure what is happening under the hood exactly, but I thought `zeros_like` was basically just "shape and device"?


A couple other observations:
1. The problem does not occur if either `dim_a` or `dim_b` is 1. (The issue is not related to the fact that B is not full-rank, however; the problem persists regardless of matrix initialization.)
2. I discovered this problem with `addcmul_` from RMSProp, and it repros with `addcdiv_` as well, but `add_` and `mul_` still work. (Can be tested by switching what is commented in `def inplace_op(tens)`.

Thanks! Let me know if I can provide any other info.

### Versions

PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.0 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.0.40.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.18 (main, Sep 11 2023, 08:25:10)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.0-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
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
Is XNNPACK available: True

CPU:
Apple M1 Max

Versions of relevant libraries:
[pip3] numpy==1.26.0
[pip3] torch==2.1.0
[pip3] torchvision==0.16.0
[conda] numpy                     1.26.0           py39h3b2db8e_0
[conda] numpy-base                1.26.0           py39ha9811e2_0
[conda] pytorch                   2.1.0                   py3.9_0    pytorch
[conda] torchvision               0.16.0                 py39_cpu    pytorch

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr