RuntimeError: derivative for aten::linear_backward is not implemented on MPS

fazaghifari: <p>Dear all,</p>
<p>I’m currently using Pytorch 1.13.1 on my Macbook M1 Pro</p>
<p>Long story short, I wanted to reproduce this notebook:<br>
<a href="https://github.com/nanditadoloi/PINN/blob/main/solve_PDE_NN.ipynb" rel="noopener nofollow ugc">https://github.com/nanditadoloi/PINN/blob/main/solve_PDE_NN.ipynb</a> on my machine. But obviously, I have to change the device to <code>"mps"</code></p>
<p>However, it returns <code>RuntimeError: derivative for aten::linear_backward is not implemented</code> if I set <code>device = torch.device("mps")</code></p>
<p>But it runs fine if I set the device to <code>"cpu"</code>. Does anyone have a similar problem?</p>
<blockquote>
<p><strong>Addition</strong>: Just found the culprit. In the training step in the given link, it combines two loss functions the <code>mse_u</code> and <code>mse_f</code>. I tried to use <code>mse_u</code> only and it works fine on MPS, but when I tried <code>mse_f</code> only, the error message appear.</p>
</blockquote>
<blockquote>
<p>But I still have no idea why.</p>
</blockquote>
<p>Thank you in advance</p>

rekarantnetwork: <p>I work with PINNs on MPS too. It seems not all gradient calculation functions have been implemented for Apple Silicon MPS backend.</p>
<p>So we can’t calculate the PDE loss in the domain using MPS.</p>
<pre><code class="lang-auto">u_x = torch.autograd.grad(u.sum(), x, create_graph=True)[0]
u_t = torch.autograd.grad(u.sum(), t, create_graph=True)[0]
</code></pre>
<p>would throw errors.</p>
<p>This is a known problem. It seems you can still make it work by using the nightly version (only for linear layers). Please check the following issue:</p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/98498">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/98498" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/98498" target="_blank" rel="noopener nofollow ugc">Higher order derivatives not working when setting compute device to `torch.device("mps")`</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2023-04-06" data-time="12:36:57" data-timezone="UTC">12:36PM - 06 Apr 23 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/HamedHemati" target="_blank" rel="noopener nofollow ugc">
          <img alt="HamedHemati" src="https://avatars.githubusercontent.com/u/11629528?v=4" class="onebox-avatar-inline" width="20" height="20">
          HamedHemati
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: autograd
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: mps
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">### 🐛 Describe the bug

I have a code example that performs fast updates, simila<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">r to what we have in MAML, to a perform inner updates and then computes the meta-loss to backpropagate through the optimization trajectory. When computing gradients in the fast updates, we can set `create_graph=True` to enable second-order derivatives when call backward on the meta-loss. When using `torch.device("mps")`, it throws an error that `derivative for aten:linear is not implemented`. It works fine when you set `create_graph=False` in the inner updates but then it wouldn't compute the higher order derivatives. I don't get the error when using `torch.device("cpu")` and `torch.device("cuda")`. Here is the code to reproduce the error:

```python

device = torch.device("mps")
model = nn.Sequential(nn.Linear(10, 5), nn.ReLU(), nn.Linear(5, 1))
model.to(device)

# Initial fast parameters
fast_params_0 = {n: deepcopy(p) for (n, p) in model.named_parameters()}

# First inner update
x = torch.randn(10, 10, device=device)
y = torch.randn(10, 1, device=device)
logits_0 = torch.func.functional_call(model, fast_params_0, x)
loss = nn.MSELoss()(logits_0, y)
grads_0 = torch.autograd.grad(loss, fast_params_0.values(),
                              create_graph=True,
                              retain_graph=True)

# Compute fast parameters after the first inner update
fast_params_1 = {n: p - 0.1 * g for ((n, p), g) in zip(fast_params_0.items(), list(grads_0))}

# Compute meta-loss and backprop through the optimization trajectory
x = torch.randn(10, 10, device=device)
y = torch.randn(10, 1, device=device)
logits_1 = torch.func.functional_call(model, fast_params_1, x)
met_loss = nn.MSELoss()(logits_1, y)
met_loss.backward()

```

And, the error I get:
```
RuntimeError: derivative for aten::linear_backward is not implemented
```

*I get the same error for any layer type. 

### Versions

```
[pip3] numpy==1.23.5
[pip3] pytorchcv==0.0.67
[pip3] torch==2.0.0
[pip3] torchaudio==2.0.0
[pip3] torchmetrics==0.11.4
[pip3] torchvision==0.15.0
[conda] numpy                     1.23.5           py39h1398885_0  
[conda] numpy-base                1.23.5           py39h90707a3_0  
[conda] pytorch                   2.0.0                   py3.9_0    pytorch
[conda] pytorchcv                 0.0.67                   pypi_0    pypi
[conda] torchaudio                2.0.0                  py39_cpu    pytorch
[conda] torchmetrics              0.11.4                   pypi_0    pypi
[conda] torchvision               0.15.0                 py39_cpu    pytorch
```

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @kulinseth @malfet @DenisVieriu97 @razarmehr @abhudev</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>