[torchax] Fix functional.max_pool

[dvhg]

Fix #8086.

Also possibly related to #8241.

To summarize the changes to max pool:

1. Dilation now correctly accounted for, including in padding calculation following the formula for output shape from here
2. Default stride is equal to kernel size
3. Indices are computed per-batch instead of across the entire input
4. Ties break towards higher indices (i.e. pick the higher index if there is more than 1 max value in a window). This is arbitrary but matches Torch's behavior on CPU.

[ManfeiBai]

Thanks, LGTM