.backward() on MSELoss fails with IndexError: Dimension out of range on MPS

[ilyarepko]

🐛 Describe the bug

import torch
import torch.nn as nn

loss_fn = torch.nn.MSELoss(reduction="sum")
model = nn.Sequential(
    nn.Linear(2, 2)
)

# Works on CPU
x = torch.tensor([1.0, 1.0], dtype=torch.float32)
y = torch.tensor([2.0, 2.0], dtype=torch.float32)
model.zero_grad()
y_predicted = model(x)
loss = loss_fn(y_predicted, y)
loss.backward()

# Doesn't on MPS
device = torch.device("mps")
x_mps = x.to(device=device)
y_mps = y.to(device=device)
model_mps = model.to(device)
model_mps.zero_grad()
y_predicted = model_mps(x_mps)
loss = loss_fn(y_predicted, y_mps)
loss.backward()

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Input In [48], in <cell line: 25>()
     23 y_predicted = model_mps(x_mps)
     24 loss = loss_fn(y_predicted, y_mps)
---> 25 loss.backward()

File /opt/homebrew/lib/python3.9/site-packages/torch/_tensor.py:400, in Tensor.backward(self, gradient, retain_graph, create_graph, inputs)
    391 if has_torch_function_unary(self):
    392     return handle_torch_function(
    393         Tensor.backward,
    394         (self,),
   (...)
    398         create_graph=create_graph,
    399         inputs=inputs)
--> 400 torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)

File /opt/homebrew/lib/python3.9/site-packages/torch/autograd/__init__.py:173, in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    168     retain_graph = create_graph
    170 # The reason we repeat same the comment below is that
    171 # some Python versions print out the first line of a multi-line function
    172 # calls in the traceback and some print out the last line
--> 173 Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    174     tensors, grad_tensors_, retain_graph, create_graph, inputs,
    175     allow_unreachable=True, accumulate_grad=True)

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

Versions

bash-5.1$ python3 collect_env.py 
Collecting environment information...
PyTorch version: 1.13.0.dev20220614
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.4 (arm64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: version 3.23.2
Libc version: N/A

Python version: 3.9.13 (main, May 24 2022, 21:13:51)  [Clang 13.1.6 (clang-1316.0.21.2)] (64-bit runtime)
Python platform: macOS-12.4-arm64-arm-64bit
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
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.0rc2
[pip3] torch==1.13.0.dev20220614
[pip3] torchaudio==0.14.0.dev20220603
[pip3] torchvision==0.14.0.dev20220614
[conda] Could not collect

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7 @kulinseth

[qqaatw]

I can reproduce. The problem stems from linear backward

def test_79784():
    import torch
    import torch.nn as nn

    device = torch.device("mps")
    weight = torch.nn.Parameter(torch.randn(2, 2))


    x = torch.tensor([1.0, 1.0], dtype=torch.float32)
    out = nn.functional.linear(x, weight)
    grad = torch.autograd.grad(out.sum(), weight)
    
    x_mps = x.to(device=device)
    weight_mps = weight.to(device)
    out_mps = x_mps @ weight_mps.T
    grad_mps = torch.autograd.grad(out_mps.sum(), weight_mps)

    torch.testing.assert_close(out, out_mps, check_device=False)
    torch.testing.assert_close(grad[0], grad_mps[0], check_device=False)

    out_mps = nn.functional.linear(x_mps, weight_mps)
    grad_mps = torch.autograd.grad(out_mps.sum(), weight_mps) # fail

Edited: seems already fixed by #80759

[ilyarepko]

Yes, I tested it with torch==1.13.0.dev20220705, and now it works. Thanks a lot!