Repository: pytorch
Issue: Incosistency with args for `nn.functional.max_poolNd` vs `nn.MaxPoolNd` functions
Labels: module: nn,triaged,module: pooling

## 🐛 Bug

Note: For this issue, I'll be taking `max_pool2d` as an example function. The same is applicable for `max_pool1d` and `max_pool3d`.

The position/index (starting from 0) of `return_indices` arg for `nn.functional.max_pool2d` as described in the [documentation](https://pytorch.org/docs/stable/generated/torch.nn.MaxPool2d.html#torch.nn.MaxPool2d) should be `5` but when used at the 5th position, it doesn't do what it should (should return a tuple of tensors). Instead, passing it as the `6th` argument/index works.

## To Reproduce

_Note the order of the args corresponds to:_ `input_shape, kernel_size, stride, padding, dilation, return_indices, ceil_mode`

`nn.functional.max_pool2d` behavior:

```python
>>> tens = torch.randn((1, 3, 9, 9))

# return_indices should be True here according to the doc
>>> len(nn.functional.max_pool2d(tens, (3, 3), None, 0, 1, True, False))
1

# return_indices should be False here according to the doc
>>> len(nn.functional.max_pool2d(tens, (3, 3), None, 0, 1, False, True))
2
```

`nn.MaxPool2d` behavior:

```python
>>> tens = torch.randn((1, 3, 9, 9))

# Note that True is passed at the 5th index, and it works fine (as expected): output length is 2
>>> mp2d = torch.nn.MaxPool2d((3, 3), None, 0, 1, True, False)
>>> len(mp2d(tens))
2

# Note that False is passed at the 5th index, and it works fine (as expected): output length is 1
>>> mp2d = torch.nn.MaxPool2d((3, 3), None, 0, 1, False, True)
>>> len(mp2d(tens))
1
```

## Expected behavior

`nn.functional.max_pool2d` behavior should be consistent with `nn.MaxPool2d` as that's where the [docs](https://pytorch.org/docs/stable/generated/torch.nn.functional.max_pool2d.html) point it to!

```python
>>> tens = torch.randn((1, 3, 9, 9))

>>> len(nn.functional.max_pool2d(tens, (3, 3), None, 0, 1, True, False))
2

>>> len(nn.functional.max_pool2d(tens, (3, 3), None, 0, 1, False, True))
1
```

## Environment

 - PyTorch Version (e.g., 1.0): 1.10.0a0+git6a6830f
 - OS (e.g., Linux): Linux
 - How you installed PyTorch (`conda`, `pip`, source): source

## Additional context

Came up with this issue when `test_variant_consistency_jit_nn_functional_max_pool2d_cuda/cpu_float32` was failing (while trying to add `OpInfo` for it).

```python
======================================================================
ERROR: test_variant_consistency_jit_nn_functional_max_pool2d_cuda_float32 (__main__.TestJitCUDA)
----------------------------------------------------------------------
... (only sharing relevant error detail)
RuntimeError: aliasOp != torch::jit::getOperatorAliasMap().end()INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":159, please report a bug to PyTorch. 
```

Trying to debug further, here is how the method graph looks:

* when `return_indices` is passed as `False` as positional arg at index `5` to `nn.functional.max_pool2d`, and index 6 which should correspond to `ceil_mode` is passed as `True`) - It returns a tuple of tensors (output and indices) - which is wrong.

```python
graph(%i0.1 : Tensor):
  %9 : bool = prim::Constant[value=0]() # <string>:3:76
  %33 : int[] = prim::Constant[value=[5, 5]]()
  %34 : int[] = prim::Constant[value=[2, 1]]()
  %35 : int[] = prim::Constant[value=[2, 2]]()
  %37 : int[] = prim::Constant[value=annotate(List[int], [])]()
  %30 : Tensor, %31 : Tensor = aten::max_pool2d_with_indices(%i0.1, %33, %37, %34, %35, %9) # /home/krshrimali/.conda/envs/pytorch-cuda-dev/lib/python3.9/site-packages/torch/nn/functional.py:693:11
  %32 : (Tensor, Tensor) = prim::TupleConstruct(%30, %31)
  return (%32)
```

* when `return_indices` is passed as `True` as positional arg at index `5` to `nn.functional.max_pool2d`, and index 6 which should correspond to `ceil_mode` is passed as `False`)

```python
graph(%i0.1 : Tensor):
  %10 : bool = prim::Constant[value=0]() # <string>:3:82
  %31 : int[] = prim::Constant[value=[5, 5]]()
  %32 : int[] = prim::Constant[value=[2, 1]]()
  %33 : int[] = prim::Constant[value=[2, 2]]()
  %35 : int[] = prim::Constant[value=annotate(List[int], [])]()
  %30 : Tensor = aten::max_pool2d(%i0.1, %31, %35, %32, %33, %10) # /home/krshrimali/.conda/envs/pytorch-cuda-dev/lib/python3.9/site-packages/torch/nn/functional.py:718:11
  return (%30)
```

The issue seems to be with the `arg_index` mentioned in the `boolean_dispatch` call. Please see https://github.com/pytorch/pytorch/pull/62544 ~for the suggested fix~ for the draft fix, WIP.

cc: @thomasjpfan @Lezcano @kshitij12345 

cc @albanD @mruberry @jbschlosser @heitorschueroff