Hessian is (incorrectly) zero when using MPS on M1 Mac, but not on cpu 
Labels: triaged, module: mps, module: functorch

Thank you for hard work writing and maintaining functorch. I ran into this issue recently. I would appreciate any help or insight :)!

### Set-up:
Using python 3.10.6, PyTorch 1.13.0, and functorch 1.13.0 on an M1 Max Mac. 

### Code: 
```
import torch
from functorch import vmap, grad, jacrev, make_functional

device =  "mps" if torch.backends.mps.is_available() else "cpu" #change to cpu to reproduce cpu output 

class DNN(torch.nn.Module):
    def __init__(self):
        super(DNN, self).__init__()

        self.net = torch.nn.Sequential(
            torch.nn.Linear(1, 2), 
            torch.nn.Tanh(),
            torch.nn.Linear(2, 2), 
            torch.nn.Tanh(),
            torch.nn.Linear(2, 2), 
            torch.nn.Tanh(),
            torch.nn.Linear(2, 1),
            )
        
    def forward(self, x):
        out =  self.net(x)
        return out
        
dnn = DNN().to(device)
fmodel, params = make_functional(dnn)

batch_size = 3
data = torch.ones(batch_size, 1).to(device)
targets = torch.ones(batch_size, 1).to(device)

#Modified from various functorch tutorials: 
def loss_fn(predictions, targets):
    return torch.nn.functional.mse_loss(predictions, targets)

def compute_loss_stateless_model (params, sample, target):
    batch = sample.unsqueeze(0)
    targets = target.unsqueeze(0)

    predictions = fmodel(params, batch) 
    loss = loss_fn(predictions, targets)
    return loss

def comp_loss (params, data, target):
    predictions = fmodel(params, data) 
    loss = loss_fn(predictions, targets)
    return loss    

ft_compute_grad = grad(compute_loss_stateless_model)
ft_compute_sample_grad = vmap(ft_compute_grad, in_dims=(None, 0, 0))
grads = ft_compute_sample_grad(params, data, targets)
compute_sample_hess = jacrev(vmap(ft_compute_grad, in_dims=(None, 0, 0)))

hess = compute_sample_hess(params,data,targets)
print(hess)
```

Output with MPS (cut off to save some space):
 
All zeros 

``` 
((tensor([[[[[0.],
           [0.]]],


         [[[0.],
           [0.]]]],



        [[[[0.],
           [0.]]],


         [[[0.],
           [0.]]]],



        [[[[0.],
           [0.]]],


         [[[0.],
           [0.]]]]], device='mps:0', grad_fn=<ViewBackward0>), tensor([[[[0., 0.]],
...

        [[0.]],

        [[0.]]], device='mps:0', grad_fn=<ViewBackward0>)))
``` 

Output with CPU (cut off to save some space):  

``` 
((tensor([[[[[-0.0468],
           [ 0.0294]]],


         [[[ 0.0294],
           [-0.0064]]]],



        [[[[-0.0468],
           [ 0.0294]]],


         [[[ 0.0294],
           [-0.0064]]]],



        [[[[-0.0468],
           [ 0.0294]]],


         [[[ 0.0294],
           [-0.0064]]]]], grad_fn=<ViewBackward0>), tensor([[[[-0.0468,  0.0294]],
...

        [[2.]],

        [[2.]]], grad_fn=<ViewBackward0>)))
```

### Other info:

`grads` gives what appears to be a correct non-zero result on both MPS and CPU 




cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev @zou3519 @Chillee @samdow @soumith