[vulkan] Replace *_size() functions with get_dim<N>() (#83423)
This diff replaces the `batch_size`, `channels_size`, etc. functions with a template function `get_dim<N>` to reduce duplicate code.
`batch_size()` has been replaced with `get_dim<Dim4D::Batch>` and so on.
Differential Revision: [D38706526](https://our.internmc.facebook.com/intern/diff/D38706526/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83423
Approved by: https://github.com/salilsdesai
