Fix batch dimension adjustment in shader indexing utils (#5399)

Summary:
Pull Request resolved: #5399

## Context

In the GLSL shader utility functions, the functions which use the axis map to obtain texture position or tensor index had a small bug where the size of the tensor dimension used to concatenate with the batch dimension was incorrectly retrieved as `sizes.w`. The correct way to retrieve this size is `sizes[axis_map[3]]`.

This diff fixes this bug.

ghstack-source-id: 242840432
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D62772113

fbshipit-source-id: 1285b1509ff652703c7ecf16dbb2cc30d6d70769

Author: SS-JIA

backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h

@@ -270,7 +270,7 @@ ivec3 to_texture_pos(
 
   // Adjust batch dim if needed
   if (sizes.w > 1) {
-    pos[axis_map[axis_map[3]]] += idx.w * sizes.w;
+    pos[axis_map[axis_map.w]] += idx.w * sizes[axis_map.w];
   }
 
   // Adjust packed dim. Moving 1 texel unit along the packed dim traverses 4
@@ -317,7 +317,7 @@ ivec4 to_texture_elem_pos(
 
   // Adjust batch dim if needed
   if (sizes.w > 1) {
-    pos[axis_map[axis_map[3]]] += idx.w * sizes.w;
+    pos[axis_map[axis_map.w]] += idx.w * sizes[axis_map.w];
   }
 
   // Adjust packed dim. Moving 1 texel unit along the packed dim traverses 4