Repository: pytorch
Issue: combining `vmap` with NN containing `MaxPool2d' leads to discrepancies in output
Labels: triaged,module: functorch

### 🐛 Describe the bug

I would like to compute the Jacobian of the output of a CNN with respect to the inputs, using `torch.func.jacrev`. For a batch of (e.g. MNIST) images, you can do this with a regular for loop or (preferably) with `torch.vmap` instead. Both methods should lead to the same outcome. However, it turns out this is not always the case when the CNN contains a `MaxPool2d`-layer.

The code snippet below gives a concrete example of the discrepancy. When you replace `MaxPool2d` by `AvgPool2d` the issue disappears. Also, the discrepancy grows when the input data has a smaller range (controlled by `scale` in the code below). When you set `scale=1.0` there is no issue, when you set `scale=0.000001` the differences between the two methods are larger. I suspect that `torch.vmap` acts incorrectly on `MaxPool2d` when two (or more) equal values are being pooled.

It would be nice if this issue gets resolved.

```python
import torch
import torch.nn as nn

class net(nn.Module):
    
    def __init__(self):
        super(net, self).__init__()
        
        self.conv1 = nn.Sequential(         
            nn.Conv2d(1, 16, 5, 1, 2),       
            nn.ReLU(),                      
            nn.MaxPool2d(kernel_size=2)) # try nn.AvgPool2d instead
        
        self.FC = nn.Linear(16 * 14 * 14, 10)
        
    def forward(self, x):
        
        x = self.conv1(x)
        x = x.flatten(start_dim=1)     
        output = self.FC(x)
        
        return output

# initialize model
model = net()
def model_single_input(x):      # forward pass on single input (image)
    return model(x.unsqueeze(dim=0))

# generate batch of input data
batch_size = 100
scale = 0.001    # smaller scale results in larger discrepancy between two methods
X = scale * torch.rand(size=(batch_size, 1, 28, 28,)) # data in MNIST format

# compute jacrev with vmap
jacrev_vmap = torch.vmap(torch.func.jacrev(model_single_input))(X)

# compute jacrev without vmap
jacrev = torch.zeros(jacrev_vmap.shape)
for i, x in enumerate(X):
    
    jacrev[i] = torch.func.jacrev(model_single_input)(x)
    
# comparison of methods
print("Equivalence of two methods: ", torch.all(jacrev == jacrev_vmap))

print("Maximum absolute difference between the methods, per input (image):")
torch.max(torch.abs(jacrev - jacrev_vmap).view(batch_size, -1), dim=(1)).values
```
Example of output:
```
Equivalence of two methods:  tensor(False)
Maximum absolute difference between the methods, per input (image):
tensor([0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0046, 0.0000, 0.0000, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0061, 0.0000, 0.0000, 0.0000, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0048, 0.0000, 0.0000, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0034, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
        0.0000, 0.0061, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
        0.0000], grad_fn=<MaxBackward0>)
```

### Versions

PyTorch version: 2.0.1
Python version: 3.10.9 (main, Mar  1 2023, 12:20:14) [Clang 14.0.6 ] (64-bit runtime)
[conda] pytorch                   2.0.1                  py3.10_0    pytorch


cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99