[iOS GPU][Kernel] Fix the softmax kernels (#54519)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/54519
The current MPSCNNSoftmax kernels operates on tensors' feature channels. Therefore, in order to use it, we need to reshape the input tensors based on the value of `dim` . Currently, I decide to limit the input to be two dimensional. I'll remove the constraint once we have shader implementations.
ghstack-source-id: 124497702
Test Plan:
- SandcastleCI
- CircleCI
Reviewed By: dhruvbird
Differential Revision: D27218823
fbshipit-source-id: 48c427ceedb42e63c183114939ca801ebfc81fd9
