Bringing Tensor::gather() behavior closer to torch.gather()

[AnubhabB]

Definition:

input and index must have the same number of dimensions. It is also required that index.size(d) <= input.size(d) for all dimensions d != dim. out will have the same shape as index. Note that input and index do not broadcast against each other.

In our case we were effectively constraining the requirement to:
index.size(d) != input.size(d); where d != dim

This PR removes the != constraint and allows the op to go through as long as:
index.dims()[d] <= input.dims()[d]; where d != dim

Added tests with data generated and validated against torch's ScatterGather tests.

Tested on arm64, x86_64, cuda and metal. I don't have access to mkl and others, so unsure if the kernels hold true!

[AnubhabB]

That's odd .. All tests are passing for my test run on Ubuntu x86_64. Let me see if I can replicate this.

[LaurentMazare]

Looks more like an issue with the github CI infrastructure, I've relaunched the failed action.

[LaurentMazare]

Thanks, I've merged the changes after running the tests locally on my linux box and it was all fine.

[AnubhabB]

Thanks a ton @LaurentMazare .. this unblocks another PR :)